import os
import cv2
import json
import numpy as np
import base64
import torch
import torchvision
import asyncio
import time
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from ultralytics import YOLO
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.transforms import functional as F

app = FastAPI()

# 1. Initialization & Device Setup
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
models = {"yolo": None, "rcnn": None}
CLASSES = ["Background"] + [chr(i) for i in range(65, 91)]

# Load YOLO Model
yolo_path = "models/best.pt"
if os.path.exists(yolo_path):
    try:
        models["yolo"] = YOLO(yolo_path)
        print("✅ YOLO model loaded.")
    except Exception as e:
        print(f"❌ Failed to load YOLO: {e}")

# Load Faster R-CNN Model
rcnn_path = "models/best_faster_rcnn.pth"
if os.path.exists(rcnn_path):
    try:
        NUM_CLASSES = 27
        rcnn = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None)
        in_features = rcnn.roi_heads.box_predictor.cls_score.in_features
        rcnn.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)
        rcnn.load_state_dict(torch.load(rcnn_path, map_location=DEVICE))
        rcnn.to(DEVICE)
        rcnn.eval()
        models["rcnn"] = rcnn
        print("✅ Faster R-CNN model loaded.")
    except Exception as e:
        print(f"❌ Failed to load R-CNN: {e}")

@app.get("/")
async def get():
    with open("index.html", "r") as f:
        return HTMLResponse(f.read())

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    
    # State for word building
    state = {
        "last_pred": None,
        "pred_count": 0,
        "frames_to_confirm": 3
    }
    
    image_queue = asyncio.Queue(maxsize=1)

    async def receiver():
        """Aggressively receives frames and keeps only the newest one."""
        try:
            while True:
                data = await websocket.receive_text()
                # If queue is full, remove the old item to make room for the new one
                if image_queue.full():
                    try:
                        image_queue.get_nowait()
                    except asyncio.QueueEmpty:
                        pass
                await image_queue.put(data)
        except Exception:
            pass

    async def processor():
        """Pulls from queue as fast as the model can handle."""
        try:
            while True:
                data = await image_queue.get()
                payload = json.loads(data)
                header, encoded = payload["image"].split(",", 1)
                active_model = payload["model"]
                
                img_bytes = base64.b64decode(encoded)
                np_arr = np.frombuffer(img_bytes, np.uint8)
                img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
                
                detected_letter = ""
                best_label_idx = None
                best_label_name = ""
                
                current_model = models.get(active_model)
                if current_model is None:
                    # Fallback logic
                    fallback = "rcnn" if active_model == "yolo" else "yolo"
                    current_model = models.get(fallback)
                    active_model = fallback if current_model else None

                if active_model == "yolo" and current_model:
                    results = current_model(img, conf=0.6, verbose=False)[0]
                    highest_conf = 0.0
                    for box in results.boxes:
                        conf = box.conf.item()
                        cls_idx = int(box.cls.item())
                        name = current_model.names[cls_idx]
                        if conf > highest_conf:
                            highest_conf, best_label_idx, best_label_name = conf, cls_idx, name
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
                        cv2.putText(img, f"YOLO: {name} {conf:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)

                elif active_model == "rcnn" and current_model:
                    rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                    img_tensor = F.to_tensor(rgb_img).unsqueeze(0).to(DEVICE)
                    with torch.no_grad():
                        predictions = current_model(img_tensor)[0]
                    
                    highest_conf = 0.0
                    for i in range(len(predictions['boxes'])):
                        score = predictions['scores'][i].item()
                        if score > 0.75:
                            box = predictions['boxes'][i].cpu().numpy().astype(int)
                            label_idx = predictions['labels'][i].item()
                            name = CLASSES[label_idx]
                            if score > highest_conf:
                                highest_conf, best_label_idx, best_label_name = score, label_idx, name
                            cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
                            cv2.putText(img, f"RCNN: {name} {score:.2f}", (box[0], box[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

                # Word Builder Logic
                if best_label_idx is not None:
                    if best_label_idx == state["last_pred"]:
                        state["pred_count"] += 1
                    else:
                        state["last_pred"] = best_label_idx
                        state["pred_count"] = 1
                    if state["pred_count"] == state["frames_to_confirm"]:
                        detected_letter = best_label_name
                else:
                    state["last_pred"], state["pred_count"] = None, 0

                # Encode and Send
                _, buffer = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 70])
                out_encoded = base64.b64encode(buffer).decode('utf-8')
                await websocket.send_json({
                    "image": f"data:image/jpeg;base64,{out_encoded}",
                    "detected": detected_letter
                })
        except Exception:
            pass

    # Run receiver and processor concurrently
    await asyncio.gather(receiver(), processor())