import os
from typing import List, Dict
from langchain_community.vectorstores import FAISS
from groq import Groq
import argparse
import json
from dotenv import load_dotenv
from app.config import config, require_env
from modules.processing.embeddings import get_hf_embeddings
load_dotenv()
GROQ_API_KEY = require_env(os.getenv("GROQ_API_KEY") or config.GROQ_API_KEY, "GROQ_API_KEY")
EMBEDDINGS_MODEL = os.getenv('EMBEDDINGS_MODEL', 'intfloat/e5-small-v2')

class SOATRetriever:
    def __init__(self, faiss_index_path: str, groq_api_key: str, groq_model: str = None):
        self.embeddings = get_hf_embeddings(EMBEDDINGS_MODEL)  # "google/embeddinggemma-300m"
        self.db = FAISS.load_local(faiss_index_path, self.embeddings, allow_dangerous_deserialization=True)
        self.groq = Groq(api_key=groq_api_key)
        self.groq_model = groq_model or config.GROQ_MODEL_SOAT

    def buscar(self, consulta: str, k: int = 20):
        return self.db.max_marginal_relevance_search(consulta, k=k, lambda_param=0.7)

    def proponer_codigos_soat(self, descripcion: str, docs) -> str:
        contexto = "\n\n".join(
            f"(p.{d.metadata.get('page')}) {d.page_content[:1200]}" for d in docs
        )
        prompt = f"""
Eres un especialista en auditoría médica en Colombia. Basado EXCLUSIVAMENTE en el Manual SOAT 2024 siguiente (contexto), propone códigos SOAT válidos para el procedimiento:

Descripción del procedimiento: {descripcion}

Contexto (Manual SOAT 2024 - extractos relevantes):
{contexto}

Reglas:
- Devuelve SOLO códigos SOAT que aparezcan explícitamente en el contexto con su descripción exacta.
- Incluye la página entre paréntesis de donde lo tomaste.
- Si no hay coincidencia exacta, devuelve vacío.

Formato por línea: <codigo_soat> - <descripcion> (p.<pagina>)
"""
        resp = self.groq.chat.completions.create(
            model=self.groq_model,
            temperature=0.0,
            messages=[
                {"role": "system", "content": "Eres estricto: no inventes códigos. Usa solo lo presente en el contexto."},
                {"role": "user", "content": prompt},
            ],
        )
        return resp.choices[0].message.content.strip()

    def asignar_codigos(self, procedimientos: List[str]) -> List[Dict[str, str]]:
        resultados = []
        for p in procedimientos:
            docs = self.buscar(p, k=12)
            propuesta = self.proponer_codigos_soat(p, docs)
            if propuesta:
                for linea in propuesta.splitlines():
                    linea = linea.strip()
                    if not linea:
                        continue
                    resultados.append({"procedimiento": p, "codigo_soat": linea})
        return resultados


def main():
    parser = argparse.ArgumentParser(description="Proponer códigos SOAT vía RAG desde un índice FAISS")
    parser.add_argument("--index", dest="index_dir", default="soat_faiss", help="Directorio del índice FAISS (por defecto: soat_faiss)")
    parser.add_argument("--groq_key", dest="groq_key", default=GROQ_API_KEY, help="API key de Groq (por defecto, toma GROQ_API_KEY del entorno)")
    parser.add_argument("--groq_model", dest="groq_model", default=os.getenv("GROQ_MODEL_SOAT") or config.GROQ_MODEL_SOAT, help="Modelo Groq (por defecto: GROQ_MODEL_SOAT)")
    parser.add_argument("--proc", dest="procedimientos", action="append", help="Descripción de procedimiento (repetible)")
    parser.add_argument("--proc_file", dest="proc_file", help="Archivo con procedimientos, uno por línea")
    args = parser.parse_args()

    if not args.groq_key:
        raise SystemExit("Falta GROQ_API_KEY (usa --groq_key o exporta variable de entorno)")

    procedimientos: List[str] = []
    if args.procedimientos:
        procedimientos.extend([p for p in args.procedimientos if p and p.strip()])
    if args.proc_file and os.path.exists(args.proc_file):
        with open(args.proc_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    procedimientos.append(line)
    if not procedimientos:
        procedimientos = [
           "REDUCCIÓN ABIERTA DE FRACTURA CON FIJACIÓN INTERNA  (DISPOSITIVOS DE FIJACIÓN U OSTEOSÍNTESIS) DE CLAVÍCULA",
           # "Desbridamiento de tejido blando de pie izquierdo",
            #"Reducción cerrada de fractura de falange proximal del pie"
        ]

    retriever = SOATRetriever(faiss_index_path=args.index_dir, groq_api_key=args.groq_key, groq_model=args.groq_model)
    resultados = retriever.asignar_codigos(procedimientos)
    print(json.dumps({"procedimientos": procedimientos, "resultados": resultados}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()