from typing import Any, Dict
import os
import json
import time
import queue
import torch
import cv2
import base64
import logging
import threading
import requests
import asyncio
import psutil
import zipfile
from urllib.parse import urlparse
from fastapi import FastAPI, WebSocket
from fastapi.websockets import WebSocketDisconnect
from websockets.exceptions import ConnectionClosedError
from ultralytics import YOLO

# Import shared pipeline functions
from siwatsystem.pympta import load_pipeline_from_zip, run_pipeline

app = FastAPI()

# Global dictionaries to keep track of models and streams
# "models" now holds a nested dict: { camera_id: { modelId: model_tree } }
models: Dict[str, Dict[str, Any]] = {}
streams: Dict[str, Dict[str, Any]] = {}

with open("config.json", "r") as f:
    config = json.load(f)

poll_interval = config.get("poll_interval_ms", 100)
reconnect_interval = config.get("reconnect_interval_sec", 5)
TARGET_FPS = config.get("target_fps", 10)
poll_interval = 1000 / TARGET_FPS
logging.info(f"Poll interval: {poll_interval}ms")
max_streams = config.get("max_streams", 5)
max_retries = config.get("max_retries", 3)

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler()
    ]
)

# Ensure the models directory exists
os.makedirs("models", exist_ok=True)

# Constants for heartbeat and timeouts
HEARTBEAT_INTERVAL = 2  # seconds
WORKER_TIMEOUT_MS = 10000

# Locks for thread-safe operations
streams_lock = threading.Lock()
models_lock = threading.Lock()

