"""
BilinguaScript HTR API
FastAPI application factory with /health and /transcribe routes.
"""
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from predict import run_inference, normalize_lang, EN_MODEL_ID, FR_MODEL_ID
import preprocess, logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="BilinguaScript HTR API", version="1.1.0")


@app.get("/")
def root():
    """Landing response for the Space page - the real work happens on /api/v1/transcribe."""
    return {
        "service": "BiLinguaScript HTR API",
        "description": "Bilingual English/French handwriting recognition for Cameroonian documents (fine-tuned TrOCR).",
        "models": {"en": EN_MODEL_ID, "fr": FR_MODEL_ID},
        "endpoints": {
            "health": "GET /health",
            "transcribe": "POST /api/v1/transcribe (multipart: file, language=en|fr|auto, segment=true|false)",
        },
        "docs": "/docs",
    }


@app.get("/health")
def health():
    return {
        "status": "ok",
        "models": {"en": EN_MODEL_ID, "fr": FR_MODEL_ID},
    }


@app.post("/api/v1/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    language: str = Form(default="auto"),
    segment: bool = Form(default=True),
):
    """Transcribe an uploaded image.

    - language: 'en'/'english', 'fr'/'french', or 'auto' (defaults to English).
    - segment:  true  -> split the image into lines first (full-page input);
                false -> treat the image as a single pre-cropped line (this is
                what the Django backend sends, since it segments with its own
                pipeline before calling us — re-segmenting a line crop here
                would waste time and risk splitting tall ascenders/descenders).
    """
    data = await file.read()
    if not data:
        raise HTTPException(status_code=400, detail="Empty file.")
    img = preprocess.decode_image(data)
    if img is None:
        raise HTTPException(status_code=422, detail="Could not decode image.")

    lang = normalize_lang(language)

    if segment:
        lines = preprocess.segment_lines(img)
    else:
        lines = [(0, img)]

    results = []
    for idx, line_img in lines:
        encoded = preprocess.encode_png(line_img)
        result = run_inference(encoded, lang)
        results.append({"line_index": idx, **result})
    return JSONResponse({"language": lang, "lines": results})


# ── Kraken blla baseline segmentation (primary engine) ───────────────────────
# kraken's blla neural segmenter is the open-source counterpart of
# Transkribus' ARU-Net baseline detection: it emits reading-order baselines
# with bounding polygons that already include ascenders and descenders, so no
# vertical-padding heuristics are needed. kraken is installed --no-deps on the
# Space (its torch~=2.1 metadata pin conflicts with our torch==2.2.2, which it
# runs on fine), so treat any import/model/segment failure as non-fatal and
# fall back to the docTR word-box grouping below.

_kraken_model = None
_kraken_disabled = False

def _get_kraken_model():
    """Load and cache kraken's default blla model (shipped as package data)."""
    global _kraken_model
    if _kraken_model is None:
        import os
        import kraken
        from kraken.lib import vgsl
        logger.info("Loading kraken blla segmentation model...")
        path = os.path.join(os.path.dirname(kraken.__file__), 'blla.mlmodel')
        _kraken_model = vgsl.TorchVGSLModel.load_model(path)
    return _kraken_model


def _kraken_lines_to_bboxes(seg, page_w, page_h):
    """Convert a kraken Segmentation (or legacy dict) to reading-order bboxes."""
    lines = getattr(seg, 'lines', None)
    if lines is None and isinstance(seg, dict):
        lines = seg.get('lines', [])
    out = []
    for ln in lines or []:
        poly = getattr(ln, 'boundary', None)
        if poly is None and isinstance(ln, dict):
            poly = ln.get('boundary')
        if not poly:
            continue
        xs = [p[0] for p in poly]
        ys = [p[1] for p in poly]
        x0 = max(0, int(min(xs)))
        y0 = max(0, int(min(ys)))
        x1 = min(page_w, int(max(xs)) + 1)
        y1 = min(page_h, int(max(ys)) + 1)
        if (x1 - x0) >= 4 and (y1 - y0) >= 4:
            out.append({'x': x0, 'y': y0, 'w': x1 - x0, 'h': y1 - y0})
    return out


def _kraken_segment(bgr_img, page_w, page_h):
    """Run kraken blla on a BGR ndarray; returns bbox lines or None on failure."""
    global _kraken_disabled
    if _kraken_disabled:
        return None
    try:
        import cv2
        from PIL import Image
        from kraken import blla
        pil = Image.fromarray(cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB))
        seg = blla.segment(pil, model=_get_kraken_model())
        lines = _kraken_lines_to_bboxes(seg, page_w, page_h)
        if lines:
            return lines
        logger.warning("kraken found no lines; falling back to docTR for this page")
        return None
    except ImportError as e:
        _kraken_disabled = True
        logger.warning("kraken unavailable (%s); docTR fallback active for this process", e)
        return None
    except Exception as e:
        logger.warning("kraken segmentation failed (%s); docTR fallback for this page", e)
        return None


