import io
import time

import cv2
import numpy as np
import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image, UnidentifiedImageError
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

MODEL_NAME = "microsoft/trocr-base-handwritten"

# Cap how many lines we OCR from one image (protects CPU / latency).
MAX_LINES = 30

app = FastAPI(title="Handwriting OCR API", version="0.2.0")

# Prototype CORS: allow any origin so a local frontend can call the deployed API.
# Tighten this to the deployed frontend origin in a later phase.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the model once at import time so requests don't pay the load cost.
processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
model.to(device)
model.eval()


def segment_lines(pil_img, min_line_height=12, pad=8):
    """Split a page image into text-line crops (top-to-bottom).

    TrOCR recognizes a SINGLE text line at a time, so multi-line photos must be
    cut into lines first. We binarize the image and use a horizontal projection
    (ink pixels per row) to find vertical bands that contain writing.

    Returns a list of PIL.Image line crops. May be empty for a blank image.
    """
    gray = cv2.cvtColor(np.array(pil_img.convert("RGB")), cv2.COLOR_RGB2GRAY)
    # Binarize so ink -> white (255). Otsu picks the threshold automatically.
    _, binimg = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    # Open with a small kernel to drop thin ruled lines / speckle noise.
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    binimg = cv2.morphologyEx(binimg, cv2.MORPH_OPEN, kernel)

    proj = binimg.sum(axis=1) / 255.0  # ink-pixel count per row
    if proj.max() <= 0:
        return []
    thresh = max(proj.max() * 0.08, 2)
    rows = proj > thresh

    bands, start = [], None
    for i, on in enumerate(rows):
        if on and start is None:
            start = i
        elif not on and start is not None:
            if i - start >= min_line_height:
                bands.append((start, i))
            start = None
    if start is not None and len(rows) - start >= min_line_height:
        bands.append((start, len(rows)))

    h, w = gray.shape
    crops = []
    for (a, b) in bands:
        a = max(0, a - pad)
        b = min(h, b + pad)
        crops.append(pil_img.crop((0, a, w, b)))
    return crops


def _ocr_images(images):
    """Run TrOCR over a batch of (line) images and return decoded strings."""
    pixel_values = processor(images=images, return_tensors="pt").pixel_values.to(device)
    with torch.inference_mode():
        generated_ids = model.generate(pixel_values, max_new_tokens=64)
    return [t.strip() for t in processor.batch_decode(generated_ids, skip_special_tokens=True)]


def _has_letters(text):
    """Keep only lines that contain at least one letter (filters splatter/noise)."""
    return any(ch.isalpha() for ch in text)


@app.get("/api")
def root():
    return {
        "name": "Handwriting OCR API",
        "status": "running",
        "model": MODEL_NAME,
        "docs": "/docs",
    }


@app.get("/health")
def health():
    return {"status": "ok", "device": device}


@app.post("/ocr")
async def ocr(file: UploadFile = File(...)):
    start_time = time.time()

    image_bytes = await file.read()
    if not image_bytes:
        raise HTTPException(status_code=400, detail="Empty file upload.")

    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except (UnidentifiedImageError, OSError):
        raise HTTPException(
            status_code=400,
            detail="Could not read the uploaded file as an image.",
        )

    # Cut the page into text lines. TrOCR is single-line, so multi-line images
    # must be segmented; if we find 0 or 1 line, just OCR the whole image.
    line_crops = segment_lines(image)
    if len(line_crops) <= 1:
        line_images = [image]
    else:
        line_images = line_crops[:MAX_LINES]

    raw_lines = _ocr_images(line_images)
    # Drop noise lines (e.g. decorative splatter that has no letters), but never
    # end up with nothing — fall back to the raw output if filtering empties it.
    lines = [t for t in raw_lines if _has_letters(t)] or raw_lines

    return {
        "text": "\n".join(lines),
        "lines": lines,
        "line_count": len(lines),
        "model": MODEL_NAME,
        "device": device,
        "latency_seconds": round(time.time() - start_time, 2),
    }


# Must be last so API routes match before the static SPA fallback.
app.mount("/", StaticFiles(directory="static", html=True), name="static")