"""
=============================================================
SURYA OCR API - Reconocimiento de Manuscritos
=============================================================
API REST para reconocer texto manuscrito usando Surya OCR
Listo para desplegar en Hugging Face Spaces
=============================================================
"""

import io
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image
from typing import Optional, List

# Surya OCR
from surya.detection import DetectionPredictor
from surya.recognition import RecognitionPredictor, FoundationPredictor

# =============================================================
# CONFIGURACION DE LA API
# =============================================================

app = FastAPI(
    title="Surya OCR API",
    description="Reconocimiento de texto manuscrito con Surya OCR - Modelo de alta precision",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# Permitir CORS (acceso desde cualquier origen)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =============================================================
# MODELOS DE RESPUESTA
# =============================================================

class LineDetail(BaseModel):
    """Detalle de cada linea reconocida"""
    texto: str
    confianza: float

class OCRResponse(BaseModel):
    """Respuesta del reconocimiento OCR"""
    exito: bool
    texto: str
    confianza: float
    lineas: List[LineDetail]
    mensaje: Optional[str] = None

# =============================================================
# GESTOR DE MODELOS SURYA (Carga diferida)
# =============================================================

class SuryaOCR:
    """
    Gestor de modelos Surya con carga diferida.
    Los modelos se cargan solo cuando se necesitan.
    """

    def __init__(self):
        self._detector = None
        self._recognizer = None
        self._foundation = None
        self._loaded = False

    def _load_models(self):
        """Carga los modelos de Surya"""
        if not self._loaded:
            print("=" * 50)
            print("Cargando modelos Surya OCR...")
            print("=" * 50)

            print("[1/3] Cargando detector de texto...")
            self._detector = DetectionPredictor()

            print("[2/3] Cargando modelo foundation...")
            self._foundation = FoundationPredictor()

            print("[3/3] Cargando reconocedor de texto...")
            self._recognizer = RecognitionPredictor(self._foundation)

            self._loaded = True
            print("=" * 50)
            print("Modelos cargados correctamente!")
            print("=" * 50)

    @property
    def detector(self):
        self._load_models()
        return self._detector

    @property
    def recognizer(self):
        self._load_models()
        return self._recognizer

    def recognize(self, image: Image.Image) -> tuple:
        """
        Reconoce texto en una imagen.

        Args:
            image: Imagen PIL

        Returns:
            tuple: (texto_completo, confianza_promedio, lista_lineas)
        """
        # Asegurar RGB
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Redimensionar si es muy grande (max 1024px)
        max_size = 1024
        if max(image.size) > max_size:
            ratio = max_size / max(image.size)
            new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
            image = image.resize(new_size, Image.Resampling.LANCZOS)

        # Procesar con Surya
        images = [image]

        predictions = self.recognizer(
            images,
            det_predictor=self.detector
        )

        # Extraer resultados
        if not predictions or not predictions[0].text_lines:
            return "", 0.0, []

        textos = []
        confianzas = []
        lineas = []

        for line in predictions[0].text_lines:
            texto = line.text.strip()
            confianza = line.confidence

            # FILTROS para eliminar ruido:
            # 1. Solo texto con confianza >= 50%
            if confianza < 0.50:
                continue

            # 2. Ignorar texto vacío
            if not texto:
                continue

            # 3. Ignorar patrones repetitivos (ruido común)
            if "the state" in texto.lower():
                continue
            if "and the" in texto.lower() and len(texto) > 50:
                continue

            # 4. Ignorar líneas muy cortas (1-2 caracteres) con baja confianza
            if len(texto) <= 2 and confianza < 0.80:
                continue

            textos.append(texto)
            confianzas.append(confianza)
            lineas.append({
                "texto": texto,
                "confianza": round(confianza, 3)
            })

        texto_completo = ' '.join(textos)
        confianza_promedio = sum(confianzas) / len(confianzas) if confianzas else 0.0

        return texto_completo, confianza_promedio, lineas

# Instancia global
surya = SuryaOCR()

# =============================================================
# ENDPOINTS DE LA API
# =============================================================

@app.get("/")
async def inicio():
    """
    Pagina de inicio con informacion de la API
    """
    return {
        "nombre": "Surya OCR API",
        "version": "1.0.0",
        "descripcion": "API para reconocimiento de texto manuscrito",
        "modelo": "Surya OCR (2024)",
        "endpoints": {
            "GET /": "Esta informacion",
            "GET /health": "Estado del servicio",
            "GET /docs": "Documentacion interactiva (Swagger)",
            "POST /recognize": "Reconocer texto en imagen"
        },
        "uso": "Sube una imagen a /recognize para obtener el texto"
    }


@app.get("/health")
async def health():
    """
    Verifica el estado del servicio
    """
    return {
        "status": "ok",
        "modelo": "Surya OCR",
        "listo": True
    }


@app.post("/recognize", response_model=OCRResponse)
async def recognize(file: UploadFile = File(...)):
    """
    Reconoce texto manuscrito en una imagen.

    Sube una imagen (JPG, PNG, etc.) y obtendras el texto reconocido.

    - **file**: Archivo de imagen a procesar

    Retorna:
    - **exito**: Si el reconocimiento fue exitoso
    - **texto**: Texto completo reconocido
    - **confianza**: Nivel de confianza (0-1)
    - **lineas**: Detalle de cada linea detectada
    """
    # Validar que sea imagen
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(
            status_code=400,
            detail="El archivo debe ser una imagen (JPG, PNG, etc.)"
        )

    try:
        # Leer imagen
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))

        # Reconocer texto
        texto, confianza, lineas = surya.recognize(image)

        # Respuesta
        if texto:
            return OCRResponse(
                exito=True,
                texto=texto,
                confianza=round(confianza, 3),
                lineas=[LineDetail(**l) for l in lineas],
                mensaje=f"Se reconocieron {len(lineas)} lineas de texto"
            )
        else:
            return OCRResponse(
                exito=False,
                texto="",
                confianza=0.0,
                lineas=[],
                mensaje="No se detecto texto en la imagen"
            )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error procesando imagen: {str(e)}"
        )


# =============================================================
# PUNTO DE ENTRADA
# =============================================================

if __name__ == "__main__":
    import uvicorn
    print("\n" + "=" * 50)
    print("  SURYA OCR API")
    print("  http://localhost:7860")
    print("  Docs: http://localhost:7860/docs")
    print("=" * 50 + "\n")
    uvicorn.run(app, host="0.0.0.0", port=7860)