# ── Layout-aware line segmentation (docTR text detection, fallback) ──────────
# The naive projection segmenter fails on phone photos, ruled paper, tables
# and diagrams. docTR's pretrained DB-ResNet detector finds word boxes
# robustly; the grouping below merges them into reading-order lines using the
# same tolerance rule validated in the dissertation corpus pipeline:
#   tau = min(alpha * median(word_height), beta * page_height)

_det_model = None

def _get_detector():
    global _det_model
    if _det_model is None:
        from doctr.models import detection_predictor
        logger.info("Loading docTR detection model (db_resnet50)...")
        _det_model = detection_predictor('db_resnet50', pretrained=True)
    return _det_model


def _group_words_into_lines(boxes, page_w, page_h,
                            alpha=0.6, beta=0.012, min_gap_px=18):
    """boxes: list of (x0, y0, x1, y1) absolute pixels -> list of line bboxes."""
    import numpy as np
    if not boxes:
        return []
    heights = np.array([b[3] - b[1] for b in boxes])
    # Exclude outlier-tall boxes (headers, diagram labels) from the estimate
    med_h = float(np.median(heights[heights <= np.percentile(heights, 90)]))         if len(heights) > 2 else float(np.median(heights))
    tau = min(alpha * med_h, max(min_gap_px, beta * page_h))

    boxes = sorted(boxes, key=lambda b: (b[1] + b[3]) / 2)
    lines = []
    for b in boxes:
        cy = (b[1] + b[3]) / 2
        placed = False
        for line in lines:
            if abs(cy - line['cy']) <= tau:
                line['boxes'].append(b)
                cys = [(x[1] + x[3]) / 2 for x in line['boxes']]
                line['cy'] = sum(cys) / len(cys)
                placed = True
                break
        if not placed:
            lines.append({'cy': cy, 'boxes': [b]})

    out = []
    # Expand each line box the way baseline-based systems (Transkribus'
    # ARU-Net, kraken's blla) build their bounding polygons: detectors fire
    # on the x-height band, so without generous vertical margins the crops
    # decapitate ascenders and cut descenders - exactly the strokes cursive
    # readers need. ~30% of the line height above and ~35% below recovers
    # them; horizontal padding stays small.
    pad_x = max(2, int(med_h * 0.15))
    for line in sorted(lines, key=lambda l: l['cy']):
        xs0 = min(b[0] for b in line['boxes']); ys0 = min(b[1] for b in line['boxes'])
        xs1 = max(b[2] for b in line['boxes']); ys1 = max(b[3] for b in line['boxes'])
        line_h = ys1 - ys0
        pad_top = int(line_h * 0.30)
        pad_bot = int(line_h * 0.35)
        y = max(0, int(ys0 - pad_top))
        out.append({
            'x': max(0, int(xs0 - pad_x)), 'y': y,
            'w': min(page_w, int(xs1 + pad_x)) - max(0, int(xs0 - pad_x)),
            'h': min(page_h, int(ys1 + pad_bot)) - y,
        })
    return out


@app.post("/api/v1/segment")
async def segment_endpoint(file: UploadFile = File(...)):
    """Detect handwriting lines on a full page; returns reading-order bboxes.

    Primary engine: kraken blla baseline segmentation. Fallback: docTR word
    boxes grouped into lines. Response:
    {"width": W, "height": H, "engine": "kraken"|"doctr", "lines": [{"x","y","w","h"}, ...]}
    """
    data = await file.read()
    if not data:
        raise HTTPException(status_code=400, detail="Empty file.")
    img = preprocess.decode_image(data)
    if img is None:
        raise HTTPException(status_code=422, detail="Could not decode image.")

    h, w = img.shape[:2]

    kraken_lines = _kraken_segment(img, w, h)
    if kraken_lines is not None:
        return JSONResponse({"width": w, "height": h, "engine": "kraken",
                             "lines": kraken_lines})

    import cv2
    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    det = _get_detector()
    result = det([rgb])
    # doctr returns normalized (x0, y0, x1, y1[, score]) per word
    raw = result[0].get('words', result[0]) if isinstance(result[0], dict) else result[0]
    abs_boxes = []
    for b in raw:
        x0, y0, x1, y1 = float(b[0]) * w, float(b[1]) * h, float(b[2]) * w, float(b[3]) * h
        if (x1 - x0) >= 4 and (y1 - y0) >= 4:
            abs_boxes.append((x0, y0, x1, y1))

    lines = _group_words_into_lines(abs_boxes, w, h)
    return JSONResponse({"width": w, "height": h, "engine": "doctr", "lines": lines})