"""Account manager: multi-provider credential vault with usage tracking.

Manages N accounts across M providers (Windsurf, Google, Anthropic, OpenRouter…).
Each account has:
  - Credentials (API keys, OAuth tokens, browser session cookies)
  - Usage metrics (tokens consumed, requests made, cost incurred)
  - Limits (daily/monthly caps from provider + user-defined)
  - Tool bindings (which IDE/tool uses which account)

Storage: Redis for live metrics, JSON file for persistent config.
"""

from __future__ import annotations

import json
import time
import uuid
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

import structlog

from proxym.observability import log_call
from proxym.storage import SQLiteNamespaceStore, is_sqlite_database

logger = structlog.get_logger()


class ProviderType(str, Enum):
    ANTHROPIC = "anthropic"
    OPENAI = "openai"
    OPENROUTER = "openrouter"
    GOOGLE = "google"
    DEEPSEEK = "deepseek"
    GROQ = "groq"
    MISTRAL = "mistral"
    WINDSURF = "windsurf"
    CURSOR = "cursor"
    GITHUB_COPILOT = "github_copilot"
    TOGETHER = "together"
    FIREWORKS = "fireworks"
    CEREBRAS = "cerebras"
    CUSTOM = "custom"


class ToolType(str, Enum):
    WINDSURF = "windsurf"
    CURSOR = "cursor"
    VSCODE = "vscode"
    JETBRAINS = "jetbrains"
    ROO_CODE = "roo_code"
    CLINE = "cline"
    CONTINUE_DEV = "continue_dev"
    AIDER = "aider"
    CLAUDE_CODE = "claude_code"
    BROWSER = "browser"          # generic browser-based tool
    CUSTOM = "custom"


@dataclass
class AccountCredentials:
    """Credentials for a single provider account."""
    api_key: str = ""
    oauth_token: str = ""
    refresh_token: str = ""
    session_cookie_path: str = ""  # path to browser profile/cookies
    extra: dict[str, str] = field(default_factory=dict)

    def to_safe_dict(self) -> dict:
        """Return dict with masked sensitive values."""
        d = {}
        if self.api_key:
            d["api_key"] = self.api_key[:8] + "..." + self.api_key[-4:] if len(self.api_key) > 12 else "***"
        if self.oauth_token:
            d["oauth_token"] = "***set***"
        if self.session_cookie_path:
            d["session_cookie_path"] = self.session_cookie_path
        return d


@dataclass
class UsageMetrics:
    """Tracked usage for a single account."""
    total_tokens_in: int = 0
    total_tokens_out: int = 0
    total_requests: int = 0
    total_cost_usd: float = 0.0
    daily_tokens_in: int = 0
    daily_tokens_out: int = 0
    daily_requests: int = 0
    daily_cost_usd: float = 0.0
    last_request_at: float = 0.0
    day_key: str = ""  # YYYY-MM-DD for daily reset

    def record(self, tokens_in: int, tokens_out: int, cost_usd: float) -> None:
        today = time.strftime("%Y-%m-%d")
        if today != self.day_key:
            self.daily_tokens_in = 0
            self.daily_tokens_out = 0
            self.daily_requests = 0
            self.daily_cost_usd = 0.0
            self.day_key = today

        self.total_tokens_in += tokens_in
        self.total_tokens_out += tokens_out
        self.total_requests += 1
        self.total_cost_usd += cost_usd
        self.daily_tokens_in += tokens_in
        self.daily_tokens_out += tokens_out
        self.daily_requests += 1
        self.daily_cost_usd += cost_usd
        self.last_request_at = time.time()


@dataclass
class AccountLimits:
    """Provider-imposed and user-defined limits."""
    # Provider limits (from subscription tier)
    provider_daily_requests: int = 0       # 0 = unlimited
    provider_monthly_tokens: int = 0
    provider_credits: float = 0.0          # e.g., Windsurf/Cursor credits
    provider_credits_remaining: float = 0.0

    # User-defined budget caps
    user_daily_budget_usd: float = 0.0     # 0 = unlimited
    user_monthly_budget_usd: float = 0.0
    user_max_request_cost: float = 0.0

    def is_over_budget(self, metrics: UsageMetrics) -> bool:
        if self.user_daily_budget_usd > 0 and metrics.daily_cost_usd >= self.user_daily_budget_usd:
            return True
        if self.user_monthly_budget_usd > 0 and metrics.total_cost_usd >= self.user_monthly_budget_usd:
            return True
        return False

    def budget_remaining(self, metrics: UsageMetrics) -> dict[str, float]:
        result = {}
        if self.user_daily_budget_usd > 0:
            result["daily_remaining_usd"] = max(0, self.user_daily_budget_usd - metrics.daily_cost_usd)
        if self.user_monthly_budget_usd > 0:
            result["monthly_remaining_usd"] = max(0, self.user_monthly_budget_usd - metrics.total_cost_usd)
        if self.provider_credits > 0:
            result["credits_remaining"] = self.provider_credits_remaining
        return result


