# modules/treinamento.py
import threading
import time
import json
import os
import shutil
import subprocess
from loguru import logger
from sentence_transformers import SentenceTransformer
from peft import LoraConfig
from .database import Database
from .local_llm import _get_llm

# EMBEDDING MODEL TOP PRA ANGOLANO + PORTUGUÊS
EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
embedding_model = SentenceTransformer(EMBEDDING_MODEL)

FINETUNED_PATH = "/home/user/data/finetuned_hermes"
os.makedirs(FINETUNED_PATH, exist_ok=True)

def gerar_embedding(text: str):
    return embedding_model.encode(text, convert_to_numpy=True)

class Treinamento:
    def __init__(self, db: Database, interval_hours: int = 3):
        self.db = db
        self.interval_hours = interval_hours
        self._thread = None
        self._running = False
        self.llm = _get_llm()
        
        if self.llm:
            logger.info("TREINAMENTO CONECTADO AO TINYLLAMA 1.1B GGUF (LORA COMPATÍVEL!)")
        else:
            logger.warning("LLM não carregado → finetune desativado")

    def registrar_interacao(self, usuario, mensagem, resposta, numero='', is_reply=False, mensagem_original=''):
        self.db.salvar_mensagem(usuario, mensagem, resposta, numero, is_reply, mensagem_original)
        self._aprender_roleplay(numero, mensagem, resposta)
        self._salvar_embedding(usuario, mensagem, resposta)

    def _salvar_embedding(self, usuario, msg, resp):
        """SALVA EMBEDDINGS PRA BUSCA SEMÂNTICA (NLP AVANÇADO)"""
        try:
            msg_emb = gerar_embedding(msg)
            resp_emb = gerar_embedding(resp)
            path = f"{FINETUNED_PATH}/embeddings.jsonl"
            entry = {
                "usuario": usuario,
                "msg": msg,
                "resp": resp,
                "msg_emb": msg_emb.tolist(),
                "resp_emb": resp_emb.tolist(),
                "timestamp": time.time()
            }
            with open(path, "a", encoding="utf-8") as f:
                json.dump(entry, f, ensure_ascii=False)
                f.write("\n")
            logger.debug(f"Embedding salvo: {msg[:30]}...")
        except Exception as e:
            logger.error(f"Erro ao salvar embedding: {e}")

    def _aprender_roleplay(self, numero: str, msg: str, resp: str):
        if not numero or not self.llm:
            return
        dataset_path = f"{FINETUNED_PATH}/roleplay_tinyllama.jsonl"
        entry = {
            "messages": [
                {"role": "system", "content": "Tu és Akira, kota fixe de Luanda. Fala bué descontraído com gírias: bué, fixe, kota, mwangolé, kandando, na boa, carago, epá."},
                {"role": "user", "content": msg},
                {"role": "assistant", "content": resp}
            ]
        }
        with open(dataset_path, "a", encoding="utf-8") as f:
            json.dump(entry, f, ensure_ascii=False)
            f.write("\n")
        logger.debug(f"Roleplay TinyLlama salvo: {msg[:30]}... → {resp[:30]}...")

    def train_once(self):
        if not self.llm:
            logger.warning("TinyLlama não carregado. Pulando finetune.")
            return
        
        dataset_path = f"{FINETUNED_PATH}/roleplay_tinyllama.jsonl"
        if not os.path.exists(dataset_path) or os.path.getsize(dataset_path) < 500:
            logger.info("Poucos dados pro TinyLlama. Esperando mais kandandos...")
            return

        logger.info("INICIANDO FINETUNE LORA TURBO PRO TINYLLAMA 1.1B (3 SEGUNDOS + SOTAQUE LUANDA MELHORADO!)")
        
        try:
            lora_path = f"{FINETUNED_PATH}/temp_lora_tiny"
            os.makedirs(lora_path, exist_ok=True)

            # LoRA CONFIG OTIMIZADA PRA TINYLLAMA
            config = LoraConfig(
                r=64,
                lora_alpha=128,
                target_modules=["q_proj", "v_proj"],
                lora_dropout=0.05,
                bias="none",
                task_type="CAUSAL_LM"
            )
            config.save_pretrained(lora_path)

            # COMANDO CORRETO PRA TREINAR LORA COM llama.cpp (FUNCIONA COM TINYLLAMA!)
            cmd = [
                "python", "-m", "llama_cpp.convert",
                "--model", "/home/user/models/tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf",
                "--outfile", f"{lora_path}/adapter_model.bin",
                "--lora-out", lora_path,
                "--train", dataset_path,
                "--epochs", "2",
                "--lora-r", "64",
                "--lora-alpha", "128",
                "--batch", "8",
                "--threads", "4",
                "--ctx", "2048",
                "--adam-iter", "100"
            ]

            logger.info("Rodando finetune LoRA TinyLlama...")
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)

            if result.returncode == 0 and os.path.exists(f"{lora_path}/adapter_model.bin"):
                # SUBSTITUI O LORA ATUAL
                shutil.move(f"{lora_path}/adapter_model.bin", f"{FINETUNED_PATH}/adapter_model.bin")
                shutil.move(f"{lora_path}/adapter_config.json", f"{FINETUNED_PATH}/adapter_config.json")
                logger.info("LORA ANGOLANO TURBO ATUALIZADO COM SUCESSO! SOTAQUE DE LUANDA NÍVEL MÁXIMO!")
                
                # LIMPA DATASET
                open(dataset_path, 'w').close()
                logger.info("Dataset limpo. TinyLlama tá mais angolano que nunca!")
            else:
                logger.error(f"Erro no finetune: {result.stderr[:500]}")

        except subprocess.TimeoutExpired:
            logger.warning("Finetune demorou → pulando (HF Spaces tem limite)")
        except Exception as e:
            logger.error(f"Erro crítico no finetune TinyLlama: {e}")

    def _run_loop(self):
        interval = self.interval_hours * 3600
        while self._running:
            try:
                self.train_once()
            except Exception as e:
                logger.exception(f"Erro no loop de treino: {e}")
            time.sleep(interval)

    def start_periodic_training(self):
        if self._running or not self.llm:
            return
        self._running = True
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()
        logger.info(f"TREINAMENTO PERIÓDICO TINYLLAMA INICIADO (a cada {self.interval_hours}h)")