from dataclasses import dataclass

from django.db import transaction
from django.utils import timezone

from .brokerage_types import BrokerageSnapshot
from .brokerages.base import BaseBrokerageAdapter
from .brokerages.registry import get_brokerage_adapter
from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun


@dataclass(frozen=True)
class IngestionResult:
    accounts_processed: int
    assets_processed: int
    holdings_processed: int


class BrokerageIngestionService:
    def __init__(
        self,
        connection: BrokerageConnection,
        adapter: BaseBrokerageAdapter | None = None,
    ):
        self.connection = connection
        self.adapter = adapter or get_brokerage_adapter(connection)

    def ingest(self) -> IngestionRun:
        run = IngestionRun.objects.create(
            connection=self.connection,
            status=IngestionRun.STATUS_RUNNING,
            started_at=timezone.now(),
        )

        try:
            snapshot = self.adapter.fetch_snapshot()
            result = self._persist_snapshot(snapshot)
        except Exception as exc:
            run.status = IngestionRun.STATUS_FAILED
            run.error_message = str(exc)
            run.finished_at = timezone.now()
            run.save(
                update_fields=["status", "error_message", "finished_at"]
            )
            raise

        run.status = IngestionRun.STATUS_SUCCEEDED
        run.finished_at = timezone.now()
        run.accounts_processed = result.accounts_processed
        run.assets_processed = result.assets_processed
        run.holdings_processed = result.holdings_processed
        run.save(
            update_fields=[
                "status",
                "finished_at",
                "accounts_processed",
                "assets_processed",
                "holdings_processed",
            ]
        )
        return run

    @transaction.atomic
    def _persist_snapshot(self, snapshot: BrokerageSnapshot) -> IngestionResult:
        accounts_by_external_id = {}
        assets_by_external_id = {}

        for payload in snapshot.accounts:
            account, _ = Account.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "name": payload.name,
                    "number": payload.number,
                    "balance": payload.balance,
                },
            )
            accounts_by_external_id[payload.external_id] = account

        for payload in snapshot.assets:
            asset, _ = Asset.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "ticker": payload.ticker,
                    "price": payload.price,
                },
            )
            assets_by_external_id[payload.external_id] = asset

        for payload in snapshot.holdings:
            account = accounts_by_external_id[payload.account_external_id]
            asset = assets_by_external_id[payload.asset_external_id]
            AccountHolding.objects.update_or_create(
                provider=self.connection.provider,
                external_id=payload.external_id,
                defaults={
                    "account": account,
                    "asset": asset,
                    "ticker": payload.ticker,
                    "ammount": payload.amount,
                    "cost_basis": payload.cost_basis,
                    "purchase_date": payload.purchase_date,
                },
            )

        return IngestionResult(
            accounts_processed=len(snapshot.accounts),
            assets_processed=len(snapshot.assets),
            holdings_processed=len(snapshot.holdings),
        )