import os
from pathlib import Path
from typing import Literal
from pinecone import Pinecone, ServerlessSpec
from langchain_voyageai import VoyageAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.retrievers import BM25Retriever
import pickle

# --- Configuración LexAI -------------------------------------------
PINECONE_INDEX = "lexai-garrigues-v1"
AREA_LEGAL = Literal["mercantil", "laboral", "penal", "fiscal"]

embeddings = VoyageAIEmbeddings(
    model="voyage-3-large",          # Recomendado por Anthropic para Claude
    voyage_api_key=os.environ["VOYAGE_API_KEY"]
)

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", ". ", " "]
)

def detect_area_legal(filepath: str) -> AREA_LEGAL:
    """Clasifica el documento por área legal según el path."""
    fp = filepath.lower()
    if "mercantil" in fp: return "mercantil"
    if "laboral"   in fp: return "laboral"
    if "penal"     in fp: return "penal"
    return "fiscal"

def load_and_chunk(doc_path: Path):
    """Carga un documento legal y lo divide en chunks con metadatos."""
    ext = doc_path.suffix.lower()
    loader = PyPDFLoader(str(doc_path)) if ext == ".pdf" else Docx2txtLoader(str(doc_path))
    docs = loader.load()
    chunks = splitter.split_documents(docs)
    # Enriquecer metadatos para filtrado posterior
    for chunk in chunks:
        chunk.metadata.update({
            "source":    doc_path.name,
            "area":      detect_area_legal(str(doc_path)),
            "expediente": doc_path.stem,
        })
    return chunks

def ingest_corpus(corpus_dir: "str" = "./corpus"):
    """Indexa todo el corpus legal en Pinecone + BM25."""
    pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])

    # Crear índice si no existe (dimensión 1024 = voyage-3-large)
    if PINECONE_INDEX not in pc.list_indexes().names():
        pc.create_index(
            name=PINECONE_INDEX, dimension=1024, metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="eu-west-1")
        )

    all_chunks = []
    for ext in ("*.pdf", "*.docx", "*.txt"):
        for fp in Path(corpus_dir).rglob(ext):
            all_chunks.extend(load_and_chunk(fp))

    print(f"[LexAI] Indexando {len(all_chunks)} chunks...")

    # Upsert en Pinecone en batches de 100
    vectorstore = PineconeVectorStore.from_documents(
        all_chunks, embeddings, index_name=PINECONE_INDEX
    )

    # Persistir BM25 para búsqueda keyword
    bm25 = BM25Retriever.from_documents(all_chunks)
    with open("bm25_index.pkl", "wb") as f:
        pickle.dump(bm25, f)

    print("[LexAI] ✓ Ingesta completada")
    return vectorstore

from langgraph.graph import StateGraph, START, END
from langchain_anthropic import ChatAnthropic
from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.documents import Document
from sentence_transformers import CrossEncoder
from pydantic import BaseModel, Field
from typing import TypedDict, Optional
import pickle

# ---- Modelos de respuesta estructurada ----------------------------

class LegalSource(BaseModel):
    expediente: str  = Field(description="ID del expediente legal")
    fragmento:  str  = Field(description="Extracto literal relevante")
    area:       str  = Field(description="Área legal: mercantil/laboral/penal/fiscal")
    relevancia: float= Field(description="Score de relevancia 0-1")

class LegalAnswer(BaseModel):
    respuesta:  str           = Field(description="Respuesta jurídica fundamentada")
    confianza:  float         = Field(description="Confianza en la respuesta 0-1")
    fuentes:    list[LegalSource]= Field(description="Fuentes citadas con expediente")
    requiere_abogado: bool   = Field(description="True si la consulta requiere revisión senior")

# ---- Estado del grafo -------------------------------------------

class LexState(TypedDict):
    pregunta: str
    area_filtro: Optional[str]        # "mercantil" | "laboral" | "penal" | "fiscal"
    context: list[Document]
    respuesta: LegalAnswer

# ---- Componentes ------------------------------------------------

llm = ChatAnthropic(model="claude-sonnet-4-6")
structured_llm = llm.with_structured_output(LegalAnswer)

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

with open("bm25_index.pkl", "rb") as f:
    bm25_retriever = pickle.load(f)
bm25_retriever.k = 10

