"""Calibration and task-aware RKNN export helpers."""

import ast
import math
from pathlib import Path

import cv2
import onnx
import yaml


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def calibration_images(path, limit):
    """Read image directories, image lists, or a dataset YAML without downloading data."""
    path = Path(path).resolve()
    if path.is_dir():
        images = sorted(p for p in path.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES)
    elif path.suffix.lower() in {".yaml", ".yml"}:
        document = yaml.safe_load(path.read_text())
        root = Path(document.get("path", "."))
        if not root.is_absolute():
            root = path.parent / root
        split = document.get("val") or document.get("train")
        if not split:
            raise ValueError("Calibration YAML needs val or train image paths")
        sources = split if isinstance(split, list) else [split]
        images = []
        for source in sources:
            images.extend(calibration_images(root / source, limit))
    elif path.suffix.lower() == ".txt":
        images = []
        for line in path.read_text().splitlines():
            line = line.strip()
            if line and not line.startswith("#"):
                image = Path(line)
                images.append(image if image.is_absolute() else path.parent / image)
    elif path.suffix.lower() in IMAGE_SUFFIXES:
        images = [path]
    else:
        raise ValueError("Calibration data must be an image directory, list, or dataset YAML")
    images = list(dict.fromkeys(p.resolve() for p in images))[:limit]
    if not images or any(not p.is_file() for p in images):
        raise ValueError("Calibration images are missing or the dataset is empty")
    return images


