"""Turn a file into pages: digital-PDF text where usable, otherwise rendered/loaded images.
Also page-level orientation correction and deskew."""

from dataclasses import dataclass, field
from pathlib import Path

import cv2
import numpy as np

from ..charset import normalize
from ..geometry import merge_line_cells

RENDER_DPI = 200
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp", ".gif"}


@dataclass
class Page:
    number: int
    image: np.ndarray | None              # RGB uint8, in the page coordinate frame
    width: int
    height: int
    text_source: str                      # "ocr" | "pdf_text"
    pdf_segments: list = field(default_factory=list)   # [(x0, y0, x1, y1, text)] when text_source == "pdf_text"


def _usable_pdf_text(segs) -> bool:
    text = "".join(s[4] for s in segs)
    if len(text.strip()) < 25:
        return False
    good = sum(c.isalnum() or c in " .,:;-$€£%/()#&@'\"" for c in text)
    return good / len(text) > 0.8 and "(cid:" not in text


def load_pages(path, dpi=RENDER_DPI, force_ocr=False):
    path = Path(path)
    suffix = path.suffix.lower()
    if suffix == ".pdf":
        import pypdfium2 as pdfium

        doc = pdfium.PdfDocument(str(path))
        scale = dpi / 72
        for i in range(len(doc)):
            page = doc[i]
            pw, ph = page.get_size()
            W, H = round(pw * scale), round(ph * scale)
            segs = []
            if not force_ocr:
                tp = page.get_textpage()
                cells = []
                for r in range(tp.count_rects()):
                    l, b, rr, t = tp.get_rect(r)
                    txt = normalize(tp.get_text_bounded(l, b, rr, t))
                    if txt:
                        cells.append((l * scale, (ph - t) * scale, rr * scale, (ph - b) * scale, txt))
                segs = [s for s in merge_line_cells(cells) if s[4].strip()]
                tp.close()
            if segs and _usable_pdf_text(segs):
                yield Page(i + 1, None, W, H, "pdf_text", segs)
            else:
                img = page.render(scale=scale).to_numpy()[:, :, :3].copy()
                yield Page(i + 1, img, img.shape[1], img.shape[0], "ocr")
            page.close()
        doc.close()
    elif suffix in IMAGE_EXTS:
        img = cv2.imread(str(path), cv2.IMREAD_COLOR)
        if img is None:
            from PIL import Image
            img = np.array(Image.open(path).convert("RGB"))[:, :, ::-1]
        img = np.ascontiguousarray(img[:, :, ::-1])
        yield Page(1, img, img.shape[1], img.shape[0], "ocr")
    else:
        raise ValueError(f"unsupported file type: {suffix} (supported: .pdf, {', '.join(sorted(IMAGE_EXTS))})")


def estimate_skew(img_rgb, max_angle=12.0) -> float:
    """Projection-profile skew estimate in degrees (positive = text rotated counter-clockwise)."""
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    s = 1000 / max(gray.shape)
    if s < 1:
        gray = cv2.resize(gray, (round(gray.shape[1] * s), round(gray.shape[0] * s)), interpolation=cv2.INTER_AREA)
    ink = cv2.adaptiveThreshold(gray, 1, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 15).astype(np.float32)
    # the 1-degree search runs at half size (a quarter of the pixels); the 0.1-degree refinement at full size
    small = cv2.resize(ink, (ink.shape[1] // 2, ink.shape[0] // 2), interpolation=cv2.INTER_AREA)

    def score(a, im):
        h, w = im.shape
        m = cv2.getRotationMatrix2D((w / 2, h / 2), -a, 1.0)
        prof = cv2.warpAffine(im, m, (w, h)).sum(1)
        return float(np.var(prof))

    coarse = np.arange(-max_angle, max_angle + 1e-6, 1.0)
    best = coarse[int(np.argmax([score(a, small) for a in coarse]))]
    fine = np.arange(best - 1, best + 1 + 1e-6, 0.1)
    return float(fine[int(np.argmax([score(a, ink) for a in fine]))])


def rotate_image(img, angle):
    """Rotate by `angle` degrees (counter-clockwise), expanding the canvas. Returns (image, 2x3 matrix)."""
    h, w = img.shape[:2]
    m = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0)
    cos, sin = abs(m[0, 0]), abs(m[0, 1])
    nw, nh = int(h * sin + w * cos), int(h * cos + w * sin)
    m[0, 2] += nw / 2 - w / 2
    m[1, 2] += nh / 2 - h / 2
    border = tuple(int(v) for v in np.median(img[::8, ::8].reshape(-1, 3), 0))   # page colour from a pixel sample
    return cv2.warpAffine(img, m, (nw, nh), flags=cv2.INTER_LINEAR, borderValue=border), m


def rot90_matrix(k, w, h):
    """2x3 matrix mapping original coords to coords after np.rot90(img, k) (counter-clockwise k times)."""
    m = np.eye(3)
    for _ in range(k % 4):
        step = np.array([[0, 1, 0], [-1, 0, w], [0, 0, 1]], np.float64)  # (x, y) -> (y, w - x)
        m = step @ m
        w, h = h, w
    return m[:2]