import io
import math
import statistics
import time
from dataclasses import dataclass
from typing import Callable

from .engine import EngineLine, OcrEngine


RENDER_DPI = 200
MAX_RENDER_PIXELS = 25_000_000


class PdfRenderError(ValueError):
    pass


@dataclass(frozen=True)
class RenderedPage:
    page: int
    width: int
    height: int
    png: bytes
    image: object


def render_pdf_pages(pdf: bytes, pages: list[int], max_pixels: int = MAX_RENDER_PIXELS) -> list[RenderedPage]:
    import pypdfium2

    if not pages or pages != sorted(set(pages)) or any(type(page) is not int or page < 1 for page in pages):
        raise PdfRenderError("PDF pages must be unique, ordered, one-based integers")
    try:
        document = pypdfium2.PdfDocument(pdf)
    except Exception as error:
        raise PdfRenderError("PDF cannot be opened for deterministic rendering") from error

    rendered: list[RenderedPage] = []
    try:
        for page_number in pages:
            if page_number > len(document):
                raise PdfRenderError(f"PDF page {page_number} does not exist")
            page = document[page_number - 1]
            try:
                page_width, page_height = page.get_size()
                width = math.ceil(page_width * RENDER_DPI / 72)
                height = math.ceil(page_height * RENDER_DPI / 72)
                if width * height > max_pixels:
                    raise PdfRenderError("Rendered page exceeds the 25 megapixels limit")
                bitmap = page.render(scale=RENDER_DPI / 72)
                try:
                    image = bitmap.to_pil()
                    output = io.BytesIO()
                    image.save(output, format="PNG")
                    rendered.append(RenderedPage(page_number, image.width, image.height, output.getvalue(), image))
                finally:
                    bitmap.close()
            finally:
                page.close()
    finally:
        document.close()
    return rendered


def _metrics(lines: list[EngineLine], text: str, image: object) -> dict[str, int | float]:
    confidences = sorted(line.confidence for line in lines)
    grayscale = image.convert("L")
    grayscale.thumbnail((256, 256))
    pixels = list(grayscale.getdata())
    return {
        "lineCount": len(lines),
        "nonWhitespaceCharacters": sum(not character.isspace() for character in text),
        "inkCoverage": sum(value < 250 for value in pixels) / len(pixels),
        "medianConfidence": statistics.median(confidences) if confidences else 0.0,
        "p10Confidence": confidences[math.floor((len(confidences) - 1) * 0.1)] if confidences else 0.0,
        "lowConfidenceLineRatio": sum(value < 0.5 for value in confidences) / len(confidences) if confidences else 0.0,
    }


def process_pdf(
    job_id: str,
    document_sha256: str,
    pdf: bytes,
    requested_pages: list[int],
    engine: OcrEngine,
    processing_ms: Callable[[int], int] | None = None,
) -> dict:
    results = []
    for rendered in render_pdf_pages(pdf, requested_pages):
        started = time.perf_counter_ns()
        indexed = list(enumerate(engine.recognize(rendered.image), start=1))
        indexed.sort(key=lambda item: (item[1].bbox[1], item[1].bbox[0], item[0]))
        lines = [line for _, line in indexed]
        text = "\n".join(line.text for line in lines)
        serialized_lines = [
            {
                "lineId": f"p{rendered.page}-l{index:02d}-{'-'.join(map(str, line.bbox))}",
                "text": line.text,
                "confidence": line.confidence,
                "bbox": list(line.bbox),
            }
            for index, line in enumerate(lines, start=1)
        ]
        elapsed = math.ceil((time.perf_counter_ns() - started) / 1_000_000)
        results.append({
            "page": rendered.page,
            "width": rendered.width,
            "height": rendered.height,
            "processingMs": processing_ms(rendered.page) if processing_ms else elapsed,
            "text": text,
            "metrics": _metrics(lines, text, rendered.image),
            "lines": serialized_lines,
        })
    return {
        "schemaVersion": "1",
        "jobId": job_id,
        "documentSha256": document_sha256,
        "engine": {
            "name": "paddleocr",
            "version": "3.4.0",
            "runtime": "paddlepaddle-3.2.2",
            "device": "cpu",
            "configVersion": "ocr-v1",
            "dpi": RENDER_DPI,
        },
        "pages": results,
    }