from __future__ import annotations

import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

from account_registry import Account, AccountRegistry
from device_auth import DeviceAuthFlow, DeviceAuthPrompt, DeviceAuthSession, terminate_process


@dataclass(frozen=True)
class FlowEvent:
    kind: str
    operation: str
    account: Account | None = None
    alias: str | None = None
    session: DeviceAuthSession | None = None
    prompt: DeviceAuthPrompt | None = None
    existing_accounts: list[tuple[str, str]] | None = None
    collision: Account | None = None
    retry_available: bool | None = None
    message: str | None = None
    error: Exception | None = None


class DeviceAuthOperation:
    def __init__(
        self,
        registry: AccountRegistry,
        device_auth_flow: DeviceAuthFlow,
    ) -> None:
        self.registry = registry
        self._device_auth_flow = device_auth_flow

    def repair(self, account: Account) -> Iterator[FlowEvent]:
        yield from self._run_repair(account)

    def run_repair(self, account: Account) -> Iterator[FlowEvent]:
        yield from self._run_repair(account)

    def repair_account(self, account: Account) -> Iterator[FlowEvent]:
        yield from self._run_repair(account)

    def add(self, alias: str, *, slug: str | None = None) -> Iterator[FlowEvent]:
        yield from self._run_add(alias, slug=slug)

    def run_add(self, alias: str, *, slug: str | None = None) -> Iterator[FlowEvent]:
        yield from self._run_add(alias, slug=slug)

    def add_account(self, alias: str, *, slug: str | None = None) -> Iterator[FlowEvent]:
        yield from self._run_add(alias, slug=slug)

    def _run_repair(self, account: Account) -> Iterator[FlowEvent]:
        session = None
        try:
            session = self._device_auth_flow.start(account.account_home, backup_existing=True)
        except Exception as exc:
            self._log_exception(exc)
            yield FlowEvent(
                kind="failure",
                operation="repair",
                account=account,
                message=self._failure_message(f"Re-authentication failed for {account.alias}", exc),
                error=exc,
            )
            return

        yield FlowEvent(
            kind="process_started",
            operation="repair",
            account=account,
            session=session,
        )

        try:
            prompt = self._device_auth_flow.read_prompt(session)
        except Exception as exc:
            self._log_exception(exc)
            self._abort_repair_session(session, account.account_home)
            yield FlowEvent(
                kind="rollback",
                operation="repair",
                account=account,
                session=session,
            )
            yield FlowEvent(
                kind="failure",
                operation="repair",
                account=account,
                session=session,
                message=self._failure_message(f"Re-authentication failed for {account.alias}", exc),
                error=exc,
            )
            return

        yield FlowEvent(
            kind="prompt_ready",
            operation="repair",
            account=account,
            session=session,
            prompt=prompt,
        )

        try:
            completed = self._device_auth_flow.await_completion(session)
        except Exception as exc:
            self._log_exception(exc)
            self._abort_repair_session(session, account.account_home)
            yield FlowEvent(
                kind="rollback",
                operation="repair",
                account=account,
                session=session,
            )
            yield FlowEvent(
                kind="failure",
                operation="repair",
                account=account,
                session=session,
                message=self._failure_message(f"Re-authentication failed for {account.alias}", exc),
                error=exc,
                retry_available=True,
            )
            return

        if completed:
            try:
                self._device_auth_flow.commit(account.account_home)
            except Exception as exc:
                self._log_exception(exc)
                self._abort_repair_session(session, account.account_home)
                yield FlowEvent(
                    kind="rollback",
                    operation="repair",
                    account=account,
                    session=session,
                )
                yield FlowEvent(
                    kind="failure",
                    operation="repair",
                    account=account,
                    session=session,
                    message=self._failure_message(f"Re-authentication failed for {account.alias}", exc),
                    error=exc,
                    retry_available=True,
                )
                return

            if self.registry.default_slug() == account.slug:
                self.registry.set_default(account)
            yield FlowEvent(
                kind="commit",
                operation="repair",
                account=account,
                session=session,
            )
            yield FlowEvent(
                kind="success",
                operation="repair",
                account=account,
                session=session,
            )
            return

        self._device_auth_flow.rollback(account.account_home)
        yield FlowEvent(
            kind="rollback",
            operation="repair",
            account=account,
            session=session,
        )
        yield FlowEvent(
            kind="failure",
            operation="repair",
            account=account,
            session=session,
            retry_available=True,
        )

    def _run_add(self, alias: str, *, slug: str | None = None) -> Iterator[FlowEvent]:
        existing_accounts = self._existing_accounts_for_add()
        try:
            if slug is None:
                slug = self.registry.new_slug(alias)
            else:
                slug = self.registry.normalize_slug(slug)
            codex_home = self.registry.add_dir(slug, alias)
        except Exception as exc:
            self._log_exception(exc)
            yield FlowEvent(
                kind="failure",
                operation="add",
                alias=alias,
                message=self._failure_message(f"Adding account {alias} failed", exc),
                error=exc,
                retry_available=False,
            )
            return

        account = self._account_by_slug(slug)

        try:
            session = self._device_auth_flow.start(codex_home, backup_existing=False)
        except Exception as exc:
            self._log_exception(exc)
            self.registry.remove(slug)
            yield FlowEvent(
                kind="rollback",
                operation="add",
                account=account,
                alias=alias,
                message=f"Removed new account {alias}",
                error=exc,
            )
            yield FlowEvent(
                kind="failure",
                operation="add",
                account=account,
                alias=alias,
                message=self._failure_message(f"Adding account {alias} failed", exc),
                error=exc,
                retry_available=False,
            )
            return

        yield FlowEvent(
            kind="process_started",
            operation="add",
            account=account,
            alias=alias,
            session=session,
        )

        try:
            prompt = self._device_auth_flow.read_prompt(session)
        except Exception as exc:
            self._log_exception(exc)
            terminate_process(session.process)
            self.registry.remove(slug)
            yield FlowEvent(
                kind="rollback",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                message=f"Removed new account {alias}",
                error=exc,
            )
            yield FlowEvent(
                kind="failure",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                message=self._failure_message(f"Adding account {alias} failed", exc),
                error=exc,
                retry_available=False,
            )
            return

        yield FlowEvent(
            kind="prompt_ready",
            operation="add",
            account=account,
            alias=alias,
            session=session,
            prompt=prompt,
            existing_accounts=existing_accounts,
        )

        try:
            completed = self._device_auth_flow.await_completion(session)
        except Exception as exc:
            self._log_exception(exc)
            terminate_process(session.process)
            self.registry.remove(slug)
            yield FlowEvent(
                kind="rollback",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                message=f"Removed new account {alias}",
                error=exc,
            )
            yield FlowEvent(
                kind="failure",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                message=self._failure_message(f"Adding account {alias} failed", exc),
                error=exc,
                retry_available=False,
            )
            return

        if not completed:
            self.registry.remove(slug)
            yield FlowEvent(
                kind="rollback",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                message=f"Removed new account {alias}",
            )
            yield FlowEvent(
                kind="failure",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                retry_available=False,
            )
            return

        collision = self._find_collision(account)
        if collision is not None:
            self.registry.remove(slug)
            yield FlowEvent(
                kind="collision",
                operation="add",
                account=account,
                alias=alias,
                session=session,
                collision=collision,
                message=f"Adding account {alias} collided with existing account {collision.alias}",
            )
            return

        yield FlowEvent(
            kind="success",
            operation="add",
            account=account,
            alias=alias,
            session=session,
        )

    def _existing_accounts_for_add(self) -> list[tuple[str, str]]:
        return [
            (account.alias, account.email)
            for account in self.registry.list()
            if account.email is not None
        ]

    def _find_collision(self, new_account: Account | None) -> Account | None:
        if new_account is None or not new_account.account_id:
            return None
        for account in self.registry.list():
            if account.slug == new_account.slug:
                continue
            if account.account_id == new_account.account_id:
                return account
        return None

    def _account_by_slug(self, slug: str) -> Account | None:
        for account in self.registry.list():
            if account.slug == slug:
                return account
        return None

    def _abort_repair_session(self, session: DeviceAuthSession, codex_home: Path) -> None:
        abort = getattr(self._device_auth_flow, "abort", None)
        if callable(abort):
            try:
                abort(session, codex_home)
                return
            except Exception as exc:
                self._log_exception(exc)
        terminate_process(session.process)
        self._device_auth_flow.rollback(codex_home)

    @staticmethod
    def _failure_message(prefix: str, exc: Exception) -> str:
        detail = str(exc)
        if not detail:
            return prefix
        return f"{prefix}: {detail}"

    @staticmethod
    def _log_exception(exc: Exception) -> None:
        traceback.print_exception(type(exc), exc, exc.__traceback__)