document.py

from __future__ import annotations

import tempfile
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path

import fitz
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageSequence

from .config import SETTINGS

SUPPORTED_SUFFIXES = {".pdf", ".png", ".jpg", ".jpeg", ".tif", ".tiff"}


@dataclass
class RenderedDocument:
    source_name: str
    pages: list[Image.Image]
    vector_text: list[str]


@dataclass
class Tile:
    page: int
    index: int
    image: Image.Image
    x: int
    y: int
    width: int
    height: int
    page_width: int
    page_height: int


def _validate_file(path: Path) -> None:
    if not path.exists() or not path.is_file():
        raise ValueError("The uploaded file is no longer available.")
    if path.suffix.lower() not in SUPPORTED_SUFFIXES:
        raise ValueError(
            "Unsupported file type. Upload PDF, PNG, JPG, JPEG, TIF, or TIFF."
        )
    size_mb = path.stat().st_size / (1024 * 1024)
    if size_mb > SETTINGS.max_upload_mb:
        raise ValueError(
            f"File is {size_mb:.1f} MB; the configured limit is "
            f"{SETTINGS.max_upload_mb} MB."
        )


def load_document(
    file_path: str | Path,
    dpi: int = 300,
    max_pages: int | None = None,
) -> RenderedDocument:
    path = Path(file_path)
    _validate_file(path)
    page_limit = min(max_pages or SETTINGS.max_pages, SETTINGS.max_pages)
    if dpi not in {150, 200, 300, 450, 600}:
        raise ValueError("DPI must be one of 150, 200, 300, 450, or 600.")

    pages: list[Image.Image] = []
    vector_text: list[str] = []

    if path.suffix.lower() == ".pdf":
        document = fitz.open(path)
        try:
            if document.page_count > page_limit:
                raise ValueError(
                    f"Drawing has {document.page_count} pages. This Space is "
                    f"configured for at most {page_limit} pages per run."
                )
            scale = dpi / 72.0
            matrix = fitz.Matrix(scale, scale)
            for page in document:
                pixmap = page.get_pixmap(matrix=matrix, alpha=False)
                if pixmap.width * pixmap.height > SETTINGS.max_image_pixels:
                    raise ValueError(
                        f"Rendered page exceeds {SETTINGS.max_image_pixels:,} pixels. "
                        "Reduce the DPI or split the drawing."
                    )
                image = Image.frombytes(
                    "RGB", (pixmap.width, pixmap.height), pixmap.samples
                )
                pages.append(image)
                vector_text.append(_extract_ordered_pdf_text(page))
        finally:
            document.close()
    else:
        with Image.open(path) as source:
            for frame_index, frame in enumerate(ImageSequence.Iterator(source)):
                if frame_index >= page_limit:
                    raise ValueError(
                        f"Image contains more than the configured {page_limit} frames."
                    )
                image = ImageOps.exif_transpose(frame.copy()).convert("RGB")
                if image.width * image.height > SETTINGS.max_image_pixels:
                    raise ValueError(
                        f"Image exceeds {SETTINGS.max_image_pixels:,} pixels."
                    )
                pages.append(image)
                vector_text.append("")

    if not pages:
        raise ValueError("No readable drawing page was found in the uploaded file.")
    return RenderedDocument(path.name, pages, vector_text)


def _extract_ordered_pdf_text(page: fitz.Page) -> str:
    words = page.get_text("words", sort=True)
    if not words:
        return ""
    lines: dict[tuple[int, int, int], list[tuple[int, str]]] = {}
    for item in words:
        block, line, word_index = int(item[5]), int(item[6]), int(item[7])
        lines.setdefault((block, line, 0), []).append((word_index, str(item[4])))
    ordered: list[str] = []
    for key in sorted(lines):
        ordered.append(" ".join(text for _, text in sorted(lines[key])))
    return "\n".join(ordered)


def tile_image(
    image: Image.Image,
    page_number: int,
    tile_size: int = 1800,
    overlap: float = 0.18,
) -> Iterator[Tile]:
    """Yield overlapping square-ish tiles without losing page-edge content."""
    tile_size = max(768, min(3072, int(tile_size)))
    overlap = max(0.0, min(0.4, float(overlap)))
    width, height = image.size
    if width <= tile_size and height <= tile_size:
        yield Tile(page_number, 0, image.copy(), 0, 0, width, height, width, height)
        return

    step = max(1, int(tile_size * (1 - overlap)))
    x_positions = _axis_positions(width, tile_size, step)
    y_positions = _axis_positions(height, tile_size, step)
    index = 0
    for y in y_positions:
        for x in x_positions:
            right = min(width, x + tile_size)
            bottom = min(height, y + tile_size)
            crop = image.crop((x, y, right, bottom))
            yield Tile(
                page_number,
                index,
                crop,
                x,
                y,
                crop.width,
                crop.height,
                width,
                height,
            )
            index += 1


