#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import json
import sys
from pathlib import Path

import cv2
import numpy as np


_MASK_CACHE = {}


def load_image(path: str):
    img = cv2.imread(path, cv2.IMREAD_COLOR)

    if img is None:
        raise RuntimeError(f"No se pudo abrir la imagen: {path}")

    return img


def rotate_image(img, rotation):
    if rotation == 0:
        return img

    if rotation == 90:
        return cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)

    if rotation == 180:
        return cv2.rotate(img, cv2.ROTATE_180)

    if rotation == 270:
        return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)

    return img


def prepare_channels(img):
    """
    Preprocesa la imagen UNA SOLA VEZ por orientación.
    Antes el código hacía cvtColor por cada burbuja, eso lo hacía muy lento.
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    return {
        "gray": gray,
        "hue": hsv[:, :, 0],
        "saturation": hsv[:, :, 1],
        "value": hsv[:, :, 2],
    }


def get_circle_mask(shape, inner_r):
    """
    Reutiliza máscaras circulares para no dibujar el círculo en cada burbuja.
    """
    h, w = shape[:2]
    key = (h, w, int(inner_r))

    if key in _MASK_CACHE:
        return _MASK_CACHE[key]

    mask = np.zeros((h, w), dtype=np.uint8)

    cv2.circle(
        mask,
        (w // 2, h // 2),
        min(int(inner_r), h // 2, w // 2),
        255,
        -1
    )

    inside = mask > 0
    _MASK_CACHE[key] = inside

    return inside


def get_anchor_centers_mm(template):
    anchors_by_name = {a["name"]: a for a in template["anchors"]}

    centers = []

    for name in ["TL", "TR", "BR", "BL"]:
        a = anchors_by_name[name]
        centers.append([
            float(a["x"]) + float(a["w"]) / 2.0,
            float(a["y"]) + float(a["h"]) / 2.0,
        ])

    return np.array(centers, dtype=np.float32)


def detect_anchors(img, template):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    h_img, w_img = gray.shape[:2]

    def build_candidates(threshold_value):
        _, th = cv2.threshold(gray, threshold_value, 255, cv2.THRESH_BINARY_INV)

        kernel = np.ones((3, 3), np.uint8)
        th = cv2.morphologyEx(th, cv2.MORPH_CLOSE, kernel, iterations=1)

        contours, _ = cv2.findContours(th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        candidates = []

        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            area = cv2.contourArea(cnt)
            box_area = max(w * h, 1)

            if w < 8 or h < 8:
                continue

            if w > w_img * 0.16 or h > h_img * 0.16:
                continue

            if area < 45:
                continue

            ratio = w / max(h, 1)
            fill = area / box_area

            if not (0.45 <= ratio <= 1.65):
                continue

            if fill < 0.25:
                continue

            cx = x + w / 2.0
            cy = y + h / 2.0

            square_score = 1.0 - min(abs(1.0 - ratio), 1.0)
            fill_score = min(fill, 1.0)
            size_score = min(area / 500.0, 4.0)

            candidates.append({
                "x": x,
                "y": y,
                "w": w,
                "h": h,
                "area": area,
                "fill": fill,
                "cx": cx,
                "cy": cy,
                "score": square_score * 2.0 + fill_score + size_score,
            })

        return candidates

    def pick_by_corner(candidates):
        corners = {
            "TL": (0.0, 0.0),
            "TR": (float(w_img), 0.0),
            "BR": (float(w_img), float(h_img)),
            "BL": (0.0, float(h_img)),
        }

        picked = {}
        used = set()

        for name, (kx, ky) in corners.items():
            best = None
            best_score = -10**9

            for idx, c in enumerate(candidates):
                if idx in used:
                    continue

                cx = c["cx"]
                cy = c["cy"]

                if name == "TL" and not (cx < w_img * 0.52 and cy < h_img * 0.52):
                    continue

                if name == "TR" and not (cx > w_img * 0.48 and cy < h_img * 0.52):
                    continue

                if name == "BR" and not (cx > w_img * 0.48 and cy > h_img * 0.48):
                    continue

                if name == "BL" and not (cx < w_img * 0.52 and cy > h_img * 0.48):
                    continue

                dist = ((cx - kx) ** 2 + (cy - ky) ** 2) ** 0.5
                max_dist = (w_img ** 2 + h_img ** 2) ** 0.5
                corner_score = 1.0 - (dist / max_dist)

                total_score = c["score"] + corner_score * 5.0

                if total_score > best_score:
                    best_score = total_score
                    best = (idx, c)

            if best is not None:
                idx, c = best
                picked[name] = c
                used.add(idx)

        return picked

    best_picked = {}

    for threshold_value in [35, 45, 55, 65, 80, 95, 110, 125, 145, 165, 185]:
        candidates = build_candidates(threshold_value)
        picked = pick_by_corner(candidates)

        if len(picked) > len(best_picked):
            best_picked = picked

        if len(picked) == 4:
            break

    if len(best_picked) < 4:
        faltan = [name for name in ["TL", "TR", "BR", "BL"] if name not in best_picked]

        raise RuntimeError(
            "No se detectaron los 4 anclajes negros. "
            f"Detectados: {len(best_picked)}. "
            f"Faltan: {', '.join(faltan)}."
        )

    ordered_img = np.array([
        [best_picked["TL"]["cx"], best_picked["TL"]["cy"]],
        [best_picked["TR"]["cx"], best_picked["TR"]["cy"]],
        [best_picked["BR"]["cx"], best_picked["BR"]["cy"]],
        [best_picked["BL"]["cx"], best_picked["BL"]["cy"]],
    ], dtype=np.float32)

    ordered_mm = get_anchor_centers_mm(template)

    H, _ = cv2.findHomography(ordered_img, ordered_mm)

    if H is None:
        raise RuntimeError("No se pudo calcular la homografía de la hoja.")

    H_inv = np.linalg.inv(H)

    return H, H_inv, ordered_img, ordered_mm


def estimate_px_per_mm_from_anchors(anchors_px, anchors_mm):
    anchors_px = np.array(anchors_px, dtype=np.float64)
    anchors_mm = np.array(anchors_mm, dtype=np.float64)

    def dist(a, b):
        return float(np.linalg.norm(a - b))

    values = []

    px_top = dist(anchors_px[0], anchors_px[1])
    px_bottom = dist(anchors_px[3], anchors_px[2])
    px_left = dist(anchors_px[0], anchors_px[3])
    px_right = dist(anchors_px[1], anchors_px[2])

    mm_top = dist(anchors_mm[0], anchors_mm[1])
    mm_bottom = dist(anchors_mm[3], anchors_mm[2])
    mm_left = dist(anchors_mm[0], anchors_mm[3])
    mm_right = dist(anchors_mm[1], anchors_mm[2])

    if mm_top > 0:
        values.append(px_top / mm_top)

    if mm_bottom > 0:
        values.append(px_bottom / mm_bottom)

    if mm_left > 0:
        values.append(px_left / mm_left)

    if mm_right > 0:
        values.append(px_right / mm_right)

    if not values:
        return 10.0

    return float(np.mean(values))


def mm_to_px(H_inv, x_mm, y_mm):
    p = np.array([float(x_mm), float(y_mm), 1.0], dtype=np.float64)
    r = H_inv @ p
    r = r / r[2]

    return float(r[0]), float(r[1])


def dark_score_at(channels, H_inv, x_mm, y_mm, radius_mm, px_per_mm):
    """
    Lector mixto/adaptativo optimizado:
    - Usa canales ya preprocesados.
    - Lee B/N, escala de grises y fichas a color.
    - Ignora borde magenta/rosado.
    - Lee solo marca negra/gris dentro de la burbuja.
    """

    gray = channels["gray"]
    hue = channels["hue"]
    saturation = channels["saturation"]
    value = channels["value"]

    cx, cy = mm_to_px(H_inv, x_mm, y_mm)

    r = max(3, int(round(float(radius_mm) * float(px_per_mm))))

    x0 = max(0, int(round(cx)) - r)
    y0 = max(0, int(round(cy)) - r)
    x1 = min(gray.shape[1], int(round(cx)) + r + 1)
    y1 = min(gray.shape[0], int(round(cy)) + r + 1)

    roi_gray = gray[y0:y1, x0:x1]
    roi_hue = hue[y0:y1, x0:x1]
    roi_sat = saturation[y0:y1, x0:x1]
    roi_val = value[y0:y1, x0:x1]

    if roi_gray.size == 0:
        return 0.0

    inner_r = max(2, int(round(r * 0.72)))
    inside = get_circle_mask(roi_gray.shape, inner_r)

    if not np.any(inside):
        return 0.0

    g = roi_gray[inside]
    h = roi_hue[inside]
    s = roi_sat[inside]
    v = roi_val[inside]

    magenta_like = (
        (s > 45) &
        (v > 80) &
        (
            ((h >= 135) & (h <= 179)) |
            ((h >= 0) & (h <= 8))
        )
    )

    bn_dark = (
        (g < 135) &
        (~magenta_like)
    )

    color_black = (
        (g < 165) &
        (v < 185) &
        (s < 125) &
        (~magenta_like)
    )

    marked = bn_dark | color_black

    return float(np.mean(marked))


def pick_single(scores, marked_threshold=0.10, double_threshold=0.08):
    if not scores:
        return {
            "value": "BLANCO",
            "status": "blank",
            "scores": {}
        }

    ordered = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)

    best_value, best_score = ordered[0]
    second_score = ordered[1][1] if len(ordered) > 1 else 0.0

    if best_score < marked_threshold:
        return {
            "value": "BLANCO",
            "status": "blank",
            "scores": scores
        }

    if second_score >= double_threshold and (best_score - second_score) < 0.08:
        return {
            "value": "DOBLE",
            "status": "double",
            "scores": scores
        }

    return {
        "value": best_value,
        "status": "ok",
        "scores": scores
    }


def read_options(channels, H_inv, options, radius_mm, px_per_mm, threshold, double_threshold):
    scores = {}

    for op in options:
        scores[str(op["value"])] = dark_score_at(
            channels,
            H_inv,
            float(op["x"]),
            float(op["y"]),
            radius_mm,
            px_per_mm
        )

    return pick_single(scores, threshold, double_threshold)


def identification_score(data):
    score = 0

    dni = str(data.get("dni", {}).get("value", ""))
    nivel = data.get("nivel", {})
    grado = data.get("grado", {})
    tipo = data.get("tipo", {})

    if len(dni) == 8 and "?" not in dni and dni.upper() != "BLANCO":
        score += 40

    if nivel.get("status") == "ok" and str(nivel.get("value")) in ["1", "2", "3"]:
        score += 15

    if grado.get("status") == "ok" and str(grado.get("value")) in ["1", "2", "3", "4", "5", "6"]:
        score += 15

    if tipo.get("status") == "ok" and str(tipo.get("value")).upper() in ["A", "B", "C", "D"]:
        score += 15

    answers = data.get("answers", {})
    ok_answers = 0

    for info in answers.values():
        if isinstance(info, dict) and info.get("status") == "ok":
            ok_answers += 1

    score += min(ok_answers, 80) * 0.10

    return score


def process_single_orientation(img, img_path, template, rotation):
    H, H_inv, anchors_img, anchors_mm = detect_anchors(img, template)

    channels = prepare_channels(img)

    px_per_mm = estimate_px_per_mm_from_anchors(anchors_img, anchors_mm)

    bubble_r = float(template.get("bubble_r_mm", 1.25))
    sample_r_tpl = float(template.get("sample_r_mm", bubble_r * 0.50))

    sample_r = min(sample_r_tpl, bubble_r * 0.50)

    threshold = float(template.get("threshold_marked", 0.10))
    double_threshold = float(template.get("threshold_double", 0.08))

    dni_tpl = template["dni"]

    dni_digits = []
    dni_cols = []

    for c in range(int(dni_tpl["cols"])):
        opts = []

        for r, label in enumerate(dni_tpl["rows"]):
            opts.append({
                "value": str(label),
                "x": float(dni_tpl["x0"]) + c * float(dni_tpl["col_gap"]),
                "y": float(dni_tpl["y0"]) + r * float(dni_tpl["row_gap"]),
            })

        picked = read_options(
            channels,
            H_inv,
            opts,
            sample_r,
            px_per_mm,
            threshold,
            double_threshold
        )

        dni_cols.append(picked)
        dni_digits.append(picked["value"] if picked["status"] == "ok" else "?")

    dni_value = "".join(dni_digits)
    dni_status = "ok" if "?" not in dni_value and len(dni_value) == 8 else "observed"

    nivel = read_options(
        channels,
        H_inv,
        template["nivel"]["options"],
        sample_r,
        px_per_mm,
        threshold,
        double_threshold
    )

    grado = read_options(
        channels,
        H_inv,
        template["grado"]["options"],
        sample_r,
        px_per_mm,
        threshold,
        double_threshold
    )

    tipo = read_options(
        channels,
        H_inv,
        template["tipo"]["options"],
        sample_r,
        px_per_mm,
        threshold,
        double_threshold
    )

    answers = {}

    for q in template["questions"]:
        picked = read_options(
            channels,
            H_inv,
            q["options"],
            sample_r,
            px_per_mm,
            threshold,
            double_threshold
        )

        answers[str(q["n"])] = picked

    data = {
        "template": template["name"],
        "image": str(img_path),
        "rotation_used": rotation,
        "anchors_px": anchors_img.tolist(),
        "anchors_mm": anchors_mm.tolist(),
        "px_per_mm_est": px_per_mm,
        "sample_r_mm_used": sample_r,
        "reader_mode": "adaptive_color_and_bw_fast",
        "dni": {
            "value": dni_value,
            "status": dni_status,
            "columns": dni_cols
        },
        "nivel": nivel,
        "grado": grado,
        "tipo": tipo,
        "answers": answers
    }

    data["_score"] = identification_score(data)

    return data


def process(img_path, template_path):
    template = json.loads(Path(template_path).read_text(encoding="utf-8"))
    original_img = load_image(img_path)

    best_data = None
    errors = []

    # Primero intenta orientación normal.
    # Si la identificación sale bien, no pierde tiempo en 90, 180 y 270.
    for rotation in [0, 90, 180, 270]:
        try:
            img = rotate_image(original_img, rotation)
            data = process_single_orientation(img, img_path, template, rotation)

            if best_data is None or data["_score"] > best_data["_score"]:
                best_data = data

            # 40 DNI + 15 nivel + 15 grado + 15 tipo = 85.
            # Si llega a 85, la identificación está completa.
            if data["_score"] >= 85:
                break

        except Exception as exc:
            errors.append(f"{rotation}°: {str(exc)}")

    if best_data is None:
        raise RuntimeError(
            "No se pudo procesar la hoja en ninguna orientación. " +
            " | ".join(errors)
        )

    best_data.pop("_score", None)

    return best_data


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--template", required=True)

    args = parser.parse_args()

    try:
        data = process(args.image, args.template)

        print(json.dumps({
            "ok": True,
            "data": data
        }, ensure_ascii=False))

    except Exception as exc:
        print(json.dumps({
            "ok": False,
            "error": str(exc)
        }, ensure_ascii=False))

        sys.exit(1)


if __name__ == "__main__":
    main()