import os
import time
from typing import List, Optional, Dict, Any
from dataclasses import dataclass, field
from loguru import logger

COOLDOWN_SECONDS = 60

ACCOUNT_NAMES = [
    "primary",
    "softedge",
    "mkultra",
]


@dataclass
class AccountQuota:
    key_index: int
    account_name: str
    api_key: str
    last_429_time: Optional[float] = None
    requests_today: int = 0
    last_reset: float = field(default_factory=time.time)
    is_exhausted: bool = False


class MistralAccountRotation:
    def __init__(self, api_keys: List[str]):
        self.api_keys = [k.strip() for k in api_keys if k and k.strip()]
        self.current_key_index = 0
        self.accounts: Dict[int, AccountQuota] = {}

        for i, key in enumerate(self.api_keys):
            account_name = ACCOUNT_NAMES[i] if i < len(ACCOUNT_NAMES) else f"account_{i}"
            self.accounts[i] = AccountQuota(
                key_index=i,
                account_name=account_name,
                api_key=key,
                requests_today=0,
            )

        self.logger = logger
        self._log_initialization()

    def _log_initialization(self):
        active_keys = len(self.api_keys)
        self.logger.success(f"✅ Mistral Rotation inicializado com {active_keys} conta(s):")
        for i, quota in self.accounts.items():
            status = "✅ ATIVA" if quota.api_key else "❌ VAZIA"
            self.logger.info(f"  [{i+1}] {quota.account_name.upper():<12} {status}")

    def get_current_key(self) -> Optional[str]:
        if not self.api_keys or self.current_key_index >= len(self.api_keys):
            return None
        return self.api_keys[self.current_key_index]

    def get_current_account_name(self) -> str:
        if not self.api_keys or self.current_key_index >= len(self.api_keys):
            return "unknown"
        if self.current_key_index < len(ACCOUNT_NAMES):
            return ACCOUNT_NAMES[self.current_key_index]
        return f"account_{self.current_key_index}"

    def _is_account_available(self, quota: AccountQuota) -> bool:
        if not quota.is_exhausted:
            return True
        if quota.last_429_time and (time.time() - quota.last_429_time) >= COOLDOWN_SECONDS:
            quota.is_exhausted = False
            self.logger.info(f"🔄 [Mistral Cooldown] Conta '{quota.account_name.upper()}' disponível novamente após {COOLDOWN_SECONDS}s")
            return True
        return False

    def handle_429_error(self) -> bool:
        quota = self.accounts[self.current_key_index]
        quota.last_429_time = time.time()
        quota.is_exhausted = True

        current_name = quota.account_name.upper()
        self.logger.warning(
            f"⚠️ [Mistral 429] Conta '{current_name}' (índice {self.current_key_index + 1}/{len(self.api_keys)}) esgotada. Buscando próxima chave..."
        )

        original_index = self.current_key_index
        for _ in range(len(self.api_keys) - 1):
            self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
            next_quota = self.accounts[self.current_key_index]
            if self._is_account_available(next_quota):
                next_name = next_quota.account_name.upper()
                self.logger.success(
                    f"✅ Mistral Rotation: mudando de '{current_name}' para '{next_name}' "
                    f"(índice {self.current_key_index + 1}/{len(self.api_keys)})"
                )
                return True

        self.logger.error(
            f"❌ Mistral Rotation: todas as {len(self.api_keys)} contas estão marcadas como esgotadas."
        )
        return False

    def reset_quotas_if_needed(self):
        now = time.time()
        reset_count = 0
        for quota in self.accounts.values():
            if now - quota.last_reset >= 24 * 3600:
                quota.requests_today = 0
                quota.is_exhausted = False
                quota.last_reset = now
                reset_count += 1
                self.logger.info(f"🔄 [Mistral Quota Reset] Conta '{quota.account_name.upper()}' resetada.")
        if reset_count:
            self.logger.success(f"✅ {reset_count} conta(s) Mistral resetada(s) e disponíveis.")

    def record_request(self):
        self.accounts[self.current_key_index].requests_today += 1

    def get_status(self) -> Dict[str, Any]:
        return {
            "current_account": self.get_current_account_name(),
            "current_index": self.current_key_index,
            "total_accounts": len(self.api_keys),
            "accounts": [
                {
                    "index": i + 1,
                    "name": quota.account_name.upper(),
                    "requests_today": quota.requests_today,
                    "exhausted": quota.is_exhausted,
                    "last_429": quota.last_429_time,
                }
                for i, quota in self.accounts.items()
            ],
        }


def _resolve_key(config, name: str) -> Optional[str]:
    for attr in (name, name.lower(), name.upper()):
        v = getattr(config, attr, None)
        if v and str(v).strip():
            return str(v).strip()
    for env in (name, name.lower(), name.upper()):
        v = os.getenv(env, "").strip()
        if v:
            if len(v) >= 2 and ((v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'"))):
                v = v[1:-1].strip()
            if v:
                return v
    return None

def get_mistral_rotation(config) -> Optional[MistralAccountRotation]:
    keys: List[str] = []
    missing: List[str] = []
    for name in ["MISTRAL_API_KEY", "SOFTEDGE_MISTRAL_API", "MKULTRA_MISTRAL_KEY"]:
        value = _resolve_key(config, name)
        if value:
            keys.append(value)
        else:
            missing.append(name)
    if missing:
        logger.warning(f"⚠️ Mistral keys missing: {', '.join(missing)} — check Secrets case (expected uppercase: {', '.join(missing)})")
        for n in missing:
            if n in ("SOFTEDGE_MISTRAL_API", "MKULTRA_MISTRAL_KEY"):
                logger.warning(f"  💡 Hint: {n} missing? check HF Secrets case — set as uppercase {n}, not {n.lower()}")
    if keys:
        logger.info(f"✅ Mistral Rotation: {len(keys)}/3 keys loaded ({', '.join([n for n in ['MISTRAL_API_KEY','SOFTEDGE_MISTRAL_API','MKULTRA_MISTRAL_KEY'] if _resolve_key(config, n)])})")
    if not keys:
        return None
    return MistralAccountRotation(keys)