main.py

from unittest import result

from fastapi import FastAPI, File, UploadFile, HTTPException
from google import genai
from httpcore import request
from google.genai import types
from dotenv import load_dotenv
from pydantic import BaseModel
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import cv2
from paddleocr import PaddleOCR
import numpy as np
import time
import os

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["FLAGS_allocator_strategy"] = "naive_best_fit"

MODEL_ID = "models/gemini-2.5-flash"
MODEL_FALLBACK = "models/gemini-3.1-flash-lite-preview"
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
modelo_local = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

load_dotenv()

api_key = os.environ.get("GEMINI_API_KEY")

client = genai.Client(api_key=api_key)


config = types.GenerateContentConfig(
    temperature=0,
)

class ChatRequest(BaseModel):
    pregunta: str
    contexto: str = ""
    historial: list[str] = []

class ChatResponse(BaseModel):
    respuesta: str
    historial: list[str] = []

class OCRResponse(BaseModel):
	texto: str

app = FastAPI(title="OCR API", version="1.0.0")
ocr = PaddleOCR(use_angle_cls=False, lang="es", det_limit_side_len=480)

try:
    local_model = pipeline("text2text-generation", model="google/flan-t5-small")
except Exception as e:
    print(f"Aviso: No se pudo cargar el modelo local: {e}")
    local_model = None

app.mount("/static", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def home():
   return FileResponse("static/index.html")



@app.post("/api/chat", response_model=ChatResponse)
def prompt(request: ChatRequest):
    max_reintentos = 2
    segundos_espera = 1.5
    e = "Error desconocido"

    try:
        pregunta = request.pregunta
        historial = request.historial
        texto_ocr = request.contexto  # <--- Tomamos el texto enviado por el frontend

        texto_prompt = f"""
                Eres un asistente que responde preguntas usando SOLO información del documento.

                REGLAS IMPORTANTES:
                - No copies el documento completo.
                - No repitas texto largo del documento.
                - Extrae SOLO la información necesaria.
                - Si la respuesta no está en el documento, di: "No aparece en el documento".
                - Responde de forma breve y directa.

                DOCUMENTO:
                \"\"\"{texto_ocr}\"\"\"  <--- Cambiado: ahora usa el OCR

                HISTORIAL:
                {historial}

                PREGUNTA:
                {pregunta}

                RESPUESTA:
                """

        for intento in range(max_reintentos):
            try:
                response = client.models.generate_content(
                        model=MODEL_ID,
                        contents=texto_prompt,
                        config=config
                )
                return ChatResponse(respuesta=response.text)
            except Exception as ex:
                e = ex
                print(f"Intento {intento+1} fallido: {e}")
                time.sleep(segundos_espera)
                
        try:
            response = client.models.generate_content(
                model=MODEL_FALLBACK,
                contents=texto_prompt,
                config=config
            )
            return ChatResponse(respuesta=response.text)
        except Exception as e_fallback:
            print(f"Fallido: {e_fallback}")
            try:
            # Simplificamos el prompt para el modelo local pequeño
                res_local = modelo_local(f"question: {pregunta} context: {texto_ocr}", max_new_tokens=50)
                return ChatResponse(respuesta=res_local[0]['generated_text'])
            except Exception as e_final:
                raise HTTPException(status_code=500, detail="Error en todos los modelos (incluyendo local)")

    except Exception as e:
        if isinstance(e, HTTPException): raise e
        raise HTTPException(status_code=500, detail=str(e))

    except Exception as e:
        if isinstance(e, HTTPException): raise e
        raise HTTPException(status_code=500, detail=str(e))
    

def extraer_texto(resultado_ocr):
    textos = []
    
    # PaddleOCR devuelve una lista de páginas. 
    # Validamos que el resultado no sea None y que la primera página tenga contenido.
    if resultado_ocr and resultado_ocr[0] is not None:
        for linea in resultado_ocr[0]:
            # linea[1][0] es donde reside el texto detectado
            textos.append(linea[1][0])

    return " ".join(textos)

      
def preprocesar_imagen(image_bytes):
    # Convertir bytes a imagen
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if img is None:
        raise ValueError("No se pudo decodificar la imagen")

    # PaddleOCR maneja internamente el binarizado, 
    # es mejor enviarle la imagen limpia o solo en gris.
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return gray
    
@app.post("/api/ocr", response_model=OCRResponse)
async def ocr_image(file: UploadFile = File(...)):
    if file.content_type not in {"image/png", "image/jpeg", "image/jpg", "image/webp"}:
        raise HTTPException(status_code=400, detail="Formato no soportado")

    try:
        # Leemos los bytes del archivo
        contenido = await file.read()
        
        # Preprocesamos y ejecutamos OCR
        img = preprocesar_imagen(contenido)
        
        # cls=True activa la clasificación de ángulo si se configuró en la instancia
        result = ocr.ocr(img) 
        
        texto_extraido = extraer_texto(result)
        
        if not texto_extraido.strip():
            return OCRResponse(texto="No se detectó texto en la imagen.")
            
        return OCRResponse(texto=texto_extraido)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error OCR: {str(e)}")