@dataclass
class ToolBinding:
    """Binds an account to a specific tool/IDE instance."""
    tool_type: ToolType
    tool_name: str              # user-friendly name, e.g. "Windsurf Main"
    vm_id: str = ""             # libvirt VM identifier (empty = host)
    project_path: str = ""      # project directory on host
    browser_profile: str = ""   # chromium/firefox profile directory
    active: bool = True
    extra_config: dict[str, Any] = field(default_factory=dict)


@dataclass
class Account:
    """Single provider account with all associated data."""
    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    name: str = ""              # user-friendly name
    provider: ProviderType = ProviderType.CUSTOM
    email: str = ""
    plan_tier: str = ""         # "free", "pro", "team", "enterprise"

    credentials: AccountCredentials = field(default_factory=AccountCredentials)
    usage: UsageMetrics = field(default_factory=UsageMetrics)
    limits: AccountLimits = field(default_factory=AccountLimits)
    tools: list[ToolBinding] = field(default_factory=list)

    browser_profile_id: str = ""    # maps to a specific host browser profile
    tags: list[str] = field(default_factory=list)   # e.g., ["work", "personal"]
    created_at: float = field(default_factory=time.time)
    notes: str = ""

    @property
    def is_active(self) -> bool:
        return bool(self.credentials.api_key or self.credentials.oauth_token)

    @property
    def is_over_budget(self) -> bool:
        return self.limits.is_over_budget(self.usage)

    def status_summary(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "provider": self.provider.value,
            "email": self.email,
            "plan": self.plan_tier,
            "active": self.is_active,
            "over_budget": self.is_over_budget,
            "credentials": self.credentials.to_safe_dict(),
            "usage": {
                "daily_cost": round(self.usage.daily_cost_usd, 4),
                "total_cost": round(self.usage.total_cost_usd, 4),
                "daily_requests": self.usage.daily_requests,
                "total_requests": self.usage.total_requests,
            },
            "budget": self.limits.budget_remaining(self.usage),
            "tools": [
                {"type": t.tool_type.value, "name": t.tool_name,
                 "vm": t.vm_id or "host", "active": t.active}
                for t in self.tools
            ],
            "tags": self.tags,
        }


