"""
OCR: per-char ONNX (case-aware) + fast ddddocr fallback.
Never loads our charsets.json into ddddocr (that caused Chinese ERR).
"""
from __future__ import annotations

import io
import json
import re
import time
from pathlib import Path

import numpy as np
from PIL import Image, ImageOps, ImageFilter

ROOT = Path(__file__).resolve().parent
MODELS = ROOT / "models"
DATASET_LABELED = ROOT / "dataset" / "labeled"
DATASET_RAW = ROOT / "dataset" / "raw"
LEARN_LOG = ROOT / "dataset" / "learn_log.jsonl"

# Don't use custom ONNX until synth val sequence accuracy is usable (anti-overfit)
MIN_ONNX_VAL_ACC = 0.45
MIN_CHAR_ACC = 0.55

_ocr_dddd = None
_char_sess = None
_char_meta = None
_onnx_sess = None
_onnx_meta = None
_mode = "none"
_loaded = False
_gemini_ok = False


def engine_mode() -> str:
    return _mode


def gemini_enabled() -> bool:
    try:
        import gemini_solver as gs

        return gs.available()
    except Exception:
        return False


def _safe_dddd():
    import ddddocr

    return ddddocr.DdddOcr(show_ad=False)


def load_engines(force: bool = False):
    global _ocr_dddd, _char_sess, _char_meta, _onnx_sess, _onnx_meta, _mode, _loaded, _gemini_ok
    if _loaded and not force and (_ocr_dddd is not None or _char_sess is not None or _onnx_sess is not None or _gemini_ok):
        return _mode

    _ocr_dddd = None
    _char_sess = None
    _char_meta = None
    _onnx_sess = None
    _onnx_meta = None
    _mode = "none"
    _gemini_ok = False

    try:
        import gemini_solver as gs

        _gemini_ok = gs.available()
        if _gemini_ok:
            _mode = "gemini"
    except Exception as e:
        print("gemini check:", e)

    try:
        _ocr_dddd = _safe_dddd()
        if _mode == "none":
            _mode = "builtin_dddd"
    except Exception as e:
        print("ddddocr load failed:", e)

    # Prefer per-character model (best for capital/small + 0/9/o)
    char_onnx = MODELS / "char62.onnx"
    char_meta_path = MODELS / "char62.json"
    if char_onnx.is_file() and char_meta_path.is_file():
        try:
            meta = json.loads(char_meta_path.read_text(encoding="utf-8"))
            if float(meta.get("val_acc") or 0) >= MIN_CHAR_ACC and meta.get("engine") == "portal_char62":
                import onnxruntime as ort

                _char_sess = ort.InferenceSession(str(char_onnx), providers=["CPUExecutionProvider"])
                _char_meta = meta
                if not _gemini_ok:
                    _mode = "portal_char62"
        except Exception as e:
            print("char62 load skipped:", e)
            _char_sess = None

    # Full captcha multihead (optional)
    onnx = MODELS / "captcha.onnx"
    meta_path = MODELS / "charsets.json"
    if _char_sess is None and onnx.is_file() and meta_path.is_file():
        try:
            meta = json.loads(meta_path.read_text(encoding="utf-8"))
        except Exception:
            meta = {}
        eng = meta.get("engine")
        val_acc = float(meta.get("val_acc") or 0)
        if not isinstance(meta.get("word"), list) and eng in ("portal_multihead", "portal_crnn_ctc") and val_acc >= MIN_ONNX_VAL_ACC:
            try:
                import onnxruntime as ort

                _onnx_sess = ort.InferenceSession(str(onnx), providers=["CPUExecutionProvider"])
                _onnx_meta = meta
                if not _gemini_ok:
                    _mode = eng
            except Exception as e:
                print("ONNX load skipped:", e)

    if _gemini_ok:
        _mode = "gemini"

    _loaded = True
    return _mode


def clean_text(text: str, only_alnum: bool = True, max_len: int = 8) -> str:
    t = (text or "").strip()
    if only_alnum:
        t = re.sub(r"[^A-Za-z0-9]", "", t)
    if max_len > 0:
        t = t[:max_len]
    return t


