# type: ignore
"""
modules/local_llm.py
================================================================================
FALLBACK LOCAL LLM - ÚLTIMA HIPÓTASE
================================================================================
Este módulo é usado SOMENTE quando TODAS as APIs externas falharem.
Implementa um modelo local leve (TinyLlama ou equivalente) para respostas
básicas em modo de emergência.

Features:
- Fallback final do sistema
- Modelo pequeno (~1.5B parâmetros)
- Respostas básicas em português/angolano
- Não requer GPU
================================================================================
"""

import os
import re
import time
from typing import Optional, List, Dict, Any
from datetime import datetime

# Imports opcionais com fallbacks
try:
    import torch  # type: ignore
    TORCH_AVAILABLE = True
except Exception:
    TORCH_AVAILABLE = False
    torch = None  # type: ignore

import requests # type: ignore
try:
    from huggingface_hub import InferenceClient # type: ignore
    HUGGINGFACE_HUB_AVAILABLE = True
except Exception:
    HUGGINGFACE_HUB_AVAILABLE = False
    InferenceClient = None

try:
    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline  # type: ignore
    TRANSFORMERS_AVAILABLE = True
except Exception:
    TRANSFORMERS_AVAILABLE = False
    AutoTokenizer = None  # type: ignore
    AutoModelForCausalLM = None  # type: ignore
    pipeline = None  # type: ignore

try:
    from loguru import logger  # type: ignore
    LOGURU_AVAILABLE = True
except Exception:
    LOGURU_AVAILABLE = False
    # Criar logger dummy
    class DummyLogger:
        def info(self, *args, **kwargs): pass
        def success(self, *args, **kwargs): pass
        def warning(self, *args, **kwargs): pass
        def error(self, *args, **kwargs): pass
        def debug(self, *args, **kwargs): pass
    logger = DummyLogger()  # type: ignore

try:
    from cachetools import TTLCache  # type: ignore
    CACHETOOLS_AVAILABLE = True
except Exception:
    CACHETOOLS_AVAILABLE = False
    # Implementação simples de cache fallback
    class TTLCache(dict):
        def __init__(self, maxsize=10, ttl=300, **kwargs):
            super().__init__(**kwargs)
            self.maxsize = maxsize
            self.ttl = ttl
            self._timestamps = {}
        
        def __setitem__(self, key, value):
            super().__setitem__(key, value)
            self._timestamps[key] = time.time()
            # Limpa itens antigos se necessário
            if len(self) > self.maxsize:
                oldest_key = min(self._timestamps.keys(), key=lambda k: self._timestamps[k])
                self.pop(oldest_key, None)
                self._timestamps.pop(oldest_key, None)
        
        def get(self, key, default=None):
            # Verifica se expirou
            if key in self._timestamps:
                if time.time() - self._timestamps[key] > self.ttl:
                    self.pop(key, None)
                    self._timestamps.pop(key, None)
                    return default
            return super().get(key, default)

# Cache de prompts
_prompt_cache: Any = None
if CACHETOOLS_AVAILABLE:
    try:
        _prompt_cache = TTLCache(maxsize=10, ttl=300)
    except Exception:
        _prompt_cache = {}

# ============================================================
# 🎯 CONFIGURAÇÕES DO FALLBACK LOCAL
# ============================================================

# Modelos locais suportados (do mais leve ao mais pesado)
LOCAL_LLM_MODELS = [
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # ~1.1GB - Mais leve
    "microsoft/Phi-3-mini-4k-instruct",     # ~2.4GB
    "TheBloke/Llama-2-7b-chat-GGUF",        # ~4GB (precisa de conversão)
    "meta-llama/Llama-2-7b-chat-hf",        # ~13GB
]

# Caminhos locais verificados
LOCAL_MODEL_PATHS = [
    "/app/models/tinyllama-1.1b",
    "/app/models/phi-3-mini",
    "/app/models/llama-2-7b",
    "/models/tinyllama",
    "/models/phi-3",
]

