import re

import cv2
import numpy as np
import torch

from PIL import Image

try:
    from paddleocr import PaddleOCR
except Exception as e:
    print(f"❌ PaddleOCR Import Error: {e}")
    PaddleOCR = None

try:
    from ultralytics import YOLO
except Exception as e:
    print(f"❌ YOLO Import Error: {e}")
    YOLO = None

try:
    from transformers import (
        AutoImageProcessor,
        AutoModelForImageClassification
    )
except Exception as e:
    print(f"❌ Transformers Import Error: {e}")
    AutoImageProcessor = None
    AutoModelForImageClassification = None

# ================= DEVICE CONFIG ================= #

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"🖥️  Using device: {DEVICE.upper()}")

# ================= STATES ================= #

STATE_CODES = {

    "TN": "Tamil Nadu",
    "KA": "Karnataka",
    "KL": "Kerala",
    "AP": "Andhra Pradesh",
    "TS": "Telangana",
    "MH": "Maharashtra",
    "DL": "Delhi",
    "GJ": "Gujarat",
    "RJ": "Rajasthan",
    "UP": "Uttar Pradesh",
    "WB": "West Bengal",
    "HR": "Haryana",
    "PB": "Punjab"
}

# ================= YOLO ================= #

yolo_model = None

if YOLO is not None:
    try:
        yolo_model = YOLO("license-plate-finetune-v1s.pt")
        
        # Move model to device and optimize
        yolo_model.to(DEVICE)
        yolo_model.overrides['conf'] = 0.5
        yolo_model.overrides['iou'] = 0.45
        yolo_model.overrides['max_det'] = 100
        
        print(f"✅ YOLO Loaded on {DEVICE.upper()}")

    except Exception as e:
        print("YOLO Error:", e)
        yolo_model = None
else:
    print("⚠️  YOLO/ultralytics not installed")

# ================= OCR ================= #

ocr = None

if PaddleOCR is not None:
    try:

        ocr = PaddleOCR(
            use_angle_cls=True,
            lang="en",
            show_log=False
        )

        print("✅ OCR Loaded")

    except Exception as e:

        print("OCR Error:", e)

        ocr = None
else:
    print("⚠️  PaddleOCR not installed")

# ================= VEHICLE MODEL ================= #

processor = None
vehicle_model = None

if AutoImageProcessor is not None and AutoModelForImageClassification is not None:
    try:

        processor = AutoImageProcessor.from_pretrained(
            "dima806/vehicle_10_types_image_detection"
        )

        vehicle_model = AutoModelForImageClassification.from_pretrained(
            "dima806/vehicle_10_types_image_detection"
        )

        vehicle_model.to(DEVICE)
        vehicle_model.eval()

        print(f"✅ Vehicle Model Loaded on {DEVICE.upper()}")

    except Exception as e:

        print("Vehicle Model Error:", e)

        processor = None
        vehicle_model = None
else:
    print("⚠️  Transformers not installed")

# ================= REGEX ================= #

plate_regex = re.compile(
    r"[A-Z]{2}[0-9]{1,2}[A-Z]{1,3}[0-9]{3,4}"
)

# ================= PREPROCESS ================= #

def preprocess_plate(crop):
    """Lightweight preprocessing - faster than full CLAHE"""
    try:
        # Skip heavy CLAHE, use simple resize + adaptive threshold
        gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
        resized = cv2.resize(gray, (320, 96))
        
        # Use adaptive thresholding instead of CLAHE (faster)
        enhanced = cv2.adaptiveThreshold(
            resized, 255, 
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 11, 2
        )
        
        # Light bilateral filter only
        filtered = cv2.bilateralFilter(enhanced, 5, 30, 30)
        
        return cv2.cvtColor(filtered, cv2.COLOR_GRAY2BGR)
    except Exception as e:
        print(f"Preprocess error: {e}")
        return crop

# ================= AUGMENT ================= #

def build_crops(crop):
    """Return only best crop variant instead of 3"""
    # Only return the original crop - no multiple variants
    # This reduces OCR calls from 3x to 1x
    return [crop]

# ================= OCR ================= #

def run_ocr(image):

    if ocr is None:
        return []

    return ocr.ocr(
        image,
        cls=True
    )