def crop_captcha_region(img: Image.Image) -> Image.Image:
    rgb = np.asarray(img.convert("RGB"), dtype=np.int16)
    h, w = rgb.shape[:2]
    if w < 40:
        return img
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    blue = (b > r + 25) & (b > g + 15) & (b > 90)
    col = blue.mean(axis=0)
    cut = w
    for x in range(w - 1, w // 2, -1):
        if col[x] >= 0.25:
            cut = x
        elif cut < w and col[x] < 0.05:
            break
    if cut < w - 8:
        cut = max(w // 2, cut - 2)
        return Image.fromarray(rgb[:, :cut, :].astype(np.uint8), mode="RGB")
    return img


def _green_mask(rgb: np.ndarray) -> np.ndarray:
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    return (g > r + 8) & (g > b + 8) & (g > 70)


def _png_bytes(im: Image.Image) -> bytes:
    buf = io.BytesIO()
    im.save(buf, format="PNG")
    return buf.getvalue()


def _fast_variants(data: bytes) -> list[bytes]:
    img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))
    rgb = np.asarray(img, dtype=np.int16)
    green = _green_mask(rgb)
    ink = np.where(green, 0, 255).astype(np.uint8)
    im = ImageOps.autocontrast(Image.fromarray(ink, mode="L"))
    mid = im.resize((im.width * 2, im.height * 2), Image.Resampling.NEAREST)
    mid = mid.filter(ImageFilter.SHARPEN)
    return [_png_bytes(img.convert("RGB")), _png_bytes(mid)]


def _segment_boxes(data: bytes, n: int = 6) -> list[tuple[int, int, int, int]]:
    """Equal-width boxes (matches char62 training crops) after refresh crop."""
    img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))
    w, h = img.size
    return [(int(i * w / n), 0, int((i + 1) * w / n), h) for i in range(n)]


def _crop_to_char_tensor(data: bytes, box, height: int, width: int) -> np.ndarray:
    img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))
    x0, y0, x1, y1 = box
    crop = img.crop((x0, y0, x1, y1))
    rgb = np.asarray(crop, dtype=np.int16)
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    green = (g > r + 8) & (g > b + 8) & (g > 70)
    gray = (0.299 * r + 0.587 * g + 0.114 * b).astype(np.float32)
    ink = np.where(green, 255.0 - np.clip(g.astype(np.float32) * 1.1, 0, 255), gray)
    im = ImageOps.autocontrast(Image.fromarray(ink.astype(np.uint8), mode="L"))
    im = im.resize((width, height), Image.Resampling.BILINEAR)
    return (np.asarray(im, dtype=np.float32) / 255.0)[None, None, :, :]


def _read_char62(data: bytes) -> str:
    assert _char_sess is not None and _char_meta is not None
    h = int(_char_meta.get("height", 48))
    w = int(_char_meta.get("width", 32))
    charset = list(_char_meta.get("charset") or [])
    boxes = _segment_boxes(data, 6)
    out = []
    for box in boxes:
        inp = _crop_to_char_tensor(data, box, h, w)
        logits = _char_sess.run(None, {"image": inp})[0]
        idx = int(np.argmax(logits[0]))
        out.append(charset[idx] if 0 <= idx < len(charset) else "?")
    return "".join(out)


def preprocess_sarathi(data: bytes, height: int, width: int) -> np.ndarray:
    img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))
    rgb = np.asarray(img, dtype=np.int16)
    arr_r, arr_g, arr_b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    green = _green_mask(rgb)
    gray = (0.299 * arr_r + 0.587 * arr_g + 0.114 * arr_b).astype(np.float32)
    ink = np.where(green, 255.0 - np.clip(arr_g.astype(np.float32) * 1.1, 0, 255), gray)
    im = ImageOps.autocontrast(Image.fromarray(ink.astype(np.uint8), mode="L"))
    im = im.resize((width, height), Image.Resampling.BILINEAR)
    return (np.asarray(im, dtype=np.float32) / 255.0)[None, None, :, :]


def _decode_multihead(logits: np.ndarray, charset: list) -> str:
    if logits.ndim == 3:
        logits = logits[0]
    idxs = logits.argmax(axis=-1).tolist()
    return "".join(charset[i] if 0 <= i < len(charset) else "" for i in idxs)


def _softmax(x: np.ndarray) -> np.ndarray:
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)