# Prompt do sistema para fallback (mais simples)
FALLBACK_SYSTEM_PROMPT = """Você é Akira, uma IA assistente angolana.
Responda de forma curta e direta (1-2 frases).
Use português brasileiro com gírias angolanas quando natural.
Se não souber a resposta, diga que está em modo de emergência.
"""


# ============================================================
# 🏗️ CLASSE PRINCIPAL - LOCAL LLM FALLBACK
# ============================================================

class LocalLLMFallback:
    """
    Fallback local para quando TODAS as APIs externas falharem.
    Carrega um modelo pequeno (TinyLlama ~1.1B) que funciona em CPU.
    
    IMPORTANTE: Esta classe só deve ser usada como ÚLTIMA opção.
    """
    
    _instance = None
    _model_lock = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
            cls._instance._model_lock = __import__('threading').Lock()
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        
        # Componentes do modelo
        self._model = None  # type: ignore
        self._tokenizer = None  # type: ignore
        self._pipeline = None  # type: ignore
        self._model_path = None  # type: ignore
        self._is_loaded = False
        
        # Configurações
        self._max_tokens = 256  # Respostas curtas para CPU
        self._temperature = 0.7
        self._max_consecutive_failures = 3
        self._consecutive_failures = 0
        self._is_hf_inference_mode = False # Se true, a class usa inferência remota do HF Hub
        self._hf_client = None
        
        # Estatísticas
        self._stats = {
            "total_calls": 0,
            "successful_calls": 0,
            "failed_calls": 0,
            "last_used": None,
            "model_loaded": False
        }
        
        # Tenta detectar e carregar modelo
        self._detect_and_load_model()
    
    def _detect_and_load_model(self) -> bool:
        """Detecta e carrega modelo local se disponível, ou fallback para HF Inference Mode."""
        hf_token = os.getenv("HF_TOKEN") or getattr(__import__('modules.config', fromlist=['HF_TOKEN']), 'HF_TOKEN', None)
        
        # Testar Hugging Face API Mode se Transformers n existirem
        if (not TORCH_AVAILABLE or not TRANSFORMERS_AVAILABLE) and HUGGINGFACE_HUB_AVAILABLE and hf_token:
            logger.info("Torch/Transformers não disponíveis, porém `huggingface_hub` sim. Usando HuggingFace Inference via API.")
            self._is_loaded = True
            self._is_hf_inference_mode = True
            self._hf_client = InferenceClient(token=hf_token)
            self._model_path = LOCAL_LLM_MODELS[0] # Usa tinyllama via API
            self._stats["model_loaded"] = True
            return True
            
        if not TORCH_AVAILABLE or not TRANSFORMERS_AVAILABLE:
            if hf_token: # Se não tem lib hub, tenta usar request simples HTTP
                logger.info("Torch/Transformers não disponíveis, tentando usar Request simples HF.")
                self._is_loaded = True
                self._is_hf_inference_mode = True
                self._hf_client = "request_fallback"
                self._model_path = LOCAL_LLM_MODELS[0]
                self._stats["model_loaded"] = True
                return True
            else:
                logger.warning("Torch/Transformers e HF_TOKEN não disponíveis. Local LLM desabilitado.")
                return False
        
        with self._model_lock:
            if self._is_loaded:
                return True
            
            # Tenta encontrar modelo local
            model_path = self._find_local_model()
            
            if model_path:
                return self._load_model(model_path)
            
            logger.info("Nenhum modelo local encontrado. Local LLM desabilitado.")
            return False
    
    def _find_local_model(self) -> Optional[str]:
        """Procura modelo local em caminhos conhecidos."""
        # 1. Verifica variável de ambiente
        env_path = os.getenv("LOCAL_LLM_PATH")
        if env_path and os.path.exists(env_path):
            logger.info(f"Modelo local encontrado via env: {env_path}")
            return env_path
        
        # 2. Verifica caminhos locais
        for path in LOCAL_MODEL_PATHS:
            if os.path.exists(path):
                logger.info(f"Modelo local encontrado: {path}")
                return path
        
        # 3. Tenta descargar TinyLlama (pequeno, ~1.1GB)
        # Só faz download se explicitly habilitado
        if os.getenv("LOCAL_LLM_AUTO_DOWNLOAD", "").lower() == "true":
            logger.info("Auto-download habilitado. TinyLlama será baixado se necessário.")
            return LOCAL_LLM_MODELS[0]
        
        return None
    
    def _load_model(self, model_path: str) -> bool:
        """Carrega modelo local."""
        try:
            logger.info(f"🔄 Carregando modelo local: {model_path}")
            
            hf_token = os.getenv("HF_TOKEN")
            
            # Carrega tokenizer
            self._tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                token=hf_token,
                padding_side="left"
            )
            
            # Configura pad_token
            if self._tokenizer.pad_token is None:
                self._tokenizer.pad_token = self._tokenizer.eos_token
            
            # Carrega modelo (CPU apenas para compatibilidade)
            self._model = AutoModelForCausalLM.from_pretrained(
                model_path,
                token=hf_token,
                torch_dtype=torch.float32 if torch else None,
                low_cpu_mem_usage=True,
                device_map="auto" if TORCH_AVAILABLE else None
            )
            
            # Cria pipeline
            self._pipeline = pipeline(
                "text-generation",
                model=self._model,
                tokenizer=self._tokenizer,
                max_new_tokens=self._max_tokens,
                temperature=self._temperature,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.1
            )
            
            self._model_path = model_path
            self._is_loaded = True
            self._stats["model_loaded"] = True
            
            logger.success(f"✅ Modelo local carregado: {model_path}")
            return True
            
        except Exception as e:
            logger.error(f"❌ Erro ao carregar modelo local: {e}")
            self._is_loaded = False
            return False
    
    def is_available(self) -> bool:
        """Verifica se o fallback local está disponível."""
        if self._is_hf_inference_mode:
            return True
        return self._is_loaded and self._pipeline is not None
    
    def is_operational(self) -> bool:
        """Verifica se está operacional (pode responder)."""
        return self.is_available() and self._consecutive_failures < self._max_consecutive_failures
    
    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None
    ) -> Optional[str]:
        """
        Gera resposta usando modelo local.
        
        Args:
            prompt: Prompt do usuário
            system_prompt: Prompt do sistema (usa default se None)
            max_tokens: Máximo de tokens (usa default se None)
            temperature: Temperatura de geração
            
        Returns:
            String da resposta ou None se falhar
        """
        self._stats["total_calls"] += 1
        
        # Verifica disponibilidade
        if not self.is_operational():
            self._stats["failed_calls"] += 1
            return None
        
        # Usa cache se disponível
        cache_key = f"{prompt[:50]}:{system_prompt or 'default'}"
        if _prompt_cache is not None:
            cached = _prompt_cache.get(cache_key)
            if cached:
                logger.debug("Resposta encontrada em cache local")
                return cached
        
        try:
            # Prepara prompts
            sys_prompt = system_prompt or FALLBACK_SYSTEM_PROMPT
            
            # Formata para modelo
            if self._tokenizer and hasattr(self._tokenizer, 'chat_template') and False:
                # Usa chat template se disponível
                messages = [
                    {"role": "system", "content": sys_prompt},
                    {"role": "user", "content": prompt}
                ]
                formatted = self._tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True
                )
            else:
                # Formato simples (funciona com a maioria dos modelos)
                formatted = f"""<|system|>
{sys_prompt}
</s>
<|user|>
{prompt}
</s>
<|assistant|>
"""
            
            # Verifica e executa via Inference API/Request se _is_hf_inference_mode for True
            if getattr(self, '_is_hf_inference_mode', False):
                hf_token = os.getenv("HF_TOKEN") or getattr(__import__('modules.config', fromlist=['HF_TOKEN']), 'HF_TOKEN', None)
                if not hf_token:
                    logger.error("❌ Token HF não encontrado para a requisição de inferência HF")
                    return None
                    
                model_route = self._model_path or LOCAL_LLM_MODELS[0]
                max_new = max_tokens or self._max_tokens
                api_url = f"https://api-inference.huggingface.co/models/{model_route}"
                
                headers = {"Authorization": f"Bearer {hf_token}", "Content-Type": "application/json"}
                
                # TinyLlama format: <|system|>\n...\n<|user|>\n...\n<|assistant|>\n
                # ou Inst format dependendo do modelo local.
                # Como a maioria dos TinyLlamas na Hf API aceita o template de chat
                messages = [
                    {"role": "system", "content": sys_prompt},
                    {"role": "user", "content": prompt}
                ]
                
                payload = {
                    "model": model_route,
                    "messages": messages,
                    "max_tokens": max_new, 
                    "temperature": temperature or self._temperature
                }
                
                # Vamos tentar a HF Chat completions API que lida nativamente com as msg:
                api_url = f"https://api-inference.huggingface.co/models/{model_route}/v1/chat/completions"
                
                try: # Requests HTTP normais ou inference client
                    response = requests.post(api_url, headers=headers, json=payload, timeout=20)
                    response.raise_for_status()
                    
                    data = response.json()
                    
                    if isinstance(data, dict) and 'choices' in data and len(data['choices']) > 0:
                        res_text = data['choices'][0].get('message', {}).get('content', '')
                        res_text = self._clean_response(res_text)
                        if res_text:
                            # Cache
                            if _prompt_cache is not None:
                                try: _prompt_cache[cache_key] = res_text
                                except Exception: pass
                            
                            self._stats["successful_calls"] += 1
                            self._stats["last_used"] = datetime.now().isoformat()
                            self._consecutive_failures = 0
                            return res_text
                            
                except Exception as ex:
                    logger.error(f"❌ Falha no Hugging Face HTTP Inference fallback: {ex}")
                    self._consecutive_failures += 1
                    self._stats["failed_calls"] += 1
                    return None
                    
            # Se chegou aqui, gerar resposta na CPU / Local com Pipeline Local Pipeline Pipeline Normal  
            max_new = max_tokens or self._max_tokens
            
            outputs = self._pipeline(
                formatted,
                max_new_tokens=max_new,
                temperature=temperature or self._temperature,
                do_sample=True,
                pad_token_id=self._tokenizer.eos_token_id if self._tokenizer else None,
                return_full_text=False # MUITO IMPORTANTE: Tira o input do output
            )
            
            # Extrai resposta
            if outputs and len(outputs) > 0:
                generated = outputs[0].get("generated_text", "")
                
                # Remove prompt da resposta
                response = self._extract_response(generated, formatted)
                response = self._clean_response(response)
                
                if response:
                    # Cache se disponível
                    if _prompt_cache is not None:
                        try:
                            _prompt_cache[cache_key] = response
                        except Exception:
                            pass
                    
                    self._stats["successful_calls"] += 1
                    self._stats["last_used"] = datetime.now().isoformat()
                    self._consecutive_failures = 0
                    
                    return response
            
            # Falha silenciosa
            self._consecutive_failures += 1
            self._stats["failed_calls"] += 1
            return None
            
        except Exception as e:
            logger.error(f"❌ Erro em fallback local: {e}")
            self._consecutive_failures += 1
            self._stats["failed_calls"] += 1
            return None
    
    def _extract_response(self, generated: str, prompt: str) -> str:
        """Extrai a resposta do texto gerado."""
        if not generated:
            return ""
        
        # Remove o prompt do início
        if prompt in generated:
            response = generated[len(prompt):]
        else:
            # Tenta encontrar padrão de separação
            if "<|assistant|>" in generated:
                response = generated.split("<|assistant|>")[-1]
            elif " [/INST]" in generated:
                response = generated.split(" [/INST]")[-1]
            elif "</s>" in generated and "<|user|>" in generated:
                # Extrai após última tag de user
                parts = generated.split("<|user|>")
                if len(parts) > 1:
                    response = parts[-1]
                else:
                    response = generated
                    response = generated
            else:
                response = generated
        
        return response.strip()
    
    def _clean_response(self, text: str) -> str:
        """Limpa a resposta gerada."""
        # Remove tags e formatação
        text = re.sub(r'<\|[^|]+\|>', '', text)
        text = re.sub(r'</?s>', '', text)
        text = re.sub(r'[\*\_\`\[\]\"]', '', text)
        
        # Normaliza espaços
        text = re.sub(r'\s+', ' ', text).strip()
        
        # Limita tamanho (1 token ≈ 4 caracteres)
        max_chars = self._max_tokens * 4
        if len(text) > max_chars:
            # Corta em sentença completa
            sentences = [s.strip() + "." for s in text.split(".") if s.strip()]
            result = ""
            for sent in sentences:
                if len(result + sent) <= max_chars:
                    result += sent + " "
                else:
                    break
            text = result.strip()
        
        return text
    
    def get_status(self) -> Dict[str, Any]:
        """Retorna status do fallback local."""
        return {
            "available": self.is_available(),
            "operational": self.is_operational(),
            "model_path": self._model_path,
            "model_loaded": self._is_loaded,
            "consecutive_failures": self._consecutive_failures,
            "max_failures_allowed": self._max_consecutive_failures,
            "stats": self._stats.copy()
        }
    
    def reset_failures(self):
        """Reseta contador de falhas."""
        self._consecutive_failures = 0
    
    def should_use_fallback(self, api_failures: int = 0) -> bool:
        """
        Decide se deve usar o fallback local.
        
        Args:
            api_failures: Número de falhas consecutivas de APIs
            
        Returns:
            True se deve usar fallback
        """
        # Só usa se:
        # 1. Modelo está operacional
        # 2. Houve pelo menos 1 falha de API OU está explicitamente habilitado
        return (
            self.is_operational() and
            (api_failures > 0 or os.getenv("USE_LOCAL_FALLBACK", "").lower() == "true")
        )


