# agent_factory.py
from langchain_asi import ASI1ChatModel
from langchain.prompts import PromptTemplate
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.chat_history import InMemoryChatMessageHistory
from sentence_transformers import SentenceTransformer
from src.agents.document_utils import tiktoken_len
import numpy as np
import faiss
import os
import tiktoken
import dotenv
import pickle
from pathlib import Path
from langdetect import detect

dotenv.load_dotenv()

class HRManualAgent:
    """
    Se aplico un patron de diseño "singletone" para evitar crear mas 
    instancias de este objeto mientras el servidor esta en ejecucion,
    """
    _instance = None  

    def __new__(cls):
        if cls._instance is None:
            print("[INFO] Creating singleton instance of HRManualAgent")
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self._llm = ASI1ChatModel(model="asi1-mini", api_key=os.getenv("ASI1_API_KEY"))
        self._history = InMemoryChatMessageHistory()
        self._embedder = SentenceTransformer("all-MiniLM-L6-v2")
        self._prompt = PromptTemplate(
            input_variables=["context", "question", "chat_history", "language"],
            template = """
            Eres Daisy, la asistente del equipo de Recursos Humanos de Fr8App. Tu función es proporcionar respuestas claras, útiles y alineadas con las políticas internas de la empresa. Utilizas múltiples fuentes oficiales, incluyendo el Manual del Empleado, el Código de Conducta, la Política de Comisiones, la Política de Viajes y Gastos, y el Proceso de Compras. Solo compartes información relevante a partir del contexto recuperado dinámicamente por el sistema.

            IMPORTANTE: Responde en {language}.

            Fuentes disponibles:
            - Manual del Empleado 2024 (Employee Handbook)
            - Código de Conducta
            - Política de Comisiones 2025
            - Política de Viajes y Gastos
            - Política y Proceso de Compras

            El sistema solo te proporciona fragmentos relevantes extraídos mediante búsqueda semántica (FAISS). Si no tienes suficiente información para responder una pregunta, indica con claridad que no cuentas con datos suficientes.

            Contexto relevante recuperado:
            --------------------
            {context}

            Historial de la conversación:
            --------------------
            {chat_history}

            Pregunta:
            ---------
            {question}

            Responde de forma profesional como personal de Recursos Humanos:
            """
        )
        self._CHUNK_SIZE = 500
        self._CHUNK_OVERLAP = 20
        self._EMBEDDING_MODEL = "all-MiniLM-L6-v2"
        self._TOP_K = 8

        self._index = None
        self._chunk_lookup = None
        self._embedder = None
        self._initialized = True

    def create_agent(self):
        if self._index and self._chunk_lookup:
            print("[INFO] Usando índice FAISS en memoria")
        else:
            print("[INFO] Cargando agente RRHH desde disco o generando...")
            self._load_or_create_index()
        
        return {
            "llm": self._llm,
            "history": self._history,
            "prompt": self._prompt,
            "get_context": self._get_context
        }
    
    def _load_or_create_index(self):
        cache_dir = Path("data/cache")
        cache_dir.mkdir(parents=True, exist_ok=True)
        index_path = cache_dir / "hr_pdf.index"
        chunks_path = cache_dir / "hr_pdf.chunks.pkl"

        self._embedder = SentenceTransformer(self._EMBEDDING_MODEL)

        # Always process PDFs, ignore cache
        print("[INFO] Procesando PDFs y creando índice FAISS (HR)...")
        # Get all PDF files from data/raw directory
        raw_dir = Path("data/raw")
        pdf_files = list(raw_dir.glob("*.pdf"))
        
        if not pdf_files:
            raise ValueError("No se encontraron archivos PDF en data/raw")
        
        print(f"[INFO] Encontrados {len(pdf_files)} archivos PDF:")
        for pdf in pdf_files:
            print(f"  - {pdf.name}")
        
        # Process all PDFs
        all_docs = []
        for pdf_file in pdf_files:
            print(f"[INFO] Procesando {pdf_file.name}...")
            loader = PyMuPDFLoader(str(pdf_file))
            docs = loader.load()
            all_docs.extend(docs)
            print(f"[INFO] Extraídos {len(docs)} documentos de {pdf_file.name}")

        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self._CHUNK_SIZE,
            chunk_overlap=self._CHUNK_OVERLAP,
            length_function=tiktoken_len
        )
        chunks = splitter.split_documents(all_docs)
        texts = [chunk.page_content for chunk in chunks]
        vectors = self._embedder.encode(texts)

        dim = vectors[0].shape[0]
        self._index = faiss.IndexFlatL2(dim)
        self._index.add(np.array(vectors))
        self._chunk_lookup = {i: texts[i] for i in range(len(texts))}

        print(f"[INFO] Índice creado con {len(texts)} chunks")
        # Save the new index and chunks
        faiss.write_index(self._index, str(index_path))
        with open(chunks_path, "wb") as f:
            pickle.dump(self._chunk_lookup, f)

    def _get_context(self, query):
        try:
            # Detect the language of the query
            query_lang = detect(query)
            print(f"[INFO] Query language detected: {query_lang}")
            query_lang = "ENGLISH" if query_lang == "en" else "ESPAÑOL"
            q_vector = self._embedder.encode([query])
            scores, indices = self._index.search(np.array(q_vector), self._TOP_K)
            
            # Get the context and return it with language information
            context = "\n\n".join([self._chunk_lookup[i] for i in indices[0]])
            
            # Print the context with clear formatting
            print("\n[INFO] Contexto obtenido de FAISS:")
            print("=" * 80)
            print(context)
            print("=" * 80)
            
            return context, query_lang
        except Exception as e:
            print(f"[WARNING] Language detection failed: {e}")
            # Fallback to Spanish if language detection fails
            return context, "es"