def prepare_calibration(images, directory, width, height, center_crop):
    """Apply the application's crop/letterbox policy before RKNN calibration."""
    import numpy as np

    directory.mkdir()
    paths = []
    for index, path in enumerate(images):
        image = cv2.imread(str(path), cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError(f"Cannot read calibration image: {path}")
        h, w = image.shape[:2]
        if center_crop:
            ratio = width / height
            crop_w = min(w, int(h * ratio)) & ~1
            crop_h = min(h, int(w / ratio)) & ~1
            if min(crop_w, crop_h) < 2:
                raise ValueError(f"Calibration image is too small: {path}")
            x = ((w - crop_w) // 2) & ~1
            y = ((h - crop_h) // 2) & ~1
            output = cv2.resize(image[y:y + crop_h, x:x + crop_w], (width, height))
        else:
            scale = min(width / w, height / h)
            rw = max(2, int(math.floor(w * scale + 0.5)) & ~1)
            rh = max(2, int(math.floor(h * scale + 0.5)) & ~1)
            x = ((width - rw) // 2) & ~1
            y = ((height - rh) // 2) & ~1
            output = np.full((height, width, 3), 114, dtype=np.uint8)
            output[y:y + rh, x:x + rw] = cv2.resize(image, (rw, rh))
        target = directory / f"{index:05d}.png"
        if not cv2.imwrite(str(target), output):
            raise RuntimeError(f"Cannot write calibration image: {target}")
        paths.append(str(target))
    dataset = directory / "images.txt"
    dataset.write_text("\n".join(paths) + "\n")
    return dataset


def floating_head_tensors(graph, head_index):
    """Keep the complete prediction head and decoding in FP16."""
    prefix = f"/model.{head_index}/"
    tensors = {value for node in graph.node if node.name.startswith(prefix)
               for value in node.output}
    if not tensors:
        raise ValueError("Cannot identify the prediction-head scope in this ONNX export")
    return tensors


def floating_decode_tensors(graph):
    """Keep terminal decoding separate from quantized learned predictions."""
    producers = {name: node for node in graph.node for name in node.output}
    pending = [value.name for value in graph.output]
    protected = set()
    learned = {"Conv", "ConvTranspose", "Gemm", "MatMul"}
    while pending:
        name = pending.pop()
        node = producers.get(name)
        if name in protected or node is None or node.op_type in learned:
            continue
        protected.add(name)
        pending.extend(node.input)
    return protected


def dense_output(graph, metadata, task):
    """Return low-resolution dense predictions; reconstruct sampled pixels in C++."""
    if task not in {"semantic", "depth"}:
        return graph
    producers = {name: node for node in graph.graph.node for name in node.output}
    node = producers[graph.graph.output[0].name]
    while node.op_type in {"Cast", "ArgMax", "Relu"}:
        node = producers[node.input[0]]
    if node.op_type != "Resize":
        raise ValueError("Dense export needs a terminal bilinear Resize")
    attributes = {a.name: onnx.helper.get_attribute_value(a) for a in node.attribute}
    mode = attributes.get("coordinate_transformation_mode", b"half_pixel").decode()
    if attributes.get("mode") != b"linear" or mode not in {"half_pixel", "align_corners"}:
        raise ValueError("Unsupported dense resize convention")
    inferred = onnx.shape_inference.infer_shapes(graph)
    tensor = next((v for v in inferred.graph.value_info if v.name == node.input[0]), None)
    if tensor is None or len(tensor.type.tensor_type.shape.dim) != 4:
        raise ValueError("Cannot determine the dense prediction shape")
    inputs = [value.name for value in inferred.graph.input]
    result = onnx.utils.Extractor(inferred).extract_model(inputs, [node.input[0]])
    metadata["dense_resize"] = mode
    return result


def export_optimized(model, chip, size, data, limit, precision):
    from rknn.api import RKNN

    exported = Path(model.export(format="onnx", imgsz=size, batch=1, device="cpu",
                                 opset=19, simplify=True, end2end=False)).resolve()
    graph = onnx.load(str(exported))
    metadata = {}
    for entry in graph.metadata_props:
        try:
            metadata[entry.key] = ast.literal_eval(entry.value)
        except (ValueError, SyntaxError):
            metadata[entry.key] = entry.value
    graph = dense_output(graph, metadata, model.task)
    onnx.save(graph, str(exported))
    result = exported.with_suffix(".rknn")

    if precision != "fp16":
        images = calibration_images(data, limit)
        dataset = prepare_calibration(images, exported.parent / "calibration", size, size,
                                      model.task == "classify")
        head_index = model.model.model[-1].i
        protected = (floating_head_tensors(graph.graph, head_index)
                     if precision == "mixed" else floating_decode_tensors(graph.graph))
        head_prefix = f"/model.{head_index}/"
        toolkit = RKNN(verbose=False)
        try:
            check(toolkit.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
                                 target_platform=chip), "configure")
            check(toolkit.load_onnx(model=str(exported)), "load ONNX")
            check(toolkit.hybrid_quantization_step1(dataset=str(dataset), proposal=False),
                  "prepare mixed precision")
        finally:
            toolkit.release()
        config_path = exported.with_suffix(".quantization.cfg")
        config = yaml.safe_load(config_path.read_text())
        # RKNN adds layout/conversion aliases such as output0-rs and output0_int8.
        # Protecting only the public output would leave its packed source quantized.
        matched = [key for key, parameter in config["quantize_parameters"].items()
                   if parameter.get("dtype") in {"float32", "float16", "int8", "uint8"}
                   and ((precision == "mixed" and key.startswith(head_prefix)) or any(
                       key == tensor or key.startswith(tensor + suffix)
                       for tensor in protected for suffix in ("_", "-", "#")))]
        if not matched and precision == "mixed":
            raise RuntimeError("Cannot identify floating-point prediction heads in this model")
        config["custom_quantize_layers"] = {key: "float16" for key in matched}
        config_path.write_text(yaml.safe_dump(config, sort_keys=False))

        toolkit = RKNN(verbose=False)
        try:
            check(toolkit.hybrid_quantization_step2(
                  model_input=str(exported.with_suffix(".model")),
                  data_input=str(exported.with_suffix(".data")),
                  model_quantization_cfg=str(config_path)), "build mixed precision")
            check(toolkit.export_rknn(str(result)), "save RKNN")
        finally:
            toolkit.release()
        metadata.update(precision=precision, head_precision="fp16" if precision == "mixed" else "int8",
                        decode_precision="fp16" if matched else "none",
                        fp16_decode_tensors=len(matched), calibration_images=len(images))
        print(f"{precision}: {len(matched)} FP16 head/decode tensors, {len(images)} calibration images")
    else:
        toolkit = RKNN(verbose=False)
        try:
            check(toolkit.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
                                 target_platform=chip), "configure")
            check(toolkit.load_onnx(model=str(exported)), "load ONNX")
            check(toolkit.build(do_quantization=False), "build FP16")
            check(toolkit.export_rknn(str(result)), "save RKNN")
        finally:
            toolkit.release()
        metadata["precision"] = "fp16"
    metadata["target_platform"] = chip
    metadata["args"] = {**metadata.get("args", {}), "format": "rknn",
                        "precision": metadata["precision"]}
    return result, metadata


def check(result, operation):
    if result != 0:
        raise RuntimeError(f"RKNN {operation} failed ({result})")