# ============================================================
# 🎯 FUNÇÃO PRINCIPAL DE FALLBACK
# ============================================================

def get_local_fallback() -> LocalLLMFallback:
    """Retorna instância singleton do fallback local."""
    return LocalLLMFallback()


def generate_fallback_response(
    prompt: str,
    system_prompt: Optional[str] = None,
    api_failures: int = 0
) -> Optional[str]:
    """
    Gera resposta de fallback se necessário.
    
    Args:
        prompt: Prompt do usuário
        system_prompt: Prompt do sistema opcional
        api_failures: Número de falhas de API
        
    Returns:
        Resposta gerada ou None
    """
    fallback = get_local_fallback()
    
    if fallback.should_use_fallback(api_failures):
        logger.info(f"🔴 Usando fallback local (API failures: {api_failures})")
        return fallback.generate(prompt, system_prompt)
    
    return None


# ============================================================
# 🧪 MOCK PARA TESTES
# ============================================================

class MockLocalLLM:
    """Mock para testes quando modelo não está disponível."""
    
    def is_available(self) -> bool:
        return False
    
    def is_operational(self) -> bool:
        return False
    
    def generate(self, prompt: str, **kwargs) -> str:
        return "🤖 Modo de emergência: Todas as APIs falharam. Tente novamente mais tarde."
    
    def get_status(self) -> Dict[str, Any]:
        return {"available": False, "mock": True}


# ============================================================
# 📤 EXPORTS
# ============================================================

__all__ = [
    "LocalLLMFallback",
    "get_local_fallback",
    "generate_fallback_response",
    "MockLocalLLM",
    "FALLBACK_SYSTEM_PROMPT",
]