# easyocr_tuning_rest.py
# FastAPI REST service for EasyOCR hyperparameter evaluation
# Usage: uvicorn easyocr_tuning_rest:app --host 0.0.0.0 --port 8000

import os
import re
import time
import threading
from typing import Optional, List
from contextlib import asynccontextmanager

import numpy as np
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

import easyocr
from jiwer import wer, cer
from dataset_manager import ImageTextDataset


def get_gpu_info() -> dict:
    """Get GPU status information from PyTorch."""
    info = {
        "cuda_available": torch.cuda.is_available(),
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "gpu_count": 0,
        "gpu_name": None,
        "gpu_memory_total": None,
        "gpu_memory_used": None,
    }

    if info["cuda_available"]:
        try:
            info["gpu_count"] = torch.cuda.device_count()
            if info["gpu_count"] > 0:
                info["gpu_name"] = torch.cuda.get_device_name(0)
                info["gpu_memory_total"] = f"{torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB"
                info["gpu_memory_used"] = f"{torch.cuda.memory_allocated(0) / (1024**3):.2f} GB"
        except Exception as e:
            info["gpu_error"] = str(e)

    return info


# Model configuration via environment variables
DEFAULT_LANGUAGES = os.environ.get("EASYOCR_LANGUAGES", "es,en").split(",")


# Global state for model and dataset
class AppState:
    reader: Optional[easyocr.Reader] = None
    dataset: Optional[ImageTextDataset] = None
    dataset_path: Optional[str] = None
    languages: List[str] = DEFAULT_LANGUAGES
    lock: threading.Lock = None  # Protects OCR model from concurrent access

    def __init__(self):
        self.lock = threading.Lock()


