#!/usr/bin/env python3
from __future__ import annotations

import inspect
import fcntl
import json
import os
import sys
import tempfile
from pathlib import Path
from typing import IO

from account_registry import AccountRegistry, AccountRegistryKind
from claude_auth_operation import ClaudeAuthOperation
from claude_health_client import ClaudeHealthClient
from device_auth import DeviceAuthFlow
from device_auth_operation import DeviceAuthOperation
from grok_auth_operation import GrokAuthOperation
from grok_health_client import GrokHealthClient
from health_client import AccountHealthClient
from indicator import Indicator
from provider_services import (
    CodexHealthAdapter,
    CodexLifecycleAdapter,
    ProviderServiceMap,
    ProviderServices,
)
from runtime_paths import runtime_dir


TRAY_DIR = runtime_dir()
TRAY_LOCK_NAME = "codex-account-switcher.lock"
def default_routing_rules(default_slug: str) -> dict:
    return {
        "projects": {},
        "default": default_slug,
        "fallback_chain": [],
        "fallback_trigger": "broken_or_quota_exhausted",
        "quota_exhausted_threshold_pct": 100,
        "account_caps": {},
    }


class RoutingRulesSetupRequiredError(RuntimeError):
    pass


def _runtime_tray_dir() -> Path:
    try:
        TRAY_DIR.mkdir(parents=True, exist_ok=True)
        return TRAY_DIR
    except OSError:
        fallback = Path(tempfile.gettempdir()) / "systray-ai"
        fallback.mkdir(parents=True, exist_ok=True)
        return fallback


def _account_registry_kwargs(kind: AccountRegistryKind) -> dict[str, object]:
    kwargs: dict[str, object] = {"kind": kind}
    if isinstance(AccountRegistry, type) and AccountRegistry.__module__ == "account_registry":
        parameters = inspect.signature(AccountRegistry).parameters
        if "base_dir" in parameters:
            kwargs["base_dir"] = _runtime_tray_dir()
        if kind == AccountRegistryKind.CODEX and "pi_agent_dir" in parameters:
            kwargs["pi_agent_dir"] = Path.home() / ".pi" / "agent"
    return kwargs


def _atomic_write(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        "w",
        encoding="utf-8",
        dir=path.parent,
        delete=False,
    ) as handle:
        handle.write(content)
        temp_path = Path(handle.name)
    temp_path.replace(path)


def seed_routing_rules(
    path: Path, known_slugs: set[str], default_slug: str | None = None
) -> None:
    if path.exists():
        return
    if not known_slugs:
        raise RoutingRulesSetupRequiredError(
            f"{path.name} setup required: no accounts registered"
        )
    if default_slug is None or default_slug not in known_slugs:
        default_slug = sorted(known_slugs)[0]
    _atomic_write(path, json.dumps(default_routing_rules(default_slug), indent=2) + "\n")


def acquire_single_instance_lock(path: Path | None = None) -> IO[str]:
    lock_path = path if path is not None else _runtime_tray_dir() / TRAY_LOCK_NAME
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    handle = lock_path.open("w")
    try:
        fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        handle.close()
        raise RuntimeError("systray-ai is already running")
    return handle


def _build_provider_services(
    registry: AccountRegistry,
    health_client: AccountHealthClient,
    claude_registry: AccountRegistry | None,
    grok_registry: AccountRegistry | None = None,
) -> tuple[ProviderServiceMap, DeviceAuthFlow, DeviceAuthOperation]:
    device_auth_flow = DeviceAuthFlow()
    device_auth_operation = DeviceAuthOperation(registry, device_auth_flow)
    providers = [
        ProviderServices(
            tool=AccountRegistryKind.CODEX,
            registry=registry,
            lifecycle=CodexLifecycleAdapter(device_auth_operation),
            health=CodexHealthAdapter(health_client),
        )
    ]
    if claude_registry is not None:
        providers.append(
            ProviderServices(
                tool=AccountRegistryKind.CLAUDE,
                registry=claude_registry,
                lifecycle=ClaudeAuthOperation(claude_registry),
                health=ClaudeHealthClient(),
            )
        )
    if grok_registry is not None:
        providers.append(
            ProviderServices(
                tool=AccountRegistryKind.GROK,
                registry=grok_registry,
                lifecycle=GrokAuthOperation(grok_registry),
                health=GrokHealthClient(),
            )
        )
    return ProviderServiceMap(providers), device_auth_flow, device_auth_operation


def main() -> None:
    os.environ.pop("CODEX_HOME", None)
    runtime_tray_dir = _runtime_tray_dir()
    try:
        _lock_handle = acquire_single_instance_lock()  # noqa: F841 - kept alive to hold the flock
    except RuntimeError as exc:
        print(f"systray_codex_switcher: {exc}", file=sys.stderr)
        sys.exit(1)

    registry = AccountRegistry(**_account_registry_kwargs(AccountRegistryKind.CODEX))
    registry.migrate_legacy()
    try:
        known_slugs = {account.slug for account in registry.list()}
        seed_routing_rules(
            runtime_tray_dir / "routing_rules.json",
            known_slugs=known_slugs,
            default_slug=registry.default_slug(),
        )
    except RoutingRulesSetupRequiredError as exc:
        print(f"systray_codex_switcher: {exc}", file=sys.stderr)
        sys.exit(1)
    claude_registry = AccountRegistry(**_account_registry_kwargs(AccountRegistryKind.CLAUDE))
    claude_registry.migrate_legacy()
    try:
        seed_routing_rules(
            runtime_tray_dir / "claude_routing_rules.json",
            known_slugs={account.slug for account in claude_registry.list()},
            default_slug=claude_registry.default_slug(),
        )
    except RoutingRulesSetupRequiredError as exc:
        print(f"systray_codex_switcher: {exc}", file=sys.stderr)
    grok_registry = AccountRegistry(**_account_registry_kwargs(AccountRegistryKind.GROK))
    # Pin current live Grok session as roy-grok (copy-only; never refreshes tokens).
    try:
        GrokAuthOperation(grok_registry).import_from_legacy_home(
            slug="roy-grok",
            alias="roy-grok",
        )
    except Exception as exc:
        print(f"systray_codex_switcher: grok bootstrap skipped: {exc}", file=sys.stderr)
    health_client = AccountHealthClient()
    provider_services, device_auth_flow, device_auth_operation = _build_provider_services(
        registry,
        health_client,
        claude_registry,
        grok_registry,
    )
    indicator = Indicator(
        registry,
        health_client,
        claude_registry=claude_registry,
        grok_registry=grok_registry,
        provider_services=provider_services,
        device_auth_flow=device_auth_flow,
        device_auth_operation=device_auth_operation,
    )
    indicator.build()
    indicator.run()


if __name__ == "__main__":
    main()
