import fastapi
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from typing import Dict
import cv2
import numpy as np
from main import NumberPlateOCR

app = FastAPI()

# Initialize the OCR pipeline
pipeline = NumberPlateOCR(use_yolo=True, verbose=False)

@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.post("/ocr-only/")
async def ocr_only(file: UploadFile = File(...)) -> Dict:
    """
    Perform OCR on an uploaded image without YOLO detection.

    Args:
        file: Uploaded image file.

    Returns:
        JSON response with OCR results.
    """
    try:
        # Read image from upload
        contents = await file.read()
        np_img = np.frombuffer(contents, np.uint8)
        image = cv2.imdecode(np_img, cv2.IMREAD_COLOR)

        if image is None:
            raise HTTPException(status_code=400, detail="Invalid image file")

        # Process the image directly (skip YOLO)
        results = pipeline.process_from_plate_image(image, show_visualization=False)

        return {
            "lines": results["lines"],
            "multiline_text": results["multiline_text"],
            "singleline_text": results["singleline_text"],
            "num_lines": results["num_lines"],
            "total_chars": results["total_chars"],
            "confidence_stats": results["confidence_stats"]
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



@app.post("/yolo-ocr/")
async def yolo_ocr(file: UploadFile = File(...)) -> Dict:
    """
    Perform YOLO-based plate detection and OCR on an uploaded image.

    Args:
        file: Uploaded image file.

    Returns:
        JSON response with OCR results.
    """
    try:
        # Read image from upload
        contents = await file.read()
        np_img = np.frombuffer(contents, np.uint8)
        image_path = "uploaded_image.jpg"

        # Save the uploaded image temporarily
        with open(image_path, "wb") as f:
            f.write(contents)

        # Process the image using the pipeline
        results = pipeline.process_image(image_path, save_contours=False, show_visualization=False)

        # Format the response
        response = {
            "image_path": results["image_path"],
            "num_plates": results["num_plates"],
            "plates": []
        }

        for plate in results["plates"]:
            response["plates"].append({
                "lines": plate["lines"],
                "multiline_text": plate["multiline_text"],
                "singleline_text": plate["singleline_text"],
                "num_lines": plate["num_lines"],
                "total_chars": plate["total_chars"],
                "confidence_stats": plate["confidence_stats"]
            })

        return response

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))