# Retrieval híbrido: 30% BM25 + 70% semántico (RRF)
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, dense_retriever],
    weights=[0.3, 0.7]
)

# ---- Prompt legal con citas -------------------------------------

LEGAL_PROMPT = ChatPromptTemplate.from_template("""
Eres un asistente jurídico de LexAI para el despacho Garrigues & Asociados.
SOLO puedes responder basándote en el contexto proporcionado.
Si no encuentras información suficiente, di exactamente: "No dispongo de información suficiente en el corpus para responder esta consulta."

CONTEXTO (expedientes del despacho):
{context}

CONSULTA JURÍDICA:
{pregunta}

INSTRUCCIONES:
1. Cita los expedientes relevantes con su ID exacto
2. Extrae fragmentos literales que soporten la respuesta
3. Indica el área legal y nivel de confianza
4. Si la consulta requiere criterio senior (estrategia procesal, riesgo > 100k€), marca requiere_abogado=true
""")

# ---- Nodos del grafo -------------------------------------------

async def retrieve_hybrid(state: LexState) -> LexState:
    """Recupera con búsqueda híbrida + filtrado por área legal."""
    docs = await ensemble_retriever.ainvoke(state["pregunta"])

    # Filtrar por área si se especificó
    if state.get("area_filtro"):
        docs = [d for d in docs if d.metadata.get("area") == state["area_filtro"]]

    return {"context": docs}

async def rerank_docs(state: LexState) -> LexState:
    """Aplica Cross-Encoder para reordenar por relevancia real."""
    query = state["pregunta"]
    docs  = state["context"]
    pairs = [[query, d.page_content] for d in docs]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
    return {"context": [d for d, _ in ranked[:5]]}  # top-5 reranked

async def generate_legal_answer(state: LexState) -> LexState:
    """Genera respuesta jurídica estructurada con Claude."""
    ctx = "\n\n---\n\n".join(
        f"[{d.metadata['expediente']}] ({d.metadata['area']})\n{d.page_content}"
        for d in state["context"]
    )
    messages = LEGAL_PROMPT.format_messages(context=ctx, pregunta=state["pregunta"])
    answer: LegalAnswer = await structured_llm.ainvoke(messages)
    return {"respuesta": answer}

# ---- Construir grafo -------------------------------------------

builder = StateGraph(LexState)
builder.add_node("retrieve", retrieve_hybrid)
builder.add_node("rerank",   rerank_docs)
builder.add_node("generate", generate_legal_answer)
builder.add_edge(START,      "retrieve")
builder.add_edge("retrieve", "rerank")
builder.add_edge("rerank",   "generate")
builder.add_edge("generate", END)

lex_rag = builder.compile()

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from lexai.rag_pipeline import lex_rag

app = FastAPI(title="LexAI API", version="1.0.0")

class QueryRequest(BaseModel):
    pregunta:    str
    area_filtro: Optional[str] = None  # "mercantil" | "laboral" | "penal" | "fiscal"
    max_fuentes: int = 5

@app.post("/api/query")
async def query_legal(req: QueryRequest):
    """
    Consulta legal con RAG híbrido.
    Retorna respuesta fundamentada con citas de expedientes.
    """
    try:
        result = await lex_rag.ainvoke({
            "pregunta":    req.pregunta,
            "area_filtro": req.area_filtro,
        })
        answer = result["respuesta"]
        return {
            "respuesta":        answer.respuesta,
            "confianza":        answer.confianza,
            "requiere_abogado": answer.requiere_abogado,
            "fuentes": [
                {
                    "expediente": f.expediente,
                    "area":       f.area,
                    "fragmento":  f.fragmento[:300],
                    "relevancia": round(f.relevancia, 3)
                }
                for f in answer.fuentes[:req.max_fuentes]
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health(): return {"status": "ok", "model": "claude-sonnet-4-6"}

# Ejemplo de respuesta de la API:
## POST /api/query
## {
##   "pregunta": "¿Precedentes de despido improcedente en sector tecnológico 2024?",
##   "area_filtro": "laboral"
## }
##
## → respuesta: "Según el expediente LAB-2024-0847, el TS en sentencia de...
##              [1] LAB-2024-0847 (relevancia: 0.94)
##              [2] LAB-2023-1203 (relevancia: 0.88)"