backend/main.py

"""
main.py — FastAPI backend for Traffic Monitoring System
Processes frames as fast as possible; frontend controls playback timing.
"""

import asyncio
import base64
import csv
import json
import re
import shutil
import subprocess
import time
import uuid
import threading
import urllib.parse
import urllib.request
from collections import defaultdict, deque
from pathlib import Path

import cv2
import numpy as np
from fastapi import FastAPI, File, Form, UploadFile, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from tracker import TrafficTracker, DEFAULT_CLASSES

# ─── App setup ────────────────────────────────────────────────────────────────
app = FastAPI(title="Traffic Monitoring API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_DIR = Path(__file__).parent.parent
UPLOAD_DIR = BASE_DIR / "uploads"; UPLOAD_DIR.mkdir(exist_ok=True)
LOG_DIR    = BASE_DIR / "logs";    LOG_DIR.mkdir(exist_ok=True)
OUTPUT_DIR = BASE_DIR / "output";  OUTPUT_DIR.mkdir(exist_ok=True)
STATIC_DIR = BASE_DIR / "frontend"

if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

# ─── Session store ─────────────────────────────────────────────────────────────
sessions: dict[str, dict] = {}
webcam_sessions: dict[str, dict] = {}
scene_lock = threading.Lock()


def _safe_log_path(filename: str) -> Path:
    path = (LOG_DIR / filename).resolve()
    if LOG_DIR.resolve() not in path.parents and path != LOG_DIR.resolve():
        raise HTTPException(400, "Invalid filename")
    return path


def _get_summary_for_session(sid: str) -> dict | None:
    if sid in sessions:
        return sessions[sid].get("summary")
    if sid in webcam_sessions:
        return webcam_sessions[sid].get("summary")

    for summary_path in LOG_DIR.glob("*_summary.json"):
        try:
            summary = json.loads(summary_path.read_text())
        except Exception:
            continue
        if summary.get("session_id") == sid:
            return summary
    return None


def _get_completed_session_summary(sid: str) -> dict:
    if sid in sessions:
        s = sessions[sid]
        if s["status"] != "done":
            raise HTTPException(202, "Processing not complete yet")

    summary = _get_summary_for_session(sid)
    if not summary:
        raise HTTPException(404, "Session not found")
    return summary


def _detection_csv_for_summary(summary: dict) -> Path | None:
    log_files = summary.get("log_files") or {}
    explicit_path = log_files.get("detections_csv")
    if explicit_path:
        path = Path(explicit_path)
        if path.exists():
            return path

    group_id = str(summary.get("group_id") or "").strip()
    scene = str(summary.get("scene") or "").strip()
    if group_id and scene:
        matches = sorted(
            LOG_DIR.glob(f"{group_id}_{scene}_*_detections.csv"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        if matches:
            return matches[0]

    return None


def _position_heatmap_for_summary(summary: dict, cols: int = 24, rows: int = 24) -> dict:
    csv_path = _detection_csv_for_summary(summary)
    cells: dict[tuple[int, int], int] = defaultdict(int)
    class_cells: dict[tuple[int, int], dict[str, int]] = defaultdict(lambda: defaultdict(int))
    total_points = 0

    if not csv_path:
        return {"cols": cols, "rows": rows, "max_density": 0, "total_points": 0, "cells": []}

    try:
        with csv_path.open(newline="", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                try:
                    cx = float(row.get("cx") or "")
                    cy = float(row.get("cy") or "")
                    width = float(row.get("frame_width") or 0)
                    height = float(row.get("frame_height") or 0)
                except ValueError:
                    continue

                if width <= 0 or height <= 0:
                    continue

                col = min(cols - 1, max(0, int((cx / width) * cols)))
                row_idx = min(rows - 1, max(0, int((cy / height) * rows)))
                class_name = str(row.get("class_name") or "unknown")
                cells[(col, row_idx)] += 1
                class_cells[(col, row_idx)][class_name] += 1
                total_points += 1
    except Exception:
        return {"cols": cols, "rows": rows, "max_density": 0, "total_points": 0, "cells": []}

    max_density = max(cells.values(), default=0)
    return {
        "cols": cols,
        "rows": rows,
        "max_density": int(max_density),
        "total_points": int(total_points),
        "cells": [
            {
                "x": x,
                "y": y,
                "count": int(count),
                "classes": dict(class_cells[(x, y)]),
            }
            for (x, y), count in sorted(cells.items(), key=lambda item: (item[0][1], item[0][0]))
        ],
    }


def _iter_known_scene_names(group_id: str):
    for session in sessions.values():
        if session.get("group_id") == group_id and session.get("scene_name"):
            yield session["scene_name"]
    for session in webcam_sessions.values():
        if session.get("group_id") == group_id and session.get("scene_name"):
            yield session["scene_name"]
    for summary_path in LOG_DIR.glob("*_summary.json"):
        try:
            summary = json.loads(summary_path.read_text())
        except Exception:
            continue
        if summary.get("group_id") == group_id and summary.get("scene"):
            yield summary["scene"]


def _next_scene_name(requested_scene: str, group_id: str) -> str:
    requested = (requested_scene or "scene").strip() or "scene"
    exact_exists = any(scene_name == requested for scene_name in _iter_known_scene_names(group_id))
    suffix_match = re.match(r"^(.+?)_(\d{2,})$", requested)
    base = suffix_match.group(1) if exact_exists and suffix_match else requested
    scene_pattern = re.compile(rf"^{re.escape(base)}(?:_(\d{{2,}}))?$")
    order_numbers = []

    for scene_name in _iter_known_scene_names(group_id):
        match = scene_pattern.match(scene_name)
        if not match:
            continue
        order_numbers.append(int(match.group(1) or 0))

    if not order_numbers:
        return base
    return f"{base}_{max(order_numbers) + 1:02d}"


def _ensure_browser_compatible_video(video_path: str) -> str:
    """
    OpenCV's mp4v output downloads fine, but many browsers will not play it in
    a <video> element. When ffmpeg is available, create an H.264/yuv420p copy.
    """
    source = Path(video_path)
    if not source.exists():
        return video_path
    if source.stem.endswith("_browser"):
        return str(source)

    target = source.with_name(f"{source.stem}_browser.mp4")
    if target.exists() and target.stat().st_mtime >= source.stat().st_mtime:
        return str(target)

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        return video_path

    cmd = [
        ffmpeg,
        "-y",
        "-i", str(source),
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-pix_fmt", "yuv420p",
        "-movflags", "+faststart",
        "-an",
        str(target),
    ]
    try:
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if target.exists() and target.stat().st_size > 0:
            return str(target)
    except (OSError, subprocess.CalledProcessError):
        pass

    return video_path


# ─── Routes ───────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def root():
    index = STATIC_DIR / "index.html"
    if index.exists():
        return HTMLResponse(index.read_text())
    return HTMLResponse("<h1>Traffic Monitoring API</h1>")


@app.get("/health")
async def health():
    return {"status": "ok", "sessions": len(sessions)}


@app.get("/classes")
async def get_classes():
    return {"classes": DEFAULT_CLASSES}


@app.post("/upload")
async def upload_video(
    file: UploadFile | None = File(None),
    video_url: str = Form(""),
    scene_name: str = Form("scene_01"),
    group_id: str = Form("Group_05"),
    classes: str = Form(""),
    conf: float = Form(0.5),
    model: str = Form("best.pt"),
):
    sid = str(uuid.uuid4())[:12]

    if file is not None and file.filename:
        suffix = Path(file.filename).suffix or ".mp4"
        video_name = file.filename
        video_path = UPLOAD_DIR / f"{sid}{suffix}"
        video_path.write_bytes(await file.read())
    elif video_url.strip():
        parsed = urllib.parse.urlparse(video_url.strip())
        if parsed.scheme not in ("http", "https"):
            raise HTTPException(400, "Video URL must start with http:// or https://")

        suffix = Path(parsed.path).suffix or ".mp4"
        video_name = Path(parsed.path).name or f"{sid}{suffix}"
        video_path = UPLOAD_DIR / f"{sid}{suffix}"
        try:
            with urllib.request.urlopen(video_url.strip(), timeout=30) as response:
                with open(video_path, "wb") as out:
                    shutil.copyfileobj(response, out)
        except Exception as exc:
            raise HTTPException(400, f"Could not download video URL: {exc}") from exc
    else:
        raise HTTPException(400, "Upload a video file or provide a video URL")

    selected = [c.strip() for c in classes.split(",") if c.strip()] or DEFAULT_CLASSES

    with scene_lock:
        scene_name = _next_scene_name(scene_name, group_id)
        sessions[sid] = {
            "status": "processing",
            "frames": [],
            "frames_lock": threading.Lock(),
            "stats": {},
            "summary": None,
            "error": None,
            "scene_name": scene_name,
            "group_id": group_id,
            "video_name": video_name,
            "progress": 0,
            "total_frames": 0,
            "fps": 30.0,
        }

    threading.Thread(
        target=_process_video_thread,
        args=(sid, str(video_path), scene_name, group_id, video_name, selected, conf, model),
        daemon=True,
    ).start()

    return {"session_id": sid, "scene_name": scene_name, "group_id": group_id, "status": "processing"}


@app.post("/webcam/start")
async def start_webcam_detection(
    scene_name: str = Form("webcam"),
    group_id: str = Form("Group_05"),
    classes: str = Form(""),
    conf: float = Form(0.5),
    model: str = Form("best.pt"),
):
    sid = str(uuid.uuid4())[:12]
    with scene_lock:
        scene_name = _next_scene_name(scene_name, group_id)
        webcam_sessions[sid] = {
            "status": "loading",
            "tracker": None,
            "tracker_lock": threading.Lock(),
            "summary": None,
            "error": None,
            "scene_name": scene_name,
            "group_id": group_id,
        }

    selected = [c.strip() for c in classes.split(",") if c.strip()] or DEFAULT_CLASSES

    try:
        webcam_sessions[sid]["tracker"] = TrafficTracker(
            model_path=model,
            selected_classes=selected,
            conf_threshold=conf,
            scene_name=scene_name,
            group_id=group_id,
            video_name="webcam",
            output_dir=str(LOG_DIR),
        )
        webcam_sessions[sid]["status"] = "processing"
    except Exception as exc:
        webcam_sessions[sid]["status"] = "error"
        webcam_sessions[sid]["error"] = str(exc)
        raise HTTPException(500, f"Could not start webcam detection: {exc}") from exc

    return {"session_id": sid, "scene_name": scene_name, "group_id": group_id, "status": "processing"}


@app.post("/webcam/frame/{sid}")
async def detect_webcam_frame(sid: str, file: UploadFile = File(...)):
    session = webcam_sessions.get(sid)
    if not session:
        raise HTTPException(404, "Webcam session not found")
    if session["status"] == "error":
        raise HTTPException(500, session.get("error") or "Webcam session failed")
    tracker = session.get("tracker")
    if tracker is None:
        raise HTTPException(202, "Webcam model is still loading")

    raw = await file.read()
    np_buf = np.frombuffer(raw, dtype=np.uint8)
    frame = cv2.imdecode(np_buf, cv2.IMREAD_COLOR)
    if frame is None:
        raise HTTPException(400, "Could not decode webcam frame")

    with session["tracker_lock"]:
        if tracker.frame_width == 0 or tracker.frame_height == 0:
            tracker.setup_frame_source(frame.shape[1], frame.shape[0], fps=5.0)

        annotated, frame_stat = tracker.process_frame(frame)

    _, buf = cv2.imencode(".jpg", annotated, [cv2.IMWRITE_JPEG_QUALITY, 80])
    cumulative = frame_stat.get("cumulative") or {}
    stats_payload = {
        "frame":       frame_stat.get("frame", 0),
        "timestamp":   round(frame_stat.get("timestamp", 0.0), 2),
        "scene":       session["scene_name"],
        "detections":  frame_stat.get("detections", 0),
        "any_visible": frame_stat.get("any_visible", False),
        "cumulative":  cumulative,
        "total_unique": int(sum(cumulative.values())),
    }

    return {
        "frame": base64.b64encode(buf).decode(),
        "stats": stats_payload,
    }


@app.post("/webcam/stop/{sid}")
async def stop_webcam_detection(sid: str):
    session = webcam_sessions.get(sid)
    if not session:
        raise HTTPException(404, "Webcam session not found")
    tracker = session.get("tracker")
    if tracker is None:
        webcam_sessions.pop(sid, None)
        return {"status": "stopped"}

    with session["tracker_lock"]:
        log_paths = tracker.save_logs()
        summary = tracker.get_summary()
        summary["session_id"] = sid
        summary["log_files"] = log_paths
        summary["total_unique"] = int(sum(summary.get("count_per_class", {}).values()))
        summary_path = Path(log_paths.get("summary", ""))
        if summary_path.exists():
            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")

    session["summary"] = summary
    session["status"] = "done"
    return {"status": "stopped", "summary": summary}


def _process_video_thread(sid, video_path, scene_name, group_id, video_name, selected_classes, conf, model_name):
    """
    Processes video as fast as possible (no artificial sleep).
    Each frame carries a `pts` (presentation timestamp in seconds) so the
    frontend can pace playback at the original video speed.
    """
    session = sessions[sid]
    try:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            session["status"] = "error"
            session["error"] = "Could not open video file."
            return

        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        session["total_frames"] = total

        tracker = TrafficTracker(
            model_path=model_name,
            selected_classes=selected_classes,
            conf_threshold=conf,
            scene_name=scene_name,
            group_id=group_id,
            video_name=video_name,
            output_dir=str(LOG_DIR),
        )
        tracker.setup_video(cap)
        session["fps"] = tracker.fps

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            annotated, frame_stat = tracker.process_frame(frame)

            # Encode to JPEG — quality 80 is a good speed/size balance
            _, buf = cv2.imencode(".jpg", annotated, [cv2.IMWRITE_JPEG_QUALITY, 80])
            b64 = base64.b64encode(buf).decode()

            cumulative = frame_stat.get("cumulative") or {}
            stats_payload = {
                "frame":       frame_stat.get("frame", 0),
                "timestamp":   round(frame_stat.get("timestamp", 0.0), 2),
                "scene":       scene_name,
                "detections":  frame_stat.get("detections", 0),
                "any_visible": frame_stat.get("any_visible", False),
                "cumulative":  cumulative,
                "total_unique": int(sum(cumulative.values())),
            }

            item = {
                "frame": b64,
                "stats": stats_payload,
                # pts tells the client when to display this frame (seconds from video start)
                "pts": stats_payload["timestamp"],
            }

            with session["frames_lock"]:
                session["frames"].append(item)

            session["stats"] = stats_payload
            session["progress"] = tracker.frame_index

        cap.release()

        log_paths = tracker.save_logs()
        summary = tracker.get_summary()
        summary["session_id"] = sid
        summary["log_files"] = log_paths
        summary["total_unique"] = int(sum(summary.get("count_per_class", {}).values()))
        summary_path = Path(log_paths.get("summary", ""))
        if summary_path.exists():
            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")

        session["summary"] = summary
        session["status"] = "done"

    except Exception as e:
        sessions[sid]["status"] = "error"
        sessions[sid]["error"] = str(e)
        raise


@app.get("/stream/{sid}")
async def stream_frames(sid: str):
    """
    SSE — pushes every processed frame to the client in order.
    Each frame event includes a `pts` field (seconds) so the browser can
    schedule rendering at the correct wall-clock time for smooth playback.
    """
    if sid not in sessions:
        raise HTTPException(404, "Session not found")

    async def generator():
        sent = 0
        while True:
            session = sessions[sid]

            with session["frames_lock"]:
                frames_snapshot = session["frames"][sent:]

            for item in frames_snapshot:
                payload = json.dumps({
                    "frame": item["frame"],
                    "stats": item["stats"],
                    "pts":   item["pts"],
                })
                yield f"data: {payload}\n\n"
                sent += 1

            if session["status"] in ("done", "error"):
                # Drain any remaining frames produced between last check and now
                with session["frames_lock"]:
                    remaining = session["frames"][sent:]
                for item in remaining:
                    payload = json.dumps({
                        "frame": item["frame"],
                        "stats": item["stats"],
                        "pts":   item["pts"],
                    })
                    yield f"data: {payload}\n\n"
                    sent += 1

                final = json.dumps({
                    "event":   "done",
                    "status":  session["status"],
                    "summary": session.get("summary"),
                    "error":   session.get("error"),
                })
                yield f"data: {final}\n\n"
                break

            await asyncio.sleep(0.01)

    return StreamingResponse(generator(), media_type="text/event-stream")


@app.get("/status/{sid}")
async def get_status(sid: str):
    if sid not in sessions:
        raise HTTPException(404, "Session not found")
    s = sessions[sid]
    return {
        "session_id":   sid,
        "status":       s["status"],
        "scene_name":   s["scene_name"],
        "progress":     s["progress"],
        "total_frames": s["total_frames"],
        "fps":          s["fps"],
        "stats":        s["stats"],
        "error":        s["error"],
    }


@app.get("/summary/{sid}")
async def get_summary(sid: str):
    try:
        return _get_completed_session_summary(sid)
    except HTTPException as exc:
        if exc.status_code == 202:
            return JSONResponse({"detail": "Not complete yet."}, status_code=202)
        raise


@app.get("/logs")
async def list_logs():
    files = [
        {
            "name":     f.name,
            "size_kb":  round(f.stat().st_size / 1024, 1),
            "modified": f.stat().st_mtime,
        }
        for f in sorted(LOG_DIR.glob("*"), key=lambda x: x.stat().st_mtime, reverse=True)
    ]
    return {"logs": files}


@app.get("/videos")
async def list_annotated_videos():
    videos = [
        {
            "name":     f.name,
            "size_kb":  round(f.stat().st_size / 1024, 1),
            "modified": f.stat().st_mtime,
        }
        for f in sorted(LOG_DIR.glob("*.mp4"), key=lambda x: x.stat().st_mtime, reverse=True)
    ]
    return {"videos": videos}


@app.get("/dashboard")
async def dashboard_data():
    completed = [
        dict(s["summary"]) for s in sessions.values()
        if s["status"] == "done" and s["summary"]
    ]
    completed.extend(
        dict(s["summary"]) for s in webcam_sessions.values()
        if s["status"] == "done" and s["summary"]
    )
    for p in sorted(LOG_DIR.glob("*_summary.json")):
        try:
            data = json.loads(p.read_text())
            if not any(c.get("session_id") == data.get("session_id") for c in completed):
                completed.append(dict(data))
        except Exception:
            pass

    global_counts: dict[str, int] = defaultdict(int)
    for s in completed:
        s["position_heatmap"] = _position_heatmap_for_summary(s)
        for cls_, cnt in s.get("count_per_class", {}).items():
            global_counts[cls_] += cnt

    return {
        "scenes":        completed,
        "global_counts": dict(global_counts),
        "total_scenes":  len(completed),
        "total_objects": int(sum(global_counts.values())),
    }


@app.get("/log/{filename}")
async def download_log(filename: str):
    p = _safe_log_path(filename)
    if not p.exists():
        raise HTTPException(404)
    return FileResponse(str(p), filename=filename)


@app.get("/download/video/{sid}")
async def download_annotated_video(sid: str):
    """
    Download the annotated output video (.mp4) for a completed session.
    Only available after status == 'done'.
    """
    summary = _get_completed_session_summary(sid)
    video_path = summary.get("annotated_video", "")
    if not video_path or not Path(video_path).exists():
        raise HTTPException(404, "Annotated video not found")
    return FileResponse(
        str(video_path),
        media_type="video/mp4",
        filename=Path(video_path).name,
    )


@app.get("/download/video-file/{filename}")
async def download_annotated_video_file(filename: str):
    video_path = _safe_log_path(filename)
    if video_path.suffix.lower() != ".mp4" or not video_path.exists():
        raise HTTPException(404, "Annotated video not found")
    return FileResponse(
        str(video_path),
        media_type="video/mp4",
        filename=video_path.name,
    )


def _stream_video_file(video_path: str, request: Request):
    video_path = _ensure_browser_compatible_video(video_path)

    file_path = Path(video_path)
    file_size = file_path.stat().st_size

    range_header = request.headers.get("range")

    def iterate_file(file_path, start=0, end=None, chunk_size=1024 * 1024):
        with open(file_path, "rb") as f:
            f.seek(start)
            remaining = None if end is None else end - start + 1
            while True:
                if remaining is not None and remaining <= 0:
                    break
                read_size = chunk_size if remaining is None else min(chunk_size, remaining)
                chunk = f.read(read_size)
                if not chunk:
                    break
                if remaining is not None:
                    remaining -= len(chunk)
                yield chunk

    headers = {
        "Accept-Ranges": "bytes",
        "Content-Disposition": f"inline; filename={file_path.name}",
    }

    if range_header:
        try:
            range_value = range_header.strip().lower().replace("bytes=", "", 1)
            start_text, end_text = range_value.split("-", 1)
            start = int(start_text) if start_text else 0
            end = int(end_text) if end_text else file_size - 1
            start = max(0, min(start, file_size - 1))
            end = max(start, min(end, file_size - 1))
        except ValueError:
            raise HTTPException(416, "Invalid range")

        content_length = end - start + 1
        headers.update({
            "Content-Range": f"bytes {start}-{end}/{file_size}",
            "Content-Length": str(content_length),
        })

        return StreamingResponse(
            iterate_file(file_path, start, end),
            status_code=206,
            media_type="video/mp4",
            headers=headers,
        )

    headers["Content-Length"] = str(file_size)
    
    return StreamingResponse(
        iterate_file(file_path),
        media_type="video/mp4",
        headers=headers,
    )


@app.get("/stream/video/{sid}")
async def stream_annotated_video(sid: str, request: Request):
    """
    Stream the annotated output video for a completed session.
    Returns video/mp4 with proper streaming headers.
    """
    summary = _get_completed_session_summary(sid)
    video_path = summary.get("annotated_video", "")
    if not video_path or not Path(video_path).exists():
        raise HTTPException(404, "Annotated video not found")
    return _stream_video_file(video_path, request)


@app.get("/stream/video-file/{filename}")
async def stream_annotated_video_file(filename: str, request: Request):
    video_path = _safe_log_path(filename)
    if video_path.suffix.lower() != ".mp4" or not video_path.exists():
        raise HTTPException(404, "Annotated video not found")
    return _stream_video_file(str(video_path), request)