import asyncio
import time

import uvicorn
from fastapi import FastAPI, Security, HTTPException
from fastapi.security import APIKeyHeader
from pydantic import BaseModel

from utils.langgraph_agent import RULES
from utils.secrets_manager import SECRETS
from services.document_extractor import parse_s3_uri, extract_text_from_s3_document
from services.authorization import evaluate_servico
from services.result_store import save_results

app = FastAPI()

_api_key_header = APIKeyHeader(name="X-API-Key")
API_KEY = SECRETS["API-KEY"]


def verify_api_key(api_key: str = Security(_api_key_header)):
    if api_key != API_KEY:
        raise HTTPException(status_code=403, detail="Invalid API key")
    return api_key


async def process_guia(guia: dict) -> dict:
    t_start = time.time()
    anexos = guia.get("anexos", [])

    async def _extract_anexo(anexo_idx: int, anexo: dict):
        s3_uri = anexo.get("urlAnexo") or anexo.get("URLAnexo", "")
        nome_arquivo = anexo.get("nomeArquivo", f"attachment_{anexo_idx}")
        if not s3_uri or not s3_uri.startswith("s3://"):
            anexo["textoExtraido"] = ""
            return None
        t0 = time.time()
        try:
            bucket, key = parse_s3_uri(s3_uri)
            extracted_text, page_count = await extract_text_from_s3_document(bucket, key)
        except Exception as e:
            extracted_text = ""
            page_count = 0
            anexo["error"] = str(e)
        anexo["textoExtraido"] = extracted_text
        anexo["pageCount"] = page_count
        anexo["tempoExtracaoSegundos"] = round(time.time() - t0, 2)
        return f"--- {nome_arquivo} ---\n{extracted_text}"

    t_extracao_start = time.time()
    parts = await asyncio.gather(*[_extract_anexo(i, a) for i, a in enumerate(anexos)])
    file_content = "\n\n".join(p for p in parts if p)
    t_extracao = round(time.time() - t_extracao_start, 2)

    servicos = guia.get("servicos", [])
    t_agent_start = time.time()
    guia["avaliacaoAgente"] = list(
        await asyncio.gather(*[evaluate_servico(s, guia, file_content) for s in servicos])
    )
    t_agent = round(time.time() - t_agent_start, 2)

    guia["tempoProcessamento"] = {
        "extracaoSegundos": t_extracao,
        "agentSegundos": t_agent,
        "totalSegundos": round(time.time() - t_start, 2),
    }
    return guia


class ProcessRequest(BaseModel):
    operadora: dict
    guias: list[dict]


@app.post("/process", dependencies=[Security(verify_api_key)])
async def process(request: ProcessRequest):
    raw_results = await asyncio.gather(
        *[process_guia(guia) for guia in request.guias],
        return_exceptions=True,
    )
    results = [
        {"error": str(r), "guia": request.guias[i].get("guia", {}).get("codigoGuiaLocal", f"index_{i}")}
        if isinstance(r, Exception) else r
        for i, r in enumerate(raw_results)
    ]

    await save_results(results)

    return {"status": "success", "operadora": request.operadora, "guias": results}


@app.get("/health")
async def health():
    return {"status": "healthy"}


@app.get("/rules")
async def get_rules():
    return {"codes": list(RULES.keys())}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)