"""Local OCR/layout stage producing independently observed lines."""

from __future__ import annotations

import csv
import io
import os
import re
import subprocess
import tempfile
import statistics
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any

from PIL import Image, ImageEnhance, ImageOps
from pdf2image import convert_from_bytes


MONEY_RE = re.compile(r"(?<!\d)(?:\$\s*)?-?\d+[.,]\d{2}(?!\d)")
SUMMARY_RE = re.compile(r"\b(?:subtotal|total|gst|tax|surcharge|rounding|change|eft|card|cash)\b", re.I)
PRODUCT_CODE_RE = re.compile(r"^\s*(\d{4,8})\b")
WEIGHT_RE = re.compile(r"^\s*\d+(?:\.\d+)?\s*kg\b.*\bNet\s*@.*\$?\s*/\s*k[ga]\b", re.I)


def extract_layout(data: bytes, mimetype: str) -> dict[str, Any]:
    image = _image(data, mimetype)
    image = ImageOps.autocontrast(ImageOps.grayscale(image))
    if image.width < 900:
        scale = 900 / image.width
        image = image.resize((900, max(1, int(image.height * scale))))
    if image.height > 4000:
        scale = 4000 / image.height
        image = image.resize((max(1, int(image.width * scale)), 4000))
    image = ImageEnhance.Contrast(image).enhance(1.35)
    tile_height = 500
    tiles = [(top, image.crop((0, top, image.width, min(image.height, top + tile_height))))
             for top in range(0, image.height, tile_height)]
    with tempfile.TemporaryDirectory(prefix="invoice-ocr-") as temporary:
        def recognize(tile):
            top, tile_image = tile
            source = Path(temporary) / f"source-{top}.png"
            tile_image.save(source, "PNG")
            completed = subprocess.run(
                ["tesseract", str(source), "stdout", "-l", "eng", "--psm", "6", "tsv"],
                check=True, capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=20,
                env={**os.environ, "OMP_THREAD_LIMIT": "1"},
            )
            return top, completed.stdout
        with ThreadPoolExecutor(max_workers=2) as executor:
            outputs = sorted(executor.map(recognize, tiles), key=lambda value: value[0])
    lines = []
    for top, output in outputs:
        lines.extend(_parse_tsv(output, image.width, image.height, offset_y=top))
    lines.sort(key=lambda line: (line["estimated_bounding_region"][1], line["estimated_bounding_region"][0]))
    for index, line in enumerate(lines, 1):
        line["source_line_id"] = f"ocr-line-{index}"
    classified = classify_lines(lines)
    monetary = associate_monetary_values(classified)
    return {
        "engine": "tesseract", "engine_version": _tesseract_version(),
        "coordinate_space": "normalized_0_1", "lines": classified,
        "item_primary_lines": [line for line in classified if line["classification"] == "item_primary"],
        "item_weight_continuations": [line for line in classified if line["classification"] == "item_weight_continuation"],
        "item_line_candidates": [line for line in classified if line["classification"] in {"item_primary", "item_weight_continuation"}],
        "line_total_column_x": monetary["line_total_column_x"],
    }


def classify_lines(lines: list[dict[str, Any]]) -> list[dict[str, Any]]:
    previous_primary = None
    result = []
    for original in lines:
        line = dict(original); raw = str(line.get("raw_text") or "")
        code_match = PRODUCT_CODE_RE.match(raw)
        money = list(line.get("monetary_candidates") or [])
        tax_matches = re.findall(r"\b([AB])\b", raw, re.I)
        description_text = PRODUCT_CODE_RE.sub("", raw, count=1)
        description_text = MONEY_RE.sub("", description_text)
        has_description = len(re.findall(r"[A-Za-z]{2,}", description_text)) >= 1
        product_code = code_match.group(1) if code_match else None
        tax_code = tax_matches[-1].upper() if tax_matches else None
        score = (3 if product_code else 0) + (2 if has_description else 0) + (2 if money else 0) + (2 if tax_code else 0)
        if WEIGHT_RE.search(raw):
            classification = "item_weight_continuation"
            attached = previous_primary
        elif not SUMMARY_RE.search(raw) and has_description and (
            bool(product_code) or (bool(money) and bool(tax_code) and score >= 6)
        ):
            classification = "item_primary"
            attached = None; previous_primary = line.get("source_line_id")
        else:
            classification = "non_item"
            attached = None
        missing_fields = []
        if classification == "item_primary":
            if not money:
                missing_fields.append("line_total")
            if not tax_code:
                missing_fields.append("tax_code")
        line.update({
            "product_code": product_code, "tax_code": tax_code,
            "final_monetary_value": money[-1] if money else None,
            "row_classification_score": score, "classification": classification,
            "attached_to_source_line_id": attached,
            "complete": classification == "item_primary" and not missing_fields,
            "missing_fields": missing_fields,
            "parsing_status": "complete" if classification == "item_primary" and not missing_fields else
                "incomplete" if classification == "item_primary" else "continuation" if classification == "item_weight_continuation" else "not_applicable",
        })
        result.append(line)
    return result


