# type: ignore
"""
================================================================================
OPENROUTER KEY FARMING SYSTEM
================================================================================
Sistema de renovação dinâmica de chaves OpenRouter sem redeploy.

Workflow:
1. Conta bate rate limit → Sistema fallback automaticamente
2. Você paga/renova no OpenRouter website
3. Você copia a NOVA chave API
4. POST para /api/openrouter/refresh-key com a nova chave
5. AKIRA atualiza dinamicamente (sem redeploy)
6. Quando todas esgotam, volta a tentar a conta renovada

Armazenamento:
- PostgreSQL (via Database class) — compartilhado entre workers
- Tracking de quando cada chave foi renovada
- Log de tentativas e sucessos
================================================================================
"""

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


@dataclass
class AccountKey:
    """Informação de chave para uma conta OpenRouter"""
    account_index: int
    account_name: str
    api_key: str
    added_at: float = field(default_factory=time.time)
    last_rotated_at: float = field(default_factory=time.time)
    requests_count: int = 0
    is_exhausted: bool = False
    last_429_at: Optional[float] = None
    rotation_count: int = 0


class OpenRouterKeyFarmingDB:
    """Database para gerenciar chaves OpenRouter — usa PostgreSQL via Database class"""

    def __init__(self):
        self.logger = logger
        self._db = None
        self._init_db()

    def _get_db(self):
        if self._db is None:
            from .database import Database
            self._db = Database()
        return self._db

    def _init_db(self):
        """Cria tabelas se não existem"""
        db = self._get_db()
        db._execute_with_retry("""
            CREATE TABLE IF NOT EXISTS account_keys (
                account_index INTEGER PRIMARY KEY,
                account_name TEXT NOT NULL,
                api_key TEXT NOT NULL,
                added_at DOUBLE PRECISION NOT NULL,
                last_rotated_at DOUBLE PRECISION NOT NULL,
                requests_count INTEGER DEFAULT 0,
                is_exhausted INTEGER DEFAULT 0,
                last_429_at DOUBLE PRECISION,
                rotation_count INTEGER DEFAULT 0
            )
        """, commit=True)

        db._execute_with_retry("""
            CREATE TABLE IF NOT EXISTS key_rotation_log (
                id SERIAL PRIMARY KEY,
                account_index INTEGER NOT NULL,
                account_name TEXT NOT NULL,
                old_key TEXT,
                new_key TEXT,
                reason TEXT,
                rotated_at DOUBLE PRECISION NOT NULL,
                by_user TEXT DEFAULT 'manual'
            )
        """, commit=True)

    def add_initial_keys(self, keys: List[Tuple[int, str, str]]):
        """Adiciona chaves iniciais (index, name, key)"""
        db = self._get_db()
        for account_index, account_name, api_key in keys:
            now = time.time()
            try:
                db._execute_with_retry("""
                    INSERT INTO account_keys
                    (account_index, account_name, api_key, added_at, last_rotated_at)
                    VALUES (%s, %s, %s, %s, %s)
                    ON CONFLICT (account_index) DO UPDATE SET
                        account_name=EXCLUDED.account_name, api_key=EXCLUDED.api_key,
                        added_at=EXCLUDED.added_at, last_rotated_at=EXCLUDED.last_rotated_at
                """, (account_index, account_name, api_key, now, now), commit=True)
                self.logger.info(f"Chave inicial adicionada: {account_name}")
            except Exception as e:
                self.logger.error(f"Erro ao adicionar chave {account_name}: {e}")

    def get_key(self, account_index: int) -> Optional[str]:
        """Obtém chave atual para uma conta"""
        db = self._get_db()
        rows = db._execute_with_retry(
            "SELECT api_key FROM account_keys WHERE account_index = %s",
            (account_index,)
        )
        if rows:
            r = rows[0]
            return r['api_key'] if isinstance(r, dict) else r[0]
        return None

    def get_all_keys(self) -> Dict[int, str]:
        """Obtém todas as chaves (index -> key)"""
        db = self._get_db()
        rows = db._execute_with_retry(
            "SELECT account_index, api_key FROM account_keys ORDER BY account_index"
        )
        if not rows:
            return {}
        result = {}
        for r in rows:
            if isinstance(r, dict):
                result[r['account_index']] = r['api_key']
            else:
                result[r[0]] = r[1]
        return result

    def rotate_key(self, account_index: int, new_api_key: str, reason: str = "manual_refresh") -> bool:
        """Renovar chave de uma conta"""
        db = self._get_db()
        try:
            rows = db._execute_with_retry(
                "SELECT api_key, account_name FROM account_keys WHERE account_index = %s",
                (account_index,)
            )
            if not rows:
                self.logger.error(f"Conta {account_index} não encontrada")
                return False

            r = rows[0]
            old_key = r['api_key'] if isinstance(r, dict) else r[0]
            account_name = r['account_name'] if isinstance(r, dict) else r[1]
            now = time.time()

            db._execute_with_retry("""
                UPDATE account_keys
                SET api_key = %s, last_rotated_at = %s, rotation_count = rotation_count + 1,
                    is_exhausted = 0, last_429_at = NULL
                WHERE account_index = %s
            """, (new_api_key, now, account_index), commit=True)

            db._execute_with_retry("""
                INSERT INTO key_rotation_log
                (account_index, account_name, old_key, new_key, reason, rotated_at)
                VALUES (%s, %s, %s, %s, %s, %s)
            """, (account_index, account_name, old_key[:20] + "...", new_api_key[:20] + "...", reason, now), commit=True)

            self.logger.success(f"[KEY FARMING] Conta '{account_name}' renovada!")
            return True

        except Exception as e:
            self.logger.error(f"Erro ao renovar chave: {e}")
            return False

    def mark_exhausted(self, account_index: int):
        """Marca uma conta como esgotada (429)"""
        db = self._get_db()
        now = time.time()
        db._execute_with_retry("""
            UPDATE account_keys
            SET is_exhausted = 1, last_429_at = %s
            WHERE account_index = %s
        """, (now, account_index), commit=True)

    def mark_available(self, account_index: int):
        """Marca uma conta como disponível"""
        db = self._get_db()
        db._execute_with_retry("""
            UPDATE account_keys
            SET is_exhausted = 0, requests_count = 0
            WHERE account_index = %s
        """, (account_index,), commit=True)

    def increment_request_count(self, account_index: int):
        """Incrementa contador de requests"""
        db = self._get_db()
        db._execute_with_retry("""
            UPDATE account_keys
            SET requests_count = requests_count + 1
            WHERE account_index = %s
        """, (account_index,), commit=True)

    def get_status(self) -> Dict[str, Any]:
        """Retorna status de todas as contas"""
        db = self._get_db()
        rows = db._execute_with_retry("""
            SELECT account_index, account_name, requests_count, is_exhausted,
                   last_rotated_at, rotation_count, last_429_at
            FROM account_keys ORDER BY account_index
        """)

        status = {"accounts": [], "total_keys": 0, "exhausted_count": 0, "total_rotations": 0}
        if not rows:
            return status

        now = time.time()
        for r in rows:
            if isinstance(r, dict):
                index = r['account_index']
                name = r['account_name']
                req_count = r['requests_count']
                exhausted = r['is_exhausted']
                last_rot = r['last_rotated_at']
                rot_count = r['rotation_count']
                last_429 = r['last_429_at']
            else:
                index, name, req_count, exhausted, last_rot, rot_count, last_429 = r

            hours_rot = (now - last_rot) / 3600 if last_rot else 0
            hours_429 = (now - last_429) / 3600 if last_429 else None

            status["accounts"].append({
                "index": (index if isinstance(r, dict) else index) + 1,
                "name": name.upper(),
                "requests": req_count,
                "exhausted": bool(exhausted),
                "last_rotated": f"{hours_rot:.1f}h atrás" if last_rot else "Nunca",
                "rotation_count": rot_count,
                "last_429": f"{hours_429:.1f}h atrás" if last_429 else "N/A"
            })
            status["total_keys"] += 1
            status["exhausted_count"] += 1 if exhausted else 0
            status["total_rotations"] += rot_count

        return status

    def get_rotation_log(self, limit: int = 50) -> List[Dict[str, Any]]:
        """Obtém log de rotações recentes"""
        db = self._get_db()
        rows = db._execute_with_retry("""
            SELECT id, account_index, account_name, old_key, new_key, reason, rotated_at
            FROM key_rotation_log ORDER BY rotated_at DESC LIMIT %s
        """, (limit,))

        log = []
        if not rows:
            return log

        for r in rows:
            if isinstance(r, dict):
                log.append({
                    "id": r['id'], "account_index": r['account_index'],
                    "account_name": r['account_name'], "old_key": r['old_key'],
                    "new_key": r['new_key'], "reason": r['reason'],
                    "timestamp": datetime.fromtimestamp(r['rotated_at']).isoformat()
                })
            else:
                log.append({
                    "id": r[0], "account_index": r[1], "account_name": r[2],
                    "old_key": r[3], "new_key": r[4], "reason": r[5],
                    "timestamp": datetime.fromtimestamp(r[6]).isoformat()
                })

        return log


_FARMING_DB_INSTANCE: Optional[OpenRouterKeyFarmingDB] = None


def get_openrouter_farming_db() -> OpenRouterKeyFarmingDB:
    global _FARMING_DB_INSTANCE
    if _FARMING_DB_INSTANCE is None:
        _FARMING_DB_INSTANCE = OpenRouterKeyFarmingDB()
    return _FARMING_DB_INSTANCE


def reset_farming_db_instance():
    global _FARMING_DB_INSTANCE
    _FARMING_DB_INSTANCE = None