def parse_ocr(ocr_out):

    texts = []
    confs = []

    if not ocr_out:
        return texts, confs

    items = (
        ocr_out[0]
        if isinstance(ocr_out[0], list)
        else ocr_out
    )

    for item in items:

        try:

            txt, conf = item[1]

            texts.append(txt)

            confs.append(float(conf))

        except:
            continue

    return texts, confs

# ================= CLEAN ================= #

def clean_text(text):

    return re.sub(
        r"[^A-Z0-9]",
        "",
        text.upper()
    )

def fix_common(text):

    return (
        text.replace("O", "0")
        .replace("I", "1")
        .replace("B", "8")
        .replace("Z", "2")
        .replace("S", "5")
    )

# ================= VEHICLE CLASSIFY ================= #

def classify_vehicle(image_np):

    try:

        if processor is None:
            return "unknown", 0.0

        image_pil = Image.fromarray(image_np)

        inputs = processor(
            images=image_pil,
            return_tensors="pt"
        )

        inputs = {
            k: v.to(DEVICE)
            for k, v in inputs.items()
        }

        with torch.no_grad():

            outputs = vehicle_model(**inputs)

        logits = outputs.logits

        probs = torch.nn.functional.softmax(
            logits,
            dim=-1
        )

        pred = probs.argmax(-1).item()

        confidence = float(
            probs.max().item()
        )

        label = vehicle_model.config.id2label[pred]

        return label, confidence

    except Exception as e:

        print("Classification Error:", e)

        return "unknown", 0.0

# ================= STATE ================= #

def extract_state(plate):

    if len(plate) < 2:
        return "UNKNOWN"

    code = plate[:2]

    return code if code in STATE_CODES else "UNKNOWN"

# ================= DETECT ================= #

def detect_plate(image):

    if isinstance(image, Image.Image):
        image = np.array(image.convert("RGB"))

    # ENABLE vehicle classification
    vehicle_type, vehicle_conf = classify_vehicle(image)

    if yolo_model is None:

        return (
            "",
            "UNKNOWN",
            vehicle_type,
            vehicle_conf,
            False
        )

    # YOLO detection with optimizations for speed
    results = yolo_model(
        image,
        conf=0.5,           # Confidence threshold
        iou=0.45,           # IoU threshold  
        imgsz=384,          # Image size
        verbose=False,      # No logging
        device=0 if DEVICE == 'cuda' else 'cpu'  # Use GPU if available
    )
    boxes = results[0].boxes

    if boxes is None or len(boxes) == 0:

        return (
            "",
            "UNKNOWN",
            vehicle_type,
            vehicle_conf,
            False
        )

    h, w = image.shape[:2]
    xyxy = boxes.xyxy.cpu().numpy()
    confs = boxes.conf.cpu().numpy()

    best_plate = ""
    best_confidence = 0.0

    for i, (x1, y1, x2, y2) in enumerate(xyxy):

        if confs[i] < 0.5:
            continue

        pad = int(0.12 * max(x2 - x1, y2 - y1))
        l = max(int(x1 - pad), 0)
        t = max(int(y1 - pad), 0)
        r = min(int(x2 + pad), w - 1)
        b = min(int(y2 + pad), h - 1)

        crop = image[t:b, l:r]
        
        # Only ONE preprocessing + OCR per detection (no variants)
        pre = preprocess_plate(crop)
        ocr_out = run_ocr(pre)
        texts, confs_ocr = parse_ocr(ocr_out)

        # Find best text in this detection
        for txt, cf in zip(texts, confs_ocr):

            if cf < 0.3:
                continue

            norm = fix_common(clean_text(txt))

            if len(norm) < 4:
                continue

            # Check if matches Indian plate regex
            match = plate_regex.search(norm)
            if match:
                plate = match.group(0)
                # Early exit on first good match
                if cf > best_confidence:
                    best_plate = plate
                    best_confidence = cf

    plate = best_plate.upper() if best_plate else ""
    state = extract_state(plate) if plate else "UNKNOWN"

    return (
        plate,
        state,
        vehicle_type,
        vehicle_conf,
        len(plate) > 0
    )