import gradio as gr
from llama_cpp import Llama
from llama_cpp.llama_chat_format import Llava15ChatHandler
from huggingface_hub import hf_hub_download
import base64
import json
import os
import io
from PIL import Image

# ---------------------------------------------------------
# IMPORTANT: REPLACE THESE WITH YOUR HUGGING FACE DETAILS
# ---------------------------------------------------------
HF_USERNAME = "bhargavcli"
REPO_NAME = "anpr-vision-api"
# For example: "Bhargav123/moondream2-anpr"
REPO_ID = f"{HF_USERNAME}/{REPO_NAME}"
# ---------------------------------------------------------

print("Downloading model files from your repository...")
try:
    mmproj_path = hf_hub_download(repo_id=REPO_ID, filename="moondream2-mmproj-f16.gguf")
    model_path = hf_hub_download(repo_id=REPO_ID, filename="moondream2-text-model-f16.gguf")
    print("Download complete!")
except Exception as e:
    print(f"Error downloading models. Make sure you entered your REPO_ID correctly! Error: {e}")
    # Fallback to public repo for testing if it fails
    mmproj_path = hf_hub_download(repo_id="moondream/moondream2-gguf", filename="moondream2-mmproj-f16.gguf")
    model_path = hf_hub_download(repo_id="moondream/moondream2-gguf", filename="moondream2-text-model-f16.gguf")

print("Loading Multimodal Projector and Text Model...")
chat_handler = Llava15ChatHandler(clip_model_path=mmproj_path)
llm = Llama(
    model_path=model_path,
    chat_handler=chat_handler,
    n_ctx=2048, # Context window
    logits_all=True
)
print("Model loaded successfully!")

SYSTEM_PROMPT = """You are an expert Automated Number Plate Recognition (ANPR) vision-language engine optimized for real-world traffic scenarios. Your task is to process the vehicle image provided in the input, apply specialized structural preprocessing filters mentally, and extract the license plate text with absolute structural fidelity.

Image Processing Instructions:
1. Isolate the high-probability license plate region (typically a rectangular plate on the lower bumper or grille).
2. Compensate for environmental distortions including sun glare, shadows, and motion blur by sharpening character boundaries.
3. If the plate uses a stacked/multi-line layout, perform spatial vertical sorting (read line 1 completely from left-to-right, then line 2 completely from left-to-right).
4. Parse the characters using standard Indian Regional Transport Office (RTO) syntax: [State Code (2 Letters)] [District Code (2 Digits)] [Series (1-2 Letters)] [Unique ID (4 Digits)].
5. Discard screw marks, plate borders, and dirt spots that mimic alphanumeric noise (e.g., mistaking a screw for a dot or '1').

Output Format:
Return your response strictly in the following JSON format. Do not include any introductory conversational text or markdown code blocks around the JSON payload.

{
  "detection_status": "Success" or "Failed",
  "layout_type": "Single-line" or "Multi-line",
  "raw_extracted_text": "Exactly what is printed including spacing",
  "cleaned_plate_number": "Sanitized alphanumeric sequence without spaces or symbols",
  "confidence_assessment": "High", "Medium", or "Low"
}"""

def extract_plate(image):
    if image is None:
        return json.dumps({"detection_status": "Failed", "error": "No image provided"})
        
    # Convert Gradio Image (numpy array or PIL Image) to base64
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
        
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    data_uri = f"data:image/jpeg;base64,{img_str}"

    print("Running inference...")
    
    # Send to Moondream2 via llama_cpp
    response = llm.create_chat_completion(
        messages=[
            {
                "role": "system",
                "content": SYSTEM_PROMPT
            },
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": "Extract the license plate JSON."}
                ]
            }
        ],
        max_tokens=150,
        temperature=0.1 # Low temperature for accurate OCR
    )
    
    output = response["choices"][0]["message"]["content"]
    
    # Try to clean it up in case the LLM hallucinates markdown blocks
    output = output.replace("```json", "").replace("```", "").strip()
    return output

# Create a Gradio interface
# Create a Gradio interface
demo = gr.Interface(
    fn=extract_plate,
    inputs=gr.Image(type="pil", label="Vehicle Image"),
    outputs=gr.JSON(label="Extracted License Plate Data"),
    title="ANPR Vision-Language API",
    description="Upload a picture of a vehicle to extract its license plate using Moondream2."
)


# Launch the app
demo.launch()