from __future__ import annotations

from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

from ocr_studio.config import ARABIC_FONT_PATH, LATIN_FONT_PATH
from ocr_studio.language import contains_arabic_script
from ocr_studio.spotting import TextSpan


def _shape(text: str) -> str:
    if not contains_arabic_script(text):
        return text
    try:
        from arabic_reshaper import reshape
        from bidi.algorithm import get_display

        return get_display(reshape(text))
    except Exception:
        return text


def _font(size: int, text: str) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    path = ARABIC_FONT_PATH if contains_arabic_script(text) else LATIN_FONT_PATH
    if path.exists():
        return ImageFont.truetype(str(path), size=size)
    return ImageFont.load_default()


def annotate_page(image: Image.Image, spans: list[TextSpan]) -> Image.Image:
    annotated = image.convert("RGB").copy()
    if not spans:
        return annotated
    draw = ImageDraw.Draw(annotated, "RGBA")
    for span in spans:
        if not span.box:
            continue
        x0, y0, x1, y1 = span.box
        draw.rectangle((x0, y0, x1, y1), outline=(31, 79, 74, 220), width=2)
        label = span.text.strip().replace("\n", " ")
        if not label:
            continue
        if len(label) > 42:
            label = label[:41] + "…"
        font = _font(max(12, int((y1 - y0) * 0.45)), label)
        shaped = _shape(label)
        draw.text((x0 + 3, max(0, y0 - 16)), shaped, fill=(31, 79, 74, 255), font=font)
    return annotated


def annotate_pages(pages: list[Image.Image], spans_by_page: list[list[TextSpan]]) -> list[Image.Image]:
    images: list[Image.Image] = []
    for index, page in enumerate(pages):
        spans = spans_by_page[index] if index < len(spans_by_page) else []
        images.append(annotate_page(page, spans))
    return images