import cv2
import os
import json
import numpy as np
from preprocess import preprocess_chassis

CONFIG_PATH = "config.json"
_ocr = None


def load_config():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH) as f:
            return json.load(f)
    return {
        "confusion_map": {},
        "error_correction": {"max_errors": 2, "window_size": 3}
    }


def get_ocr():
    global _ocr
    if _ocr is None:
        from paddleocr import PaddleOCR
        _ocr = PaddleOCR(use_angle_cls=True, lang='en')
    return _ocr


def ocr_image(img):
    ocr = get_ocr()
    result = ocr.ocr(img, cls=True)
    if not result or not result[0]:
        return "", 0.0
    texts = [line[1][0] for line in result[0]]
    confs = [line[1][1] for line in result[0]]
    full_text = "".join(texts).upper()
    full_text = "".join(c for c in full_text if c.isalnum())
    avg_conf = sum(confs) / len(confs) if confs else 0.0
    return full_text, avg_conf


def read_chassis(image_path, save_comparison=False):
    variations = preprocess_chassis(image_path, save_comparison=save_comparison)
    best_text, best_conf, best_score = "", 0.0, -1
    for var in variations:
        text, conf = ocr_image(var)
        score = conf * max(len(text), 1)
        if score > best_score:
            best_text, best_conf, best_score = text, conf, score
    return best_text, best_conf


def can_substitute(got, want, confusion_map):
    return want in confusion_map.get(got, []) or got in confusion_map.get(want, [])


def apply_substitutions(ocr_text, expected_text, confusion_map, max_errors):
    if len(ocr_text) != len(expected_text):
        return ocr_text, False
    diffs = [(i, ocr_text[i], expected_text[i])
             for i in range(len(ocr_text)) if ocr_text[i] != expected_text[i]]
    if len(diffs) > max_errors:
        return ocr_text, False
    corrected = list(ocr_text)
    for i, got, want in diffs:
        if can_substitute(got, want, confusion_map):
            corrected[i] = want
        else:
            return ocr_text, False
    return "".join(corrected), True


def best_window_match(ocr_text, expected_text, window_size):
    exp_len = len(expected_text)
    best, best_diffs = None, exp_len + 1
    for start in range(max(0, len(ocr_text) - exp_len) + 1):
        candidate = ocr_text[start:start + exp_len]
        if len(candidate) != exp_len:
            continue
        diffs = sum(1 for a, b in zip(candidate, expected_text) if a != b)
        if diffs < best_diffs:
            best_diffs = diffs
            best = (candidate, diffs)
    if best and best[1] <= window_size:
        return best
    return None


def postprocess_with_hint(ocr_text, expected_text):
    config        = load_config()
    confusion_map = config.get("confusion_map", {})
    ec            = config.get("error_correction", {"max_errors": 2, "window_size": 3})
    max_errors    = ec["max_errors"]
    window_size   = ec["window_size"]

    if not ocr_text:
        return ocr_text, False
    if ocr_text == expected_text:
        return ocr_text, True
    if expected_text in ocr_text:
        return expected_text, True
    if len(ocr_text) == len(expected_text):
        corrected, fixed = apply_substitutions(
            ocr_text, expected_text, confusion_map, max_errors)
        if fixed:
            return corrected, True
    if abs(len(ocr_text) - len(expected_text)) <= window_size:
        match = best_window_match(ocr_text, expected_text, window_size)
        if match:
            candidate, diffs = match
            if diffs == 0:
                return candidate, True
            corrected, fixed = apply_substitutions(
                candidate, expected_text, confusion_map, max_errors)
            if fixed:
                return corrected, True
    if len(ocr_text) < len(expected_text):
        suffix = expected_text[-len(ocr_text):]
        corrected, fixed = apply_substitutions(
            ocr_text, suffix, confusion_map, max_errors=1)
        if fixed or ocr_text == suffix:
            return expected_text, True
    return ocr_text, False


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python ocr.py <chassis_image_path>")
    else:
        text, conf = read_chassis(sys.argv[1], save_comparison=True)
        print(f"Result     : {text}")
        print(f"Confidence : {conf:.2%}")