def transcript(layout: dict[str, Any], *, candidates_only: bool = False) -> str:
    values = layout.get("item_line_candidates" if candidates_only else "lines") or []
    return "\n".join(
        f'{line["source_line_id"]} {line["raw_text"]}' for line in values
        if isinstance(line, dict) and line.get("raw_text")
    )[:16000]


def associate_items(layout: dict[str, Any], items: list[dict[str, Any]]) -> dict[str, Any]:
    candidates = [dict(line) for line in layout.get("item_primary_lines") or layout.get("item_line_candidates") or []
                  if line.get("classification") != "item_weight_continuation"]
    for line in candidates:
        line.update({"consumed_by_item_id": None, "parsing_status": "unconsumed"})
    used: set[int] = set()
    for item_index, item in enumerate(items):
        description = _tokens(str(item.get("description") or ""))
        source = _tokens(str(item.get("source_text") or ""))
        expected = description | source
        expected_code = str(item.get("product_code") or "")
        best_index, best_score = None, 0.0
        for index, line in enumerate(candidates):
            if index in used:
                continue
            actual = _tokens(line["raw_text"])
            score = len(expected & actual) / max(1, len(description or expected))
            if expected_code and expected_code == str(line.get("product_code") or ""):
                score += 1.0
            if item.get("line_total") is not None and str(item.get("line_total")) in {
                str(value).replace("$", "") for value in line.get("monetary_candidates") or []
            }:
                score += 0.5
            if score > best_score:
                best_index, best_score = index, score
        if best_index is not None and best_score >= 0.35:
            used.add(best_index)
            candidates[best_index]["consumed_by_item_id"] = f"item-{item_index + 1}"
            candidates[best_index]["parsing_status"] = "consumed"
            candidates[best_index]["association_confidence"] = round(best_score, 4)
    unconsumed = [line for line in candidates if line["parsing_status"] == "unconsumed"]
    return {"visual_item_lines": candidates, "unconsumed_source_lines": unconsumed}


def associate_monetary_values(lines: list[dict[str, Any]]) -> dict[str, Any]:
    """Select line totals by row geometry; never borrow from adjacent rows."""
    ordered = sorted(lines, key=lambda line: _line_center_y(line))
    centers = [_line_center_y(line) for line in ordered]
    for index, line in enumerate(ordered):
        lower = 0.0 if index == 0 else (centers[index - 1] + centers[index]) / 2
        upper = 1.0 if index == len(ordered) - 1 else (centers[index] + centers[index + 1]) / 2
        line["exclusive_vertical_interval"] = [round(lower, 6), round(upper, 6)]
    for source in lines:
        for token in source.get("monetary_tokens") or []:
            if token.get("center_y") is None:
                box = token.get("bbox") or [0, _line_center_y(source), 0, _line_center_y(source)]
                token["center_y"] = (float(box[1]) + float(box[3])) / 2
            owners = [line for line in ordered
                      if line["exclusive_vertical_interval"][0] <= float(token["center_y"]) < line["exclusive_vertical_interval"][1]]
            owner = owners[0] if owners else min(ordered, key=lambda line: abs(_line_center_y(line) - float(token["center_y"])))
            token["owner_source_line_id"] = owner.get("source_line_id")
            token["ownership"] = "exclusive"
    anchors = []
    for line in lines:
        tokens = _owned_monetary_tokens(lines, line.get("source_line_id"))
        if line.get("classification") == "item_primary" and line.get("tax_code") and line.get("confidence", 0) >= 0.6 and tokens:
            anchors.append(float(max(tokens, key=lambda token: token["bbox"][2])["center_x"]))
    column_x = statistics.median(anchors) if anchors else None
    by_id = {line.get("source_line_id"): line for line in lines}
    for line in lines:
        tokens = _owned_monetary_tokens(lines, line.get("source_line_id"))
        line["owned_monetary_tokens"] = tokens
        for token in tokens:
            token["semantic_role"] = "unknown"
        if line.get("classification") == "item_primary":
            valid = [token for token in tokens if column_x is None or abs(float(token["center_x"]) - column_x) <= 0.10]
            selected = max(valid, key=lambda token: token["bbox"][2]) if valid else None
            if selected:
                selected["semantic_role"] = "line_total"
            line["selected_line_total_token"] = selected
        elif line.get("classification") == "item_weight_continuation":
            weight_match = re.search(r"(\d+(?:\.\d+)?)\s*kg", str(line.get("raw_text") or ""), re.I)
            line["weight_quantity"] = float(weight_match.group(1)) if weight_match else None
            rate = tokens[-1] if tokens else None
            if rate:
                rate["semantic_role"] = "unit_price"
                rate["exclusive_to_weight_continuation"] = True
            line["selected_unit_rate_token"] = rate
            parent = by_id.get(line.get("attached_to_source_line_id"))
            if parent is not None:
                parent["weight_continuation"] = line
    return {"line_total_column_x": None if column_x is None else round(column_x, 5)}


