"""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

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 = "~/.config/proxym/accounts.json"):
        self.config_path = Path(config_path).expanduser()
        self.accounts: dict[str, Account] = {}
        self._load()

    def _load(self) -> None:
        if self.config_path.exists():
            try:
                data = json.loads(self.config_path.read_text())
                for item in data.get("accounts", []):
                    acct = self._dict_to_account(item)
                    self.accounts[acct.id] = acct
                logger.info("accounts_loaded", count=len(self.accounts))
            except (json.JSONDecodeError, KeyError) as e:
                logger.error("accounts_load_failed", error=str(e))

    def _save(self) -> None:
        self.config_path.parent.mkdir(parents=True, exist_ok=True)
        data = {"accounts": [self._account_to_dict(a) for a in self.accounts.values()]}
        self.config_path.write_text(json.dumps(data, indent=2, default=str))

    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)
        ))

    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