# Add helper to download mpta ZIP file from a remote URL
def download_mpta(url: str, dest_path: str) -> str:
    try:
        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
        response = requests.get(url, stream=True)
        if response.status_code == 200:
            with open(dest_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            logging.info(f"Downloaded mpta file from {url} to {dest_path}")
            return dest_path
        else:
            logging.error(f"Failed to download mpta file (status code {response.status_code})")
            return None
    except Exception as e:
        logging.error(f"Exception downloading mpta file from {url}: {e}")
        return None

####################################################
# Detection and frame processing functions
####################################################
@app.websocket("/")
async def detect(websocket: WebSocket):
    logging.info("WebSocket connection accepted")
    persistent_data_dict = {}

    async def handle_detection(camera_id, stream, frame, websocket, model_tree, persistent_data):
        try:
            detection_result = run_pipeline(frame, model_tree)
            detection_data = {
                "type": "imageDetection",
                "cameraIdentifier": camera_id,
                "timestamp": time.time(),
                "data": {
                    "detection": detection_result if detection_result else None,
                    "modelId": stream["modelId"],
                    "modelName": stream["modelName"]
                }
            }
            logging.debug(f"Sending detection data for camera {camera_id}: {detection_data}")
            await websocket.send_json(detection_data)
            return persistent_data
        except Exception as e:
            logging.error(f"Error in handle_detection for camera {camera_id}: {e}")
            return persistent_data

    def frame_reader(camera_id, cap, buffer, stop_event):
        retries = 0
        try:
            while not stop_event.is_set():
                try:
                    ret, frame = cap.read()
                    if not ret:
                        logging.warning(f"Connection lost for camera: {camera_id}, retry {retries+1}/{max_retries}")
                        cap.release()
                        time.sleep(reconnect_interval)
                        retries += 1
                        if retries > max_retries and max_retries != -1:
                            logging.error(f"Max retries reached for camera: {camera_id}")
                            break
                        # Re-open
                        cap = cv2.VideoCapture(streams[camera_id]["rtsp_url"])
                        if not cap.isOpened():
                            logging.error(f"Failed to reopen RTSP stream for camera: {camera_id}")
                            continue
                        continue
                    retries = 0
                    # Overwrite old frame if buffer is full
                    if not buffer.empty():
                        try:
                            buffer.get_nowait()
                        except queue.Empty:
                            pass
                    buffer.put(frame)
                except cv2.error as e:
                    logging.error(f"OpenCV error for camera {camera_id}: {e}")
                    cap.release()
                    time.sleep(reconnect_interval)
                    retries += 1
                    if retries > max_retries and max_retries != -1:
                        logging.error(f"Max retries reached after OpenCV error for camera {camera_id}")
                        break
                    cap = cv2.VideoCapture(streams[camera_id]["rtsp_url"])
                    if not cap.isOpened():
                        logging.error(f"Failed to reopen RTSP stream for camera {camera_id} after OpenCV error")
                        continue
                except Exception as e:
                    logging.error(f"Unexpected error for camera {camera_id}: {e}")
                    cap.release()
                    break
        except Exception as e:
            logging.error(f"Error in frame_reader thread for camera {camera_id}: {e}")

    async def process_streams():
        logging.info("Started processing streams")
        try:
            while True:
                start_time = time.time()
                with streams_lock:
                    current_streams = list(streams.items())
                for camera_id, stream in current_streams:
                    buffer = stream["buffer"]
                    if not buffer.empty():
                        frame = buffer.get()
                        with models_lock:
                            model_tree = models.get(camera_id, {}).get(stream["modelId"])
                        key = (camera_id, stream["modelId"])
                        persistent_data = persistent_data_dict.get(key, {})
                        updated_persistent_data = await handle_detection(
                            camera_id, stream, frame, websocket, model_tree, persistent_data
                        )
                        persistent_data_dict[key] = updated_persistent_data
                elapsed_time = (time.time() - start_time) * 1000  # ms
                sleep_time = max(poll_interval - elapsed_time, 0)
                logging.debug(f"Elapsed time: {elapsed_time}ms, sleeping for: {sleep_time}ms")
                await asyncio.sleep(sleep_time / 1000.0)
        except asyncio.CancelledError:
            logging.info("Stream processing task cancelled")
        except Exception as e:
            logging.error(f"Error in process_streams: {e}")

    async def send_heartbeat():
        while True:
            try:
                cpu_usage = psutil.cpu_percent()
                memory_usage = psutil.virtual_memory().percent
                if torch.cuda.is_available():
                    gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2)  # MB
                    gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2)  # MB
                else:
                    gpu_usage = None
                    gpu_memory_usage = None

                camera_connections = [
                    {
                        "cameraIdentifier": camera_id,
                        "modelId": stream["modelId"],
                        "modelName": stream["modelName"],
                        "online": True
                    }
                    for camera_id, stream in streams.items()
                ]

                state_report = {
                    "type": "stateReport",
                    "cpuUsage": cpu_usage,
                    "memoryUsage": memory_usage,
                    "gpuUsage": gpu_usage,
                    "gpuMemoryUsage": gpu_memory_usage,
                    "cameraConnections": camera_connections
                }
                await websocket.send_text(json.dumps(state_report))
                logging.debug("Sent stateReport as heartbeat")
                await asyncio.sleep(HEARTBEAT_INTERVAL)
            except Exception as e:
                logging.error(f"Error sending stateReport heartbeat: {e}")
                break

    async def on_message():
        while True:
            try:
                msg = await websocket.receive_text()
                logging.debug(f"Received message: {msg}")
                data = json.loads(msg)
                msg_type = data.get("type")

                if msg_type == "subscribe":
                    payload = data.get("payload", {})
                    camera_id = payload.get("cameraIdentifier")
                    rtsp_url = payload.get("rtspUrl")
                    model_url = payload.get("modelUrl")  # may be remote or local
                    modelId = payload.get("modelId")
                    modelName = payload.get("modelName")

                    if model_url:
                        with models_lock:
                            if camera_id not in models:
                                models[camera_id] = {}
                            if modelId not in models[camera_id]:
                                logging.info(f"Loading model from {model_url}")
                                extraction_dir = os.path.join("models", camera_id, str(modelId))
                                os.makedirs(extraction_dir, exist_ok=True)
                                # If model_url is remote, download it first.
                                parsed = urlparse(model_url)
                                if parsed.scheme in ("http", "https"):
                                    local_mpta = os.path.join(extraction_dir, os.path.basename(parsed.path))
                                    local_path = download_mpta(model_url, local_mpta)
                                    if not local_path:
                                        logging.error("Failed to download the remote mpta file.")
                                        continue
                                    model_tree = load_pipeline_from_zip(local_path, extraction_dir)
                                else:
                                    model_tree = load_pipeline_from_zip(model_url, extraction_dir)
                                if model_tree is None:
                                    logging.error("Failed to load model from mpta file.")
                                    continue
                                models[camera_id][modelId] = model_tree
                                logging.info(f"Loaded model {modelId} for camera {camera_id}")

                    if camera_id and rtsp_url:
                        with streams_lock:
                            if camera_id not in streams and len(streams) < max_streams:
                                cap = cv2.VideoCapture(rtsp_url)
                                if not cap.isOpened():
                                    logging.error(f"Failed to open RTSP stream for camera {camera_id}")
                                    continue
                                buffer = queue.Queue(maxsize=1)
                                stop_event = threading.Event()
                                thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer, stop_event))
                                thread.daemon = True
                                thread.start()
                                streams[camera_id] = {
                                    "cap": cap,
                                    "buffer": buffer,
                                    "thread": thread,
                                    "rtsp_url": rtsp_url,
                                    "stop_event": stop_event,
                                    "modelId": modelId,
                                    "modelName": modelName
                                }
                                logging.info(f"Subscribed to camera {camera_id} with modelId {modelId}, modelName {modelName}, URL {rtsp_url}")
                            elif camera_id and camera_id in streams:
                                # If already subscribed, unsubscribe
                                stream = streams.pop(camera_id)
                                stream["cap"].release()
                                logging.info(f"Unsubscribed from camera {camera_id}")
                                with models_lock:
                                    if camera_id in models and modelId in models[camera_id]:
                                        del models[camera_id][modelId]
                                        if not models[camera_id]:
                                            del models[camera_id]
                elif msg_type == "unsubscribe":
                    payload = data.get("payload", {})
                    camera_id = payload.get("cameraIdentifier")
                    logging.debug(f"Unsubscribing from camera {camera_id}")
                    with streams_lock:
                        if camera_id and camera_id in streams:
                            stream = streams.pop(camera_id)
                            stream["stop_event"].set()
                            stream["thread"].join()
                            stream["cap"].release()
                            logging.info(f"Unsubscribed from camera {camera_id}")
                            with models_lock:
                                if camera_id in models:
                                    del models[camera_id]
                elif msg_type == "requestState":
                    cpu_usage = psutil.cpu_percent()
                    memory_usage = psutil.virtual_memory().percent
                    if torch.cuda.is_available():
                        gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2)
                        gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2)
                    else:
                        gpu_usage = None
                        gpu_memory_usage = None

                    camera_connections = [
                        {
                            "cameraIdentifier": camera_id,
                            "modelId": stream["modelId"],
                            "modelName": stream["modelName"],
                            "online": True
                        }
                        for camera_id, stream in streams.items()
                    ]

                    state_report = {
                        "type": "stateReport",
                        "cpuUsage": cpu_usage,
                        "memoryUsage": memory_usage,
                        "gpuUsage": gpu_usage,
                        "gpuMemoryUsage": gpu_memory_usage,
                        "cameraConnections": camera_connections
                    }
                    await websocket.send_text(json.dumps(state_report))
                else:
                    logging.error(f"Unknown message type: {msg_type}")
            except json.JSONDecodeError:
                logging.error("Received invalid JSON message")
            except (WebSocketDisconnect, ConnectionClosedError) as e:
                logging.warning(f"WebSocket disconnected: {e}")
                break
            except Exception as e:
                logging.error(f"Error handling message: {e}")
                break

    try:
        await websocket.accept()
        stream_task = asyncio.create_task(process_streams())
        heartbeat_task = asyncio.create_task(send_heartbeat())
        message_task = asyncio.create_task(on_message())
        await asyncio.gather(heartbeat_task, message_task)
    except Exception as e:
        logging.error(f"Error in detect websocket: {e}")
    finally:
        stream_task.cancel()
        await stream_task
        with streams_lock:
            for camera_id, stream in streams.items():
                stream["stop_event"].set()
                stream["thread"].join()
                stream["cap"].release()
                while not stream["buffer"].empty():
                    try:
                        stream["buffer"].get_nowait()
                    except queue.Empty:
                        pass
                logging.info(f"Released camera {camera_id} and cleaned up resources")
            streams.clear()
        with models_lock:
            models.clear()
        logging.info("WebSocket connection closed")