def _line_center_y(line: dict[str, Any]) -> float:
    box = line.get("estimated_bounding_region") or [0, 0, 0, 0]
    return (float(box[1]) + float(box[3])) / 2


def _owned_monetary_tokens(lines: list[dict[str, Any]], source_line_id: Any) -> list[dict[str, Any]]:
    return [token for line in lines for token in line.get("monetary_tokens") or []
            if token.get("owner_source_line_id") == source_line_id]


def crop_exclusive_row(data: bytes, mimetype: str, interval: list[float]) -> bytes:
    """Crop exactly one normalized row interval; adjacent rows are excluded."""
    image = _image(data, mimetype)
    lower = max(0.0, min(1.0, float(interval[0]))); upper = max(lower, min(1.0, float(interval[1])))
    cropped = image.crop((0, int(image.height * lower), image.width, max(int(image.height * upper), int(image.height * lower) + 1)))
    stream = io.BytesIO(); cropped.save(stream, "PNG")
    return stream.getvalue()


def _parse_tsv(value: str, width: int, height: int, *, offset_y: int = 0) -> list[dict[str, Any]]:
    grouped: dict[tuple[str, ...], list[dict[str, str]]] = {}
    for row in csv.DictReader(io.StringIO(value), delimiter="\t"):
        text = str(row.get("text") or "").strip()
        if not text or int(float(row.get("conf") or -1)) < 0:
            continue
        key = tuple(str(row.get(name) or "0") for name in ("page_num", "block_num", "par_num", "line_num"))
        grouped.setdefault(key, []).append(row)
    result = []
    for index, words in enumerate(grouped.values(), 1):
        raw = " ".join(str(word["text"]).strip() for word in words).strip()
        left = min(int(word["left"]) for word in words); top = min(int(word["top"]) for word in words)
        right = max(int(word["left"]) + int(word["width"]) for word in words)
        bottom = max(int(word["top"]) + int(word["height"]) for word in words)
        confidence = sum(max(0.0, float(word["conf"])) for word in words) / (100 * len(words))
        monetary_tokens = []
        for word in words:
            for match in MONEY_RE.finditer(str(word.get("text") or "")):
                word_left = int(word["left"]); word_top = int(word["top"]) + offset_y
                word_right = word_left + int(word["width"]); word_bottom = word_top + int(word["height"])
                bbox = [round(word_left / width, 5), round(word_top / height, 5),
                        round(word_right / width, 5), round(word_bottom / height, 5)]
                monetary_tokens.append({"raw": match.group(0), "bbox": bbox,
                                        "center_x": round((bbox[0] + bbox[2]) / 2, 5),
                                        "center_y": round((bbox[1] + bbox[3]) / 2, 5),
                                        "confidence": round(max(0.0, float(word["conf"])) / 100, 4)})
        result.append({
            "source_line_id": f"ocr-line-{index}", "raw_text": raw,
            "estimated_bounding_region": [round(left / width, 5), round((top + offset_y) / height, 5), round(right / width, 5), round((bottom + offset_y) / height, 5)],
            "bounding_region_origin": "tesseract_observed", "confidence": round(confidence, 4),
            "monetary_candidates": MONEY_RE.findall(raw), "alternative_monetary_candidates": MONEY_RE.findall(raw),
            "monetary_tokens": monetary_tokens,
        })
    return result


def _image(data: bytes, mimetype: str) -> Image.Image:
    if mimetype == "application/pdf":
        pages = convert_from_bytes(data, dpi=250, first_page=1, last_page=1)
        if not pages:
            raise ValueError("ocr_pdf_has_no_pages")
        return pages[0].convert("RGB")
    return Image.open(io.BytesIO(data)).convert("RGB")


def _tokens(value: str) -> set[str]:
    return {token for token in re.findall(r"[a-z0-9]+", value.lower()) if len(token) > 1 and not MONEY_RE.fullmatch(token)}


def _tesseract_version() -> str:
    try:
        return subprocess.run(["tesseract", "--version"], capture_output=True, text=True, timeout=5).stdout.splitlines()[0]
    except Exception:
        return "unknown"