import os
os.environ['OMP_NUM_THREADS'] = '1'

import re
from datetime import datetime
import asyncio

import cv2
import gradio as gr
import numpy as np
import torch
from PIL import Image

from paddleocr import PaddleOCR
from ultralytics import YOLO
from transformers import AutoImageProcessor, AutoModelForImageClassification
from sqlalchemy import create_engine, text
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

# Fix async warning
asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())

# Load environment variables
load_dotenv()

DATABASE_URL = os.getenv("DATABASE_URL")
engine = create_engine(DATABASE_URL)

HF_TOKEN = os.getenv("HF_TOKEN")
client = InferenceClient(
    model="defog/sqlcoder-7b-2",
    token=HF_TOKEN
)

# ---------------- LOAD MODELS ---------------- #

# YOLO (plate detection)
try:
    yolo_model = YOLO("license-plate-finetune-v1s.pt")
except Exception as e:
    print(f"Warning: Failed to load YOLO model: {e}")
    yolo_model = None

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

# Vehicle classification
try:
    processor = AutoImageProcessor.from_pretrained("dima806/vehicle_10_types_image_detection")
    vehicle_model = AutoModelForImageClassification.from_pretrained("dima806/vehicle_10_types_image_detection")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    vehicle_model.to(device)
    vehicle_model.eval()
except Exception as e:
    print(f"Warning: Failed to load vehicle classification model: {e}")
    processor = None
    vehicle_model = None
    device = "cpu"

# Regex
plate_regex = re.compile(r"[A-Z]{2}\d{1,2}[A-Z]{1,3}\d{3,4}")

# Database initialization
def init_db():
    with engine.connect() as conn:
        conn.execute(text("""
        CREATE TABLE IF NOT EXISTS vehicle_logs (
            id SERIAL PRIMARY KEY,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            plate TEXT,
            vehicle_type TEXT,
            vehicle_conf FLOAT,
            date TEXT,
            time TEXT
        )
        """))
        conn.commit()

init_db()

# ---------------- PREPROCESS ---------------- #

def preprocess_plate(crop):
    gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
    resized = cv2.resize(gray, (320, 96))

    clahe = cv2.createCLAHE(2.0, (8, 8))
    enhanced = clahe.apply(resized)

    filtered = cv2.bilateralFilter(enhanced, 7, 50, 50)

    kernel = np.array([[0,-1,0],[-1,5,-1],[0,-1,0]])
    sharp = cv2.filter2D(filtered, -1, kernel)

    return cv2.cvtColor(sharp, cv2.COLOR_GRAY2BGR)

# ---------------- AUGMENT ---------------- #

def build_crops(crop):
    return [
        crop,
        cv2.resize(crop, None, fx=1.2, fy=1.2),
        cv2.GaussianBlur(crop, (3,3), 0)
    ]

# ---------------- OCR ---------------- #

def run_ocr(image):
    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:
            text, conf = item[1]
            texts.append(text)
            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 CLASSIFICATION ---------------- #

def classify_vehicle(image_np):
    try:
        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("Vehicle classification error:", e)
        return "unknown", 0.0

# ---------------- MAIN PIPELINE ---------------- #