def _read_multihead_tta(data: bytes) -> str:
    """Multiple slight preprocesses + per-position vote (reduces overfit misses)."""
    assert _onnx_sess is not None and _onnx_meta is not None
    h = int(_onnx_meta.get("height", 64))
    w = int(_onnx_meta.get("width", 192))
    charset = list(_onnx_meta.get("charset") or [])
    img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))

    variants = [img]
    # slight shifts
    for dx, dy in ((-2, 0), (2, 0), (0, -1), (0, 1), (-1, -1), (1, 1)):
        shifted = img.transform(
            img.size, Image.AFFINE, (1, 0, dx, 0, 1, dy), fillcolor=(220, 225, 220)
        )
        variants.append(shifted)
    # contrast
    from PIL import ImageEnhance

    variants.append(ImageEnhance.Contrast(img).enhance(1.25))
    variants.append(ImageEnhance.Contrast(img).enhance(0.85))

    # accumulate softmax votes per slot
    votes = [np.zeros(len(charset), dtype=np.float64) for _ in range(6)]
    for im in variants:
        buf = io.BytesIO()
        im.save(buf, format="PNG")
        # reuse preprocess on bytes
        arr = preprocess_sarathi(buf.getvalue(), h, w)
        logits = _onnx_sess.run(None, {"image": arr})[0]
        if logits.ndim == 3:
            logits = logits[0]
        probs = _softmax(logits.astype(np.float64))
        for i in range(min(6, probs.shape[0])):
            votes[i] += probs[i]

    return "".join(charset[int(v.argmax())] for v in votes)


def read_image_bytes(data: bytes) -> str:
    global _mode
    if not _loaded or (
        _ocr_dddd is None and _char_sess is None and _onnx_sess is None and not _gemini_ok
    ):
        load_engines()

    # 0) Gemini first — best for capital/small (friend-style ~90%+)
    try:
        import gemini_solver as gs

        if gs.available():
            # send cropped captcha only (no refresh button)
            img = crop_captcha_region(Image.open(io.BytesIO(data)).convert("RGB"))
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            text, model = gs.solve_bytes(buf.getvalue())
            t = clean_text(text, max_len=8)
            if t:
                _mode = f"gemini:{model}"
                return t
    except Exception as e:
        print("gemini skip:", e)

    # 1) Per-char model — designed for case + digits
    if _char_sess is not None:
        try:
            t = clean_text(_read_char62(data), max_len=8)
            if len(t) == 6:
                return t
            if t:
                return t
        except Exception as e:
            print("char62 error:", e)

    # 2) Full captcha ONNX (single pass — TTA was hurting some samples)
    if _onnx_sess is not None and _onnx_meta is not None:
        try:
            h = int(_onnx_meta.get("height", 64))
            w = int(_onnx_meta.get("width", 192))
            charset = list(_onnx_meta.get("charset") or [])
            out = _onnx_sess.run(None, {"image": preprocess_sarathi(data, h, w)})[0]
            t = clean_text(_decode_multihead(out, charset), max_len=8)
            if t:
                return t
        except Exception as e:
            print("ONNX infer error:", e)

    # 3) Fast ddddocr (weak on case / 0-o)
    if _ocr_dddd is not None:
        best = ""
        for variant in _fast_variants(data):
            try:
                t = clean_text(str(_ocr_dddd.classification(variant) or ""), max_len=8)
            except Exception:
                continue
            if len(t) == 6:
                return t
            if len(t) > len(best):
                best = t
        return best
    return ""


def save_labeled_bytes(data: bytes, label: str, ext: str = ".png") -> Path:
    DATASET_LABELED.mkdir(parents=True, exist_ok=True)
    lab = clean_text(label, max_len=12)
    if len(lab) < 3:
        raise ValueError("label too short")
    path = DATASET_LABELED / f"{lab}_{int(time.time() * 1000)}{ext}"
    path.write_bytes(data)
    return path


def learn_log(event: dict) -> None:
    LEARN_LOG.parent.mkdir(parents=True, exist_ok=True)
    event = dict(event)
    event["ts"] = int(time.time())
    with LEARN_LOG.open("a", encoding="utf-8") as f:
        f.write(json.dumps(event, ensure_ascii=False) + "\n")


def labeled_count() -> int:
    if not DATASET_LABELED.is_dir():
        return 0
    return sum(
        1
        for p in DATASET_LABELED.glob("*.*")
        if p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
    )