state = AppState()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load EasyOCR model at startup."""
    gpu_info = get_gpu_info()
    print("=" * 50)
    print("GPU STATUS")
    print("=" * 50)
    print(f"  CUDA available: {gpu_info['cuda_available']}")
    print(f"  Device: {gpu_info['device']}")
    if gpu_info['cuda_available']:
        print(f"  GPU count: {gpu_info['gpu_count']}")
        print(f"  GPU name: {gpu_info['gpu_name']}")
        print(f"  GPU memory total: {gpu_info['gpu_memory_total']}")
    print("=" * 50)

    print(f"Loading EasyOCR models...")
    print(f"  Languages: {state.languages}")
    state.reader = easyocr.Reader(
        state.languages,
        gpu=gpu_info['cuda_available'],
    )

    if gpu_info['cuda_available']:
        gpu_after = get_gpu_info()
        print(f"  GPU memory after load: {gpu_after.get('gpu_memory_used', 'N/A')}")

    print("Model loaded successfully!")
    yield
    state.reader = None
    state.dataset = None


app = FastAPI(
    title="EasyOCR Tuning API",
    description="REST API for EasyOCR hyperparameter evaluation",
    version="1.0.0",
    lifespan=lifespan,
)


class EvaluateRequest(BaseModel):
    """Request schema with all tunable EasyOCR hyperparameters."""
    pdf_folder: str = Field("/app/dataset", description="Path to dataset folder")

    # Detection thresholds (CRAFT algorithm)
    text_threshold: float = Field(0.7, ge=0.0, le=1.0, description="Text confidence threshold")
    low_text: float = Field(0.4, ge=0.0, le=1.0, description="Text lower-bound score")
    link_threshold: float = Field(0.4, ge=0.0, le=1.0, description="Link confidence threshold")

    # Bounding box merging
    slope_ths: float = Field(0.1, ge=0.0, le=1.0, description="Maximum slope for box merging")
    ycenter_ths: float = Field(0.5, ge=0.0, le=2.0, description="Maximum vertical shift for merging")
    height_ths: float = Field(0.5, ge=0.0, le=2.0, description="Maximum height variance for merging")
    width_ths: float = Field(0.5, ge=0.0, le=2.0, description="Maximum horizontal distance for merging")
    add_margin: float = Field(0.1, ge=0.0, le=1.0, description="Bounding box extension margin")

    # Contrast handling
    contrast_ths: float = Field(0.1, ge=0.0, le=1.0, description="Contrast threshold for dual-pass")
    adjust_contrast: float = Field(0.5, ge=0.0, le=1.0, description="Target contrast adjustment level")

    # Decoder options
    decoder: str = Field("greedy", description="Decoder type: greedy, beamsearch, wordbeamsearch")
    beamWidth: int = Field(5, ge=1, le=20, description="Beam width for beam search decoders")

    # Other
    min_size: int = Field(10, ge=1, description="Minimum text box size in pixels")
    rotation_info: Optional[List[int]] = Field(None, description="Rotation angles to try: [90, 180, 270]")

    # Page range
    start_page: int = Field(5, ge=0, description="Start page index (inclusive)")
    end_page: int = Field(10, ge=1, description="End page index (exclusive)")
    save_output: bool = Field(False, description="Save OCR predictions to debugset folder")


class EvaluateResponse(BaseModel):
    """Response schema matching CLI output."""
    CER: float
    WER: float
    TIME: float
    PAGES: int
    TIME_PER_PAGE: float


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    dataset_loaded: bool
    dataset_size: Optional[int] = None
    languages: Optional[List[str]] = None
    cuda_available: Optional[bool] = None
    device: Optional[str] = None
    gpu_name: Optional[str] = None
    gpu_memory_used: Optional[str] = None
    gpu_memory_total: Optional[str] = None


def assemble_easyocr_result(result: list) -> str:
    """
    Assemble EasyOCR result into text.
    EasyOCR returns: [(bbox, text, confidence), ...]
    """
    if not result:
        return ""

    # Sort by vertical position (y), then horizontal (x)
    # bbox format: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    def get_y_center(item):
        bbox = item[0]
        return (bbox[0][1] + bbox[2][1]) / 2

    def get_x(item):
        return item[0][0][0]

    # Group by lines based on y-center
    sorted_items = sorted(result, key=lambda x: (get_y_center(x), get_x(x)))

    if not sorted_items:
        return ""

    # Adaptive line tolerance
    heights = []
    for item in sorted_items:
        bbox = item[0]
        h = abs(bbox[2][1] - bbox[0][1])
        heights.append(h)

    median_h = float(np.median(heights)) if heights else 20.0
    line_tol = max(8.0, 0.6 * median_h)

    lines, cur_line, last_y = [], [], None
    for item in sorted_items:
        y_center = get_y_center(item)
        text = item[1]

        if last_y is None or abs(y_center - last_y) <= line_tol:
            cur_line.append((get_x(item), text))
        else:
            cur_line.sort(key=lambda t: t[0])
            lines.append(" ".join(t[1] for t in cur_line))
            cur_line = [(get_x(item), text)]
        last_y = y_center

    if cur_line:
        cur_line.sort(key=lambda t: t[0])
        lines.append(" ".join(t[1] for t in cur_line))

    text = " ".join(lines)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def evaluate_text(reference: str, prediction: str) -> dict:
    """Calculate WER and CER metrics."""
    return {"WER": wer(reference, prediction), "CER": cer(reference, prediction)}


@app.get("/health", response_model=HealthResponse)
def health_check():
    """Check if the service is ready."""
    gpu_info = get_gpu_info()
    return HealthResponse(
        status="ok" if state.reader is not None else "initializing",
        model_loaded=state.reader is not None,
        dataset_loaded=state.dataset is not None,
        dataset_size=len(state.dataset) if state.dataset else None,
        languages=state.languages,
        cuda_available=gpu_info.get("cuda_available"),
        device=gpu_info.get("device"),
        gpu_name=gpu_info.get("gpu_name"),
        gpu_memory_used=gpu_info.get("gpu_memory_used"),
        gpu_memory_total=gpu_info.get("gpu_memory_total"),
    )


@app.post("/evaluate", response_model=EvaluateResponse)
def evaluate(request: EvaluateRequest):
    """
    Evaluate OCR with given hyperparameters.
    Returns CER, WER, and timing metrics.
    """
    if state.reader is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    # Validate decoder
    if request.decoder not in ["greedy", "beamsearch", "wordbeamsearch"]:
        raise HTTPException(status_code=400, detail=f"Invalid decoder: {request.decoder}")

    # Load or reload dataset if path changed
    if state.dataset is None or state.dataset_path != request.pdf_folder:
        if not os.path.isdir(request.pdf_folder):
            raise HTTPException(status_code=400, detail=f"Dataset folder not found: {request.pdf_folder}")
        state.dataset = ImageTextDataset(request.pdf_folder)
        state.dataset_path = request.pdf_folder

    if len(state.dataset) == 0:
        raise HTTPException(status_code=400, detail="Dataset is empty")

    # Validate page range
    start = request.start_page
    end = min(request.end_page, len(state.dataset))
    if start >= end:
        raise HTTPException(status_code=400, detail=f"Invalid page range: {start}-{end}")

    cer_list, wer_list = [], []
    time_per_page_list = []
    t0 = time.time()

    # Lock to prevent concurrent OCR access (model is not thread-safe)
    with state.lock:
        for idx in range(start, end):
            img, ref = state.dataset[idx]
            arr = np.array(img)

            tp0 = time.time()
            result = state.reader.readtext(
                arr,
                # Detection thresholds
                text_threshold=request.text_threshold,
                low_text=request.low_text,
                link_threshold=request.link_threshold,
                # Bounding box merging
                slope_ths=request.slope_ths,
                ycenter_ths=request.ycenter_ths,
                height_ths=request.height_ths,
                width_ths=request.width_ths,
                add_margin=request.add_margin,
                # Contrast
                contrast_ths=request.contrast_ths,
                adjust_contrast=request.adjust_contrast,
                # Decoder
                decoder=request.decoder,
                beamWidth=request.beamWidth,
                # Other
                min_size=request.min_size,
                rotation_info=request.rotation_info,
            )

            pred = assemble_easyocr_result(result)
            time_per_page_list.append(float(time.time() - tp0))

            # Save prediction to debugset if requested
            if request.save_output:
                out_path = state.dataset.get_output_path(idx, "easyocr_text")
                with open(out_path, "w", encoding="utf-8") as f:
                    f.write(pred)

            m = evaluate_text(ref, pred)
            cer_list.append(m["CER"])
            wer_list.append(m["WER"])

    return EvaluateResponse(
        CER=float(np.mean(cer_list)) if cer_list else 1.0,
        WER=float(np.mean(wer_list)) if wer_list else 1.0,
        TIME=float(time.time() - t0),
        PAGES=len(cer_list),
        TIME_PER_PAGE=float(np.mean(time_per_page_list)) if time_per_page_list else 0.0,
    )


@app.post("/evaluate_full", response_model=EvaluateResponse)
def evaluate_full(request: EvaluateRequest):
    """Evaluate on ALL pages (ignores start_page/end_page)."""
    request.start_page = 0
    request.end_page = 9999
    return evaluate(request)


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