class AccountManager:
    """Central account registry.

    Persists to a JSON file, with live metrics in memory (optionally Redis).
    """

    def __init__(
        self,
        config_path: str | Path | None = None,
        *,
        config_dir: str | Path | None = None,
    ):
        if config_path is None:
            if config_dir is not None:
                config_path = Path(config_dir).expanduser() / "proxym.sqlite"
        self.config_path = Path(config_path).expanduser() if config_path is not None else None
        self._use_sqlite = self.config_path is not None
        self._store = (
            SQLiteNamespaceStore(self.config_path)
            if self._use_sqlite and self.config_path is not None
            else None
        )
        self.accounts: dict[str, Account] = {}
        self._load()

    def _load(self) -> None:
        if self.config_path is None or not self.config_path.exists():
            return

        if self._use_sqlite:
            if is_sqlite_database(self.config_path):
                data = self._store.load_namespace("accounts") if self._store else {}
                for item in data.values():
                    acct = self._dict_to_account(dict(item))
                    self.accounts[acct.id] = acct
                logger.info("accounts_loaded", count=len(self.accounts))
                return

            try:
                raw = json.loads(self.config_path.read_text())
                items = raw.get("accounts", []) if isinstance(raw, dict) else raw
                for item in items:
                    acct = self._dict_to_account(dict(item))
                    self.accounts[acct.id] = acct
                logger.info("accounts_loaded", count=len(self.accounts))
                if self.accounts:
                    self._save()
                return
            except (json.JSONDecodeError, KeyError, TypeError) as e:
                logger.error("accounts_load_failed", error=str(e))
                return

        try:
            raw = json.loads(self.config_path.read_text())
            items = raw.get("accounts", []) if isinstance(raw, dict) else raw
            for item in items:
                acct = self._dict_to_account(dict(item))
                self.accounts[acct.id] = acct
            logger.info("accounts_loaded", count=len(self.accounts))
            if self.accounts:
                self._save()
        except (json.JSONDecodeError, KeyError, TypeError) as e:
            logger.error("accounts_load_failed", error=str(e))

    def _save(self) -> None:
        if self.config_path is None:
            return
        if self._store is None:
            self._store = SQLiteNamespaceStore(self.config_path)
        self._store.save_namespace(
            "accounts",
            {acct.id: self._account_to_dict(acct) for acct in self.accounts.values()},
        )

    @log_call()
    def add_account(self, account: Account) -> Account:
        self.accounts[account.id] = account
        self._save()
        logger.info("account_added", id=account.id, provider=account.provider.value)
        return account

    def remove_account(self, account_id: str) -> bool:
        if account_id in self.accounts:
            del self.accounts[account_id]
            self._save()
            return True
        return False

    def get_account(self, account_id: str) -> Account | None:
        return self.accounts.get(account_id)

    def list_accounts(self, provider: ProviderType | None = None,
                      tag: str | None = None) -> list[Account]:
        results = list(self.accounts.values())
        if provider:
            results = [a for a in results if a.provider == provider]
        if tag:
            results = [a for a in results if tag in a.tags]
        return sorted(results, key=lambda a: a.name)

    def find_cheapest_active(self, provider: ProviderType) -> Account | None:
        """Find the cheapest active account for a provider (not over budget)."""
        candidates = [
            a for a in self.accounts.values()
            if a.provider == provider and a.is_active and not a.is_over_budget
        ]
        if not candidates:
            return None
        # Prefer accounts with most remaining budget
        return max(candidates, key=lambda a: (
            a.limits.budget_remaining(a.usage).get("daily_remaining_usd", 999)
        ))

    @log_call()
    def record_usage(self, account_id: str, tokens_in: int, tokens_out: int,
                     cost_usd: float) -> None:
        acct = self.accounts.get(account_id)
        if acct:
            acct.usage.record(tokens_in, tokens_out, cost_usd)
            # Periodic save (every 100 requests)
            if acct.usage.total_requests % 100 == 0:
                self._save()

    def bind_tool(self, account_id: str, binding: ToolBinding) -> bool:
        acct = self.accounts.get(account_id)
        if not acct:
            return False
        acct.tools.append(binding)
        self._save()
        return True

    def get_all_status(self) -> list[dict]:
        return [a.status_summary() for a in self.accounts.values()]

    def get_cost_summary(self) -> dict:
        """Aggregate cost summary across all accounts."""
        total_daily = sum(a.usage.daily_cost_usd for a in self.accounts.values())
        total_monthly = sum(a.usage.total_cost_usd for a in self.accounts.values())
        total_requests = sum(a.usage.total_requests for a in self.accounts.values())
        by_provider: dict[str, float] = {}
        for a in self.accounts.values():
            p = a.provider.value
            by_provider[p] = by_provider.get(p, 0) + a.usage.total_cost_usd
        return {
            "daily_total_usd": round(total_daily, 4),
            "monthly_total_usd": round(total_monthly, 4),
            "total_requests": total_requests,
            "by_provider": {k: round(v, 4) for k, v in sorted(by_provider.items())},
            "accounts_count": len(self.accounts),
            "accounts_active": sum(1 for a in self.accounts.values() if a.is_active),
            "accounts_over_budget": sum(1 for a in self.accounts.values() if a.is_over_budget),
        }

    @staticmethod
    def _account_to_dict(acct: Account) -> dict:
        d = asdict(acct)
        d["provider"] = acct.provider.value
        d["tools"] = [
            {**asdict(t), "tool_type": t.tool_type.value}
            for t in acct.tools
        ]
        return d

    @staticmethod
    def _from_dict(cls, d: dict):
        """Create a dataclass instance from a dict, ignoring unknown keys."""
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})

    @staticmethod
    def _dict_to_account(d: dict) -> Account:
        tools_data = d.pop("tools", [])
        creds = d.pop("credentials", {})
        usage = d.pop("usage", {})
        limits = d.pop("limits", {})
        d["provider"] = ProviderType(d.get("provider", "custom"))

        fd = AccountManager._from_dict
        acct = fd(Account, d)
        acct.credentials = fd(AccountCredentials, creds)
        acct.usage = fd(UsageMetrics, usage)
        acct.limits = fd(AccountLimits, limits)
        for td in tools_data:
            td["tool_type"] = ToolType(td.get("tool_type", "custom"))
            acct.tools.append(fd(ToolBinding, td))
        return acct
