from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import requests
import pandas as pd
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

INDEX_PATH = "/data/indice_documentos.faiss"
CHUNKS_PATH = "/data/dataset_chunks.json"

LLAMA_URL = "http://192.168.1.39:8089/v1/chat/completions"
MODEL_NAME = "Qwen3-1.7B-Instruct-Q4_K_M.gguf"

modelo_embeddings = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
index = faiss.read_index(INDEX_PATH)
dataset_chunks = pd.read_json(CHUNKS_PATH)

class PerguntaRequest(BaseModel):
    pergunta: str

@app.get("/")
def root():
    return {"status": "ok"}

def pesquisar_chunks(query, modelo_embeddings, index, dataset_chunks, top_k=5):
    query_embedding = modelo_embeddings.encode([query], convert_to_numpy=True)
    query_embedding = np.array(query_embedding, dtype="float32")

    distancias, indices = index.search(query_embedding, top_k)

    resultados = dataset_chunks.iloc[indices[0]].copy()
    resultados["score_distancia"] = distancias[0]
    return resultados

def construir_contexto_rag(query, resultados, max_chars_por_chunk=500):
    partes = []

    for _, row in resultados.iterrows():
        chunk_curto = row["texto_chunk"][:max_chars_por_chunk]
        partes.append(
            f"[Documento: {row['nome_pdf']} | Categoria: {row['categoria']}]\n{chunk_curto}"
        )

    contexto = "\n\n---\n\n".join(partes)

    return f"""
És um assistente especializado em documentos do Exército Português.

Responde apenas com base no contexto fornecido.

Pergunta:
{query}

Contexto:
{contexto}

Regras:
- Responde em português.
- Sê objetivo.
- Dá uma resposta curta.
- Se a resposta não estiver no contexto, diz isso explicitamente.
- Se possível, menciona o nome do documento.
"""

def perguntar_qwen(prompt):
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.2,
        "max_tokens": 256,
        "stream": False
    }

    try:
        response = requests.post(
            LLAMA_URL,
            json=payload,
            timeout=18000
        )

        print("STATUS LLAMA:", response.status_code)
        print("BODY LLAMA:", response.text[:1000])

        response.raise_for_status()
        data = response.json()

        return data["choices"][0]["message"]["content"]

    except Exception as e:
        print("ERRO AO CHAMAR LLAMA:", repr(e))
        raise HTTPException(status_code=500, detail=f"Erro llama.cpp: {str(e)}")

@app.post("/perguntar")
def perguntar(request: PerguntaRequest):
    try:
        query = request.pergunta
        print("PERGUNTA:", query)

        resultados = pesquisar_chunks(
            query=query,
            modelo_embeddings=modelo_embeddings,
            index=index,
            dataset_chunks=dataset_chunks,
            top_k=1
        )

        print("RESULTADOS ENCONTRADOS:", len(resultados))

        prompt = construir_contexto_rag(query, resultados)
        resposta = perguntar_qwen(prompt)

        fontes = []
        for _, row in resultados.iterrows():
            fontes.append({
                "nome_pdf": row["nome_pdf"],
                "categoria": row["categoria"],
                "chunk_id": int(row["chunk_id"]),
                "score_distancia": float(row["score_distancia"])
            })

        return {
            "pergunta": query,
            "resposta": resposta,
            "fontes": fontes
        }

    except HTTPException:
        raise
    except Exception as e:
        print("ERRO NO ENDPOINT:", repr(e))
        raise HTTPException(status_code=500, detail=str(e))