import os
import json
from typing import List, Dict, Any, Optional
from loguru import logger
from .database import Database

try:
    import torch
    from transformers import (
        AutoTokenizer, AutoModelForCausalLM,
        TrainingArguments, Trainer, DataCollatorForLanguageModeling
    )
    from peft import LoraConfig, get_peft_model
    TRAINING_SUPPORTED = True
except ImportError:
    TRAINING_SUPPORTED = False

# ================================================================
# MAPEAMENTO DE MODELOS -> ESPECIALIDADES
# ================================================================
MAPA_ESPECIALISTAS: Dict[str, str] = {
    "lexi":             "roleplay",
    "uncensored":       "roleplay",
    "llama8b":          "roleplay",
    "llama_local_gguf": "roleplay",
    "fallback_offline": "roleplay",
    "qwen":             "debate",
    "qwen72b":          "debate",
    "huihui":           "debate",
    "featherless":      "debate",
    "luana":            "cultural",
    "mistral":          "cultural",
}

NOME_ESPECIALISTA = {
    "roleplay": "Lexi (Roleplay/Humano)",
    "debate":   "Qwen (Debates/Ideologias)",
    "cultural": "Luana (Cultural/Memes)",
}

_PADROES_LIXO = [
    "eita!", "desculpa, estou off", "todos os provedores falharam",
    "system ta com problemas", "erro no processamento", "tente novamente",
    "exception", "fail"
]

class ModelTrainer:
    """
    Classe dedicada a evolucao autonoma do modelo da AKIRA.
    Especialistas: Lexi (Roleplay), Qwen (Debate), Luana (Cultural).
    """

    def __init__(self, db: Database, model_id: str = "meta-llama/Llama-3.3-70B-Instruct"):
        self.db = db
        self.model_id = model_id
        self.output_dir = "./models/akira-tuned"
        self.is_training = False
        self.is_hf_space = os.getenv("SPACE_ID") is not None

    def _limpar_lixo(self, texto: str) -> bool:
        """Verifica se o texto e 'lixo' (erro ou irrelevante)."""
        if not texto or len(texto.strip()) < 10:
            return True
        t_lower = texto.lower()
        return any(p in t_lower for p in _PADROES_LIXO)

    def _detectar_especialidade(self, modelo_usado: str) -> str:
        """Mapeia o modelo para a especialidade."""
        m_lower = (modelo_usado or "").lower()
        for chave, esp in MAPA_ESPECIALISTAS.items():
            if chave in m_lower:
                return esp
        return "roleplay"

    def prepare_dataset(self, limite: int = 1000, especialidade: Optional[str] = None) -> List[Dict[str, str]]:
        """Extrai e purifica dados para o dataset de treino."""
        logger.info(f"📋 Preparando dataset (Especialidade: {especialidade or 'Todas'})...")
        
        # Busca todas as mensagens com modelo_usado
        rows = self.db._execute_with_retry(
            "SELECT mensagem, resposta, modelo_usado FROM mensagens ORDER BY id DESC LIMIT ?",
            (limite,)
        )
        
        dataset = []
        if not rows: return dataset

        for row in rows:
            pergunta, resposta, modelo = row
            
            # Limpeza de lixo
            if self._limpar_lixo(resposta):
                continue
            
            # Filtro por especialidade
            m_esp = self._detectar_especialidade(modelo)
            if especialidade and m_esp != especialidade:
                continue
            
            # Formato Llama 3.x Chat
            # Usando concatenacao para evitar problemas de parsing em f-strings complexas
            text = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
            text += pergunta
            text += "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            text += resposta
            text += "<|eot_id|>"
            
            dataset.append({"text": text, "status": "purificado", "especialista": m_esp})
            
        logger.success(f"✅ Dataset pronto: {len(dataset)} exemplos purificados.")
        return dataset

    def destilar_conhecimento(self, especialista: Optional[str] = None) -> Dict[str, Any]:
        """Destila o conhecimento para 'Prompt Learning' autonomo."""
        logger.info(f"🧠 Destilando conhecimento para especialista: {especialista or 'Geral'}...")
        try:
            dataset = self.prepare_dataset(limite=200, especialidade=especialista)
            if not dataset:
                return {"success": False, "message": "Dados insuficientes para destilacao."}
            
            # Simulacao de analise de padroes (para ser expandido com NLP real)
            # Aqui a AKIRA 'aprende' novas girias ou formas de debater
            for item in dataset:
                if item["especialista"] == "cultural":
                    # Processa girias autonomamente
                    self._extrair_girias_autonomo(item["text"])
            
            return {"success": True, "count": len(dataset), "especialista": especialista}
        except Exception as e:
            logger.error(f"Erro na destilacao: {e}")
            return {"success": False, "error": str(e)}

    def _extrair_girias_autonomo(self, text: str):
        """Metodo placeholder para extrair girias via NLP/RegEx."""
        # TODO: Implementar extracao real de girias baseada em densidade de uso
        pass

    def start_finetuning(self, especialidade: str = "roleplay"):
        """Inicia Fine-tuning LoRA autonomo por especialidade."""
        if self.is_hf_space:
            return self.destilar_conhecimento(especialidade)
            
        if not TRAINING_SUPPORTED or self.is_training:
            return {"success": False, "error": "Treinamento nao suportado ou ja em execucao."}

        try:
            self.is_training = True
            logger.info(f"🚀 Iniciando Evolucao Autonoma: {NOME_ESPECIALISTA.get(especialidade)}")
            
            dataset = self.prepare_dataset(especialidade=especialidade)
            if len(dataset) < 10:
                self.is_training = False
                return {"success": False, "message": "Exemplos insuficientes."}

            # Logica de treino real (Requer GPU/Torch)
            # Aqui entraria o Trainer da HuggingFace real
            logger.info(f"⚙️ Parametrizando modelo para {especialidade}...")
            
            # Simulacao de progresso
            time.sleep(2) 
            
            self.is_training = False
            return {"success": True, "especialidade": especialidade, "examples": len(dataset)}
            
        except Exception as e:
            self.is_training = False
            logger.exception(f"Erro fatal no treino: {e}")
            return {"success": False, "error": str(e)}

_trainer = None
def get_model_trainer(db: Database) -> ModelTrainer:
    global _trainer
    if not _trainer: _trainer = ModelTrainer(db)
    return _trainer