def detect(image):
    now = datetime.now()
    date = now.strftime("%Y-%m-%d")
    time = now.strftime("%H:%M:%S")

    # 🔹 Vehicle classification
    vehicle_type, vehicle_conf = classify_vehicle(image)

    # 🔹 Plate detection
    results = yolo_model(image)
    boxes = results[0].boxes

    if boxes is None or len(boxes) == 0:
        return f"{date} {time} | {vehicle_type} |", {
            "date": date,
            "time": time,
            "vehicle_type": vehicle_type,
            "vehicle_confidence": round(vehicle_conf, 3),
            "plate": ""
        }

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

    collected = []

    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]

        for variant in build_crops(crop):
            pre = preprocess_plate(variant)

            ocr_out = run_ocr(pre)
            texts, confs_ocr = parse_ocr(ocr_out)

            for txt, cf in zip(texts, confs_ocr):

                if cf < 0.3:
                    continue

                norm = fix_common(clean_text(txt))

                if len(norm) < 4:
                    continue

                collected.append(norm)

    if not collected:
        plate = ""
    else:
        combined = "".join(collected)
        match = plate_regex.search(combined)
        plate = match.group(0) if match else combined

    # SAVE TO DATABASE
    try:
        with engine.connect() as conn:
            conn.execute(text("""
                INSERT INTO vehicle_logs
                (plate, vehicle_type, vehicle_conf, date, time)
                VALUES (:plate, :vehicle_type, :vehicle_conf, :date, :time)
            """), {
                "plate": plate,
                "vehicle_type": vehicle_type,
                "vehicle_conf": float(vehicle_conf),
                "date": date,
                "time": time
            })
            conn.commit()
    except Exception as e:
        print("Database insert error:", e)

    # 🔹 Final Output
    return f"{date} {time} | {vehicle_type} | {plate}", {
        "date": date,
        "time": time,
        "vehicle_type": vehicle_type,
        "vehicle_confidence": round(vehicle_conf, 3),
        "plate": plate
    }

# ---------------- LLM SQL LAYER (SQLCoder) ---------------- #

def ask_llm(user_query):
    schema = """
    Table: vehicle_logs
    
    Columns:
    - id (SERIAL PRIMARY KEY)
    - timestamp (TIMESTAMP)
    - plate (TEXT) - License plate number
    - vehicle_type (TEXT) - Type of vehicle
    - vehicle_conf (FLOAT) - Detection confidence
    - date (TEXT) - Date in YYYY-MM-DD format
    - time (TEXT) - Time in HH:MM:SS format
    """

    prompt = f"""### Task
Generate PostgreSQL SQL query for the following question.

### Rules
- Only SELECT queries allowed
- Use vehicle_logs table
- No markdown, no explanation
- Output SQL only
- Use appropriate WHERE clauses for filtering
- Use COUNT(*), SUM(), AVG() for aggregations if needed
- Order by timestamp DESC for chronological queries

### Schema
{schema}

### User Question
{user_query}

### SQL Query
"""

    try:
        response = client.text_generation(
            prompt,
            max_new_tokens=150,
            temperature=0.1
        )

        sql_query = response.strip()
        
        # Clean up markdown formatting if present
        sql_query = sql_query.replace("```sql", "")
        sql_query = sql_query.replace("```", "")
        sql_query = sql_query.strip()

        return sql_query
    
    except Exception as e:
        print(f"LLM error: {e}")
        return f"SELECT * FROM vehicle_logs LIMIT 10; -- Error: {e}"

def run_query(user_query):
    sql_query = ask_llm(user_query)

    try:
        with engine.connect() as conn:
            result = conn.execute(text(sql_query))
            rows = [dict(row._mapping) for row in result]

        return {
            "query": user_query,
            "sql": sql_query,
            "result": rows
        }

    except Exception as e:
        return {
            "error": str(e),
            "sql": sql_query
        }

# ---------------- UI ---------------- #

with gr.Blocks() as demo:
    gr.Markdown("# 🚗 Vehicle Intelligence System")

    with gr.Tab("Detection"):
        img = gr.Image(type="numpy")
        out1 = gr.Textbox(label="Result")
        out2 = gr.JSON(label="Structured Output")
        btn = gr.Button("Detect")

        btn.click(
            fn=detect,
            inputs=img,
            outputs=[out1, out2]
        )

    with gr.Tab("Ask Database"):
        query_input = gr.Textbox(
            label="Ask Anything",
            placeholder="How many cars today?"
        )
        query_output = gr.JSON()
        ask_btn = gr.Button("Ask")

        ask_btn.click(
            fn=run_query,
            inputs=query_input,
            outputs=query_output
        )

demo.launch(server_name="0.0.0.0", server_port=7860)