def _axis_positions(length: int, window: int, step: int) -> list[int]:
    if length <= window:
        return [0]
    positions = list(range(0, max(1, length - window + 1), step))
    last = length - window
    if positions[-1] != last:
        positions.append(last)
    return positions


def create_sample_drawing() -> str:
    """Create a deterministic test drawing; it is clearly labelled as a sample."""
    output = Path(tempfile.gettempdir()) / "drawing_intelligence_sample.png"
    if output.exists():
        return str(output)

    canvas = Image.new("RGB", (1800, 1180), "white")
    draw = ImageDraw.Draw(canvas)
    font = _font(26)
    small = _font(21)
    title = _font(34, bold=True)
    red = "#c9202e"
    charcoal = "#151b24"
    gray = "#75808f"

    draw.text((55, 35), "DRAWING INTELLIGENCE — DEMONSTRATION DRAWING", fill=charcoal, font=title)
    draw.text(
        (55, 82),
        "Synthetic training-free sample · not a production engineering drawing",
        fill=gray,
        font=small,
    )
    draw.rectangle((45, 125, 1750, 1125), outline="#aeb5bf", width=3)

    # Simplified cylinder-head-like orthographic view.
    body = (250, 310, 1260, 805)
    draw.rounded_rectangle(body, radius=35, outline=charcoal, width=5)
    for cx in (450, 750, 1050):
        draw.ellipse((cx - 125, 410, cx + 125, 660), outline=charcoal, width=4)
        for dx, dy in ((-48, -48), (48, -48), (-48, 48), (48, 48)):
            draw.ellipse(
                (cx + dx - 28, 535 + dy - 28, cx + dx + 28, 535 + dy + 28),
                outline=charcoal,
                width=3,
            )
    for x in range(310, 1230, 115):
        draw.ellipse((x - 16, 330, x + 16, 362), outline=charcoal, width=3)
        draw.ellipse((x - 16, 752, x + 16, 784), outline=charcoal, width=3)
    draw.line((250, 850, 1260, 850), fill=charcoal, width=3)
    draw.line((250, 835, 250, 865), fill=charcoal, width=3)
    draw.line((1260, 835, 1260, 865), fill=charcoal, width=3)
    draw.text((715, 860), "215", fill=charcoal, font=small)

    callouts = [
        ("1", (385, 208), "4X ⌀12 H7 THRU", (440, 204)),
        ("5", (1030, 225), "M14×1.5-6H", (1090, 220)),
        ("7", (650, 185), "⌖ ⌀0.10 M | A | B | C", (710, 180)),
        ("14", (1340, 415), "Ra 1.6", (1410, 410)),
        ("27", (710, 705), "⌀12 H7", (760, 742)),
        ("32", (350, 980), "BREAK ALL SHARP EDGES 0.5×45°", (410, 975)),
    ]
    anchors = {
        "1": (420, 340),
        "5": (1100, 345),
        "7": (750, 310),
        "14": (1255, 520),
        "27": (735, 745),
        "32": (680, 800),
    }
    for balloon_id, center, text, text_at in callouts:
        _draw_balloon(draw, center, balloon_id, red, small)
        draw.line((center[0], center[1] + 22, *anchors[balloon_id]), fill=charcoal, width=3)
        draw.text(text_at, text, fill=charcoal, font=font)

    # Title block.
    draw.rectangle((1265, 805, 1750, 1125), outline=charcoal, width=3)
    for y in (875, 945, 1015):
        draw.line((1265, y, 1750, y), fill=charcoal, width=2)
    draw.text((1290, 825), "PART: SAMPLE CYLINDER HEAD", fill=charcoal, font=small)
    draw.text((1290, 895), "MATERIAL: FG260", fill=charcoal, font=small)
    draw.text((1290, 965), "UNLESS SPECIFIED: ±0.20", fill=charcoal, font=small)
    draw.text((1290, 1035), "SHEET 1 OF 1 · mm", fill=charcoal, font=small)
    canvas.save(output, optimize=True)
    return str(output)


def _draw_balloon(
    draw: ImageDraw.ImageDraw,
    center: tuple[int, int],
    label: str,
    color: str,
    font: ImageFont.FreeTypeFont | ImageFont.ImageFont,
) -> None:
    x, y = center
    radius = 24
    draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color)
    box = draw.textbbox((0, 0), label, font=font)
    text_width, text_height = box[2] - box[0], box[3] - box[1]
    draw.text(
        (x - text_width / 2, y - text_height / 2 - 3),
        label,
        fill="white",
        font=font,
    )


def _font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    candidates = [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
        if bold
        else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation2/LiberationSans-Bold.ttf"
        if bold
        else "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf",
    ]
    for candidate in candidates:
        try:
            return ImageFont.truetype(candidate, size)
        except OSError:
            continue
    return ImageFont.load_default()