backend/server.py

import os
import json
import uuid
import asyncio
import tempfile
import shutil
from pathlib import Path
import zipfile

import cv2
import time
from fastapi import FastAPI, WebSocket, UploadFile, File, BackgroundTasks, Request
from fastapi.responses import FileResponse, Response, JSONResponse
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware
import resend

from model import load_model
from config import get_optimal_config
from engine import run
from pcu import MODEL_CLASSES
from auth import verify_google_token, get_or_create_user, set_username, GOOGLE_CLIENT_ID
from visualize import generate_all

BUSINESS_MAP = {
    "Cars": [0, 1, 2, 3, 12],
    "Buses": [4, 9, 10],
    "Two-wheelers": [7, 11],
    "Three-wheelers": [6],
    "Trucks": [5, 8],
    "Others": [13]
}

from contextlib import asynccontextmanager
import numpy as np

BASE = Path(__file__).parent.parent
FRONTEND = BASE / "frontend"
UPLOAD_DIR = Path(tempfile.gettempdir()) / "funky_uploads"
UPLOAD_DIR.mkdir(exist_ok=True)
REPORT_DIR = Path(tempfile.gettempdir()) / "funky_reports"
REPORT_DIR.mkdir(exist_ok=True)

videos = {}
video_info = {}
run_results = {}
model = None

MAX_MEMORY_ENTRIES = 50
def evict_old(d):
    while len(d) > MAX_MEMORY_ENTRIES:
        d.pop(next(iter(d)))

UPLOAD_LIMITS = {}
def is_rate_limited(ip: str):
    now = time.time()
    stamps = [t for t in UPLOAD_LIMITS.get(ip, []) if now - t < 60]
    if len(stamps) >= 5:
        UPLOAD_LIMITS[ip] = stamps
        return True
    stamps.append(now)
    UPLOAD_LIMITS[ip] = stamps
    return False

MAX_UPLOAD_BYTES = 500 * 1024 * 1024

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model = load_model()
    # Warm up: run a dummy inference so OpenVINO compiles its graph now,
    # not on the first real user request
    try:
        dummy_img = np.zeros((736, 736, 3), dtype=np.uint8)
        model([dummy_img, dummy_img], verbose=False)  # list of 2 imgs triggers batch=2
        print("[BACKEND] Model warm-up complete.")
    except Exception as e:
        print(f"[BACKEND] Warm-up skipped: {e}")
    yield
    # Shutdown: nothing to clean up


app = FastAPI(lifespan=lifespan)

class SecurityHeaders(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        # response.headers["X-Frame-Options"] = "SAMEORIGIN"  # Removed for HF Spaces compatibility
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        return response

app.add_middleware(SecurityHeaders)

@app.get("/")
def index():
    return FileResponse(FRONTEND / "initial.html")


@app.get("/vehicles.html")
@app.get("/vehicles")
def vehicles():
    return FileResponse(FRONTEND / "vehicles.html")


@app.post("/upload")
async def upload(request: Request, file: UploadFile = File(...)):
    client_ip = request.client.host if request.client else "unknown"
    if is_rate_limited(client_ip):
        return JSONResponse({"error": "Rate limit exceeded. Please wait a minute."}, status_code=429)

    if not file.content_type.startswith("video/"):
        return JSONResponse({"error": "Invalid file type. Only videos are allowed."}, status_code=400)

    if hasattr(file, 'size') and file.size and file.size > MAX_UPLOAD_BYTES:
        return JSONResponse({"error": "File too large. Maximum size is 500MB."}, status_code=413)

    video_id = str(uuid.uuid4())[:8]
    path = UPLOAD_DIR / f"{video_id}.mp4"

    # Clean up any previous temp uploads to avoid stale state
    for old_path in UPLOAD_DIR.glob("*.mp4"):
        try:
            old_path.unlink()
        except Exception:
            pass

    print(f"[BACKEND] Received upload request: {file.filename}")
    try:
        with open(path, "wb") as f:
            shutil.copyfileobj(file.file, f)

        file_size = os.path.getsize(path)
        if file_size > MAX_UPLOAD_BYTES:
            path.unlink()
            return JSONResponse({"error": "File too large. Maximum size is 500MB."}, status_code=413)

        print(f"[BACKEND] Successfully stored: {path} ({file_size} bytes)")

        videos[video_id] = str(path)
        video_info[video_id] = file.filename
        evict_old(videos)
        evict_old(video_info)
        return {"video_id": video_id}
    except Exception as e:
        print(f"[BACKEND] Upload failed: {str(e)}")
        return JSONResponse({"error": str(e)}, status_code=500)


@app.get("/config/{video_id}")
def config_endpoint(video_id: str):
    path = videos.get(video_id)
    cfg = get_optimal_config(path)
    return cfg


@app.get("/first-frame/{video_id}")
def first_frame(video_id: str):
    path = videos.get(video_id)
    cap = cv2.VideoCapture(path)
    ret, frame = cap.read()
    cap.release()
    _, buf = cv2.imencode(".jpg", frame)
    return Response(content=buf.tobytes(), media_type="image/jpeg")


@app.get("/constants")
def constants():
    return {"classes": MODEL_CLASSES, "business_map": BUSINESS_MAP}


@app.post("/reports/{video_id}")
def generate_reports(video_id: str):
    data = run_results.get(video_id)
    if not data:
        return {"error": "no results", "files": []}
    out_dir = str(REPORT_DIR / video_id)
    report_format = data.get("report_format", "png")
    files = generate_all(data, MODEL_CLASSES, out_dir, report_format)
    # Include annotated video if it exists
    annotated_src = data.get("annotated_video")
    if annotated_src and os.path.exists(annotated_src):
        dest = os.path.join(out_dir, "annotated.mp4")
        shutil.copy2(annotated_src, dest)  # Always overwrite — never skip
        files.append("annotated.mp4")
    return {"files": files}


@app.get("/reports/{video_id}/{name}")
def get_report(video_id: str, name: str):
    safe_name = Path(name).name
    path = REPORT_DIR / video_id / safe_name
    if not path.resolve().is_relative_to(REPORT_DIR.resolve()):
        return JSONResponse({"error": "Invalid path"}, status_code=400)
    if not path.exists():
        return JSONResponse({"error": "File not found"}, status_code=404)
    media = "image/png"
    if name.endswith(".pdf"):
        media = "application/pdf"
    elif name.endswith(".mp4"):
        media = "video/mp4"
    elif name.endswith(".json"):
        media = "application/json"
    elif name.endswith(".csv"):
        media = "text/csv"

    return FileResponse(str(path), media_type=media)


@app.get("/bundle/{video_id}")
def download_all_reports(video_id: str):
    print(f"[BACKEND] ZIP request for {video_id}")
    base_path = REPORT_DIR / video_id
    if not base_path.exists():
        print(f"[BACKEND] Error: {base_path} not found")
        return JSONResponse({"error": f"Report directory not found for {video_id}"}, status_code=404)
    
    try:
        zip_filename = f"bundle_{video_id}.zip"
        zip_path = REPORT_DIR / zip_filename
        
        if zip_path.exists():
            zip_path.unlink()
            
        print(f"[BACKEND] Creating ZIP: {zip_path}")
        # shutil.make_archive adds the .zip extension, so we strip it from the target path
        shutil.make_archive(str(REPORT_DIR / f"bundle_{video_id}"), 'zip', str(base_path))
        
        if not zip_path.exists():
             raise Exception("Zip file was not created")
             
        original_name = video_info.get(video_id, "UrbanFlow_Analysis")
        safe_name = "".join(x for x in original_name if x.isalnum() or x in "._-").rsplit(".", 1)[0]
        
        print(f"[BACKEND] Serving ZIP: {zip_path}")
        return FileResponse(
            str(zip_path), 
            media_type="application/zip", 
            filename=f"{safe_name}_UrbanFlow.zip"
        )
    except Exception as e:
        import traceback
        print(f"[BACKEND] ZIP Error: {str(e)}\n{traceback.format_exc()}")
        return JSONResponse({"error": str(e)}, status_code=500)


FEEDBACK_PATH = Path(tempfile.gettempdir()) / "urbanflow_feedback.json"


def send_feedback_email(api_key, feedback):
    try:
        resend.api_key = api_key
        
        fb_type = feedback.get('type') or 'General'
        rating = feedback.get('rating', 0)
        details = feedback.get('details') or ""
        usecase = feedback.get('usecase') or "Not specified"
        emojis = feedback.get('emojis', {})
        priorities = feedback.get('priorities', [])
        
        # Check if it's stars-only (no emojis, no priorities, no text)
        has_emojis = any(v for v in emojis.values())
        has_priorities = len(priorities) > 0
        has_text = bool(details and details.strip())
        is_stars_only = not (has_emojis or has_priorities or has_text)

        def get_emoji_row(label, choice, custom_options=None):
            options = custom_options or ["Poor", "Fair", "Good", "Great"]
            c = (choice or "").title()
            
            row = f"<div style='margin-bottom: 16px;'><span style='font-size: 10px; font-weight: 800; color: #888; text-transform: uppercase; letter-spacing: 1px;'>{label}</span><br><div style='margin-top: 6px;'>"
            for opt in options:
                if c == opt:
                    row += f"<span style='display: inline-block; background: #c89a6c; color: #000; font-size: 10px; font-weight: bold; padding: 4px 12px; border-radius: 6px; margin-right: 8px;'>{opt}</span>"
                else:
                    row += f"<span style='display: inline-block; background: #fff; color: #ccc; font-size: 10px; padding: 3px 11px; border-radius: 6px; border: 1px solid #eee; margin-right: 8px;'>{opt}</span>"
            row += "</div></div>"
            return row

        # Section 1: Experience Metrics (Above)
        metrics_html = ""
        if not is_stars_only:
            metrics_html += get_emoji_row("Recommend Product", emojis.get("fb-recommend"), ["Unlikely", "Maybe", "Likely", "Highly"])
            metrics_html += get_emoji_row("Security Assessment", emojis.get("fb-security"))
            metrics_html += get_emoji_row("Integration Readiness", emojis.get("fb-integration"))
            metrics_html += get_emoji_row("Ease of Use", emojis.get("fb-ease"))

        # Section 2: Priorities (Below)
        priority_section = ""
        if has_priorities:
            priority_section = "<p style='font-size: 10px; font-weight: 800; color: #888; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 10px;'>Feature Prioritization:</p>"
            for p in priorities:
                priority_section += f"<div style='background: #fff; border-left: 4px solid #c89a6c; padding: 10px 16px; font-size: 13px; font-weight: 600; margin-bottom: 8px; border-top: 1px solid #f0f0f0; border-right: 1px solid #f0f0f0; border-bottom: 1px solid #f0f0f0; border-radius: 0 8px 8px 0;'>{p}</div>"

        # Categorization & Feedback Content
        categorization_html = ""
        if not is_stars_only:
            categorization_html = f"""
            <div style="background: #fafafa; padding: 25px; border-radius: 16px; margin: 30px 0; border: 1px solid #f0f0f0;">
                <table style="width: 100%; font-size: 13px;">
                    <tr><td style="color: #999; width: 40%; padding: 6px 0; text-transform: uppercase; font-size: 9px; font-weight: 800; letter-spacing: 0.5px;">Primary Use Case</td><td style="font-weight: 700; color: #333;">{usecase}</td></tr>
                    <tr><td style="color: #999; padding: 6px 0; text-transform: uppercase; font-size: 9px; font-weight: 800; letter-spacing: 0.5px;">Feedback Category</td><td style="font-weight: 700; color: #333;">{fb_type}</td></tr>
                </table>
            </div>
            """

        detailed_feedback_html = ""
        if has_text:
            detailed_feedback_html = f"""
            <div style="margin-top: 30px;">
                <p style="font-size: 10px; font-weight: 800; color: #888; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 12px;">Detailed Word Feedback</p>
                <div style="background: #fff; border: 1px solid #eee; padding: 25px; border-radius: 14px; font-size: 14px; color: #444; white-space: pre-wrap; line-height: 1.8; box-shadow: 0 2px 10px rgba(0,0,0,0.02);">{details}</div>
            </div>
            """

        # Header with Rating
        user_email = feedback.get('user_email', '')
        user_email_html = ''
        if user_email:
            user_email_html = f"""
            <div style="margin-top: 14px; padding: 10px 20px; background: #fff; border: 1px solid #eee; border-radius: 10px; display: inline-block;">
                <span style="font-size: 9px; font-weight: 800; color: #999; text-transform: uppercase; letter-spacing: 1px;">Submitted by</span><br>
                <a href="mailto:{user_email}" style="color: #8b5e3c; font-size: 13px; font-weight: 700; text-decoration: none;">{user_email}</a>
            </div>
            """

        header_rating = f"""
        <div style="text-align: center; margin-bottom: 40px; padding: 20px; background: linear-gradient(180deg, #fff 0%, #fafafa 100%); border-radius: 20px;">
            <h2 style="color: #8b5e3c; margin: 0; font-size: 26px; font-weight: 900; letter-spacing: -1px;">UrbanFlow Intelligence</h2>
            <div style="margin-top: 15px; font-size: 22px; color: #c89a6c; letter-spacing: 4px;">{'★' * rating}{'☆' * (5-rating)}</div>
            <p style="color: #aaa; font-size: 10px; text-transform: uppercase; letter-spacing: 2px; margin-top: 10px; font-weight: 700;">Overall Experience: {rating}/5 Stars</p>
            {user_email_html}
        </div>
        """

        # Assemble the body
        final_content = header_rating
        if is_stars_only:
            final_content += "<div style='text-align: center; padding: 40px; color: #777; font-style: italic; font-size: 14px;'>User submitted an overall star rating with no additional details.</div>"
        else:
            final_content += f"""
                <div style="margin-bottom: 40px;">{metrics_html}</div>
                {categorization_html}
                <div style="margin-bottom: 40px;">{priority_section}</div>
                {detailed_feedback_html}
            """

        template_path = BASE / "backend" / "email_template.html"
        html_template = template_path.read_text(encoding="utf-8") if template_path.exists() else "{final_content}"
        html_body = html_template.format(final_content=final_content)
        
        resend.Emails.send({
            "from": "UrbanFlow <onboarding@resend.dev>",
            "to": "support.urbanflow365@gmail.com",
            "subject": f"Feedback: {fb_type} - {rating}/5 Stars" + (f" [{feedback.get('user_email', '')}]" if feedback.get('user_email') else ""),
            "html": html_body
        })
        print(f"[BACKEND] Feedback email successfully transmitted via Resend.")
    except Exception as e:
        print(f"[BACKEND] Resend Error: {str(e)}")

# =========== Auth Endpoints ===========

@app.get("/api/auth/client-id")
def auth_client_id():
    if not GOOGLE_CLIENT_ID:
        return JSONResponse({"error": "GOOGLE_CLIENT_ID not configured"}, status_code=500)
    return {"client_id": GOOGLE_CLIENT_ID}


@app.post("/api/auth/verify")
async def auth_verify(request_data: dict):
    credential = request_data.get("credential", "")
    if not credential:
        return JSONResponse({"error": "Missing credential"}, status_code=400)
    try:
        token_info = verify_google_token(credential)
    except ValueError as e:
        return JSONResponse({"error": str(e)}, status_code=401)
    user = get_or_create_user(token_info["email"], token_info["name"], token_info["picture"])
    return user


@app.post("/api/auth/onboard")
async def auth_onboard(request_data: dict):
    email = request_data.get("email", "")
    username = request_data.get("username", "")
    if not email or not username:
        return JSONResponse({"error": "Email and username required"}, status_code=400)
    ok = set_username(email, username)
    if not ok:
        return JSONResponse({"error": "User not found"}, status_code=404)
    return {"status": "ok"}

@app.post("/api/event")
async def track_event(request_data: dict):
    event = request_data.get("event", "UNKNOWN")
    meta = request_data.get("meta", {})
    print(f"[ANALYTICS] EVENT: {event} | {meta}")
    return {"status": "ok"}


# =========== Feedback ===========

@app.post("/api/feedback")
async def submit_feedback(background_tasks: BackgroundTasks, request_data: dict = None):
    from datetime import datetime, timezone
    feedback = request_data or {}
    feedback["timestamp"] = datetime.now(timezone.utc).isoformat()
    
    def write_feedback(fb):
        entries = []
        if FEEDBACK_PATH.exists():
            try:
                entries = json.loads(FEEDBACK_PATH.read_text())
            except Exception:
                pass
        entries.append(fb)
        FEEDBACK_PATH.write_text(json.dumps(entries, indent=2))
        
    background_tasks.add_task(write_feedback, feedback)
    
    # Trigger Email via Resend if API key is present
    resend_key = os.getenv("RESEND_API_KEY")
    if resend_key:
        background_tasks.add_task(send_feedback_email, resend_key, feedback)
        print(f"[BACKEND] Email task queued for support.urbanflow365@gmail.com")
    else:
        print(f"[BACKEND] Resend skipped: RESEND_API_KEY secret not found.")
        
    return {"status": "ok"}


@app.websocket("/ws/run")
async def ws_run(ws: WebSocket):
    await ws.accept()
    data = json.loads(await ws.receive_text())

    video_id = data["video_id"]
    line = data["line"]
    cfg = data["config"]
    save_annotated = data.get("annotated_video", False)
    annotated_options = data.get("annotated_options", {"bbox": True, "track_id": True, "spatial": True})
    report_format = data.get("report_format", "png")

    path = videos.get(video_id)
    if not path:
        await ws.send_text(json.dumps({"error": f"Unknown video_id: {video_id}"}))
        await ws.close()
        return

    loop = asyncio.get_running_loop()
    queue = asyncio.Queue()

    def on_frame(update):
        loop.call_soon_threadsafe(queue.put_nowait, update)

    task = loop.run_in_executor(
        None, run, model, path, line, cfg, on_frame, save_annotated, annotated_options
    )

    try:
        while True:
            done = task.done()
            while not queue.empty():
                update = queue.get_nowait()
                await ws.send_text(json.dumps(update))

            if done:
                break

            await asyncio.sleep(0.05)

        result = task.result()  # re-raises any exception from the engine
        result["report_format"] = report_format
        result["video_path"] = path
        result["video_meta"] = {
            "filename": video_info.get(video_id, "unknown"),
            "fps": cfg.get("video_fps", 0),
            "frames": cfg.get("frames", 0),
            "duration": cfg.get("duration", 0),
            "resolution": cfg.get("resolution", [0, 0]),
            "pixels": cfg.get("pixels", 0),
        }
        result["engine_config"] = cfg
        result["export_json"] = data.get("export_json", False)
        result["export_csv"] = data.get("export_csv", False)
        run_results[video_id] = result
        evict_old(run_results)
        await ws.send_text(json.dumps({
            "done": True,
            "video_id": video_id,
            "processing_time": result["processing_time"],
            "actual_fps": result["actual_fps"],
            "speed_vs_realtime": result["speed_vs_realtime"],
            "pcu": result.get("pcu", {}),
        }))
        await ws.close()

    except Exception as e:
        import traceback
        err_msg = traceback.format_exc()
        print(f"[BACKEND] Engine error for {video_id}:\n{err_msg}")
        try:
            await ws.send_text(json.dumps({"error": str(e), "detail": err_msg}))
            await ws.close()
        except Exception:
            pass


class SmartCacheStaticFiles(StaticFiles):
    def is_not_modified(self, response_headers, request_headers) -> bool:
        return False

    async def get_response(self, path: str, scope):
        resp = await super().get_response(path, scope)
        if path.endswith(".html") or path in ("", "/"):
            resp.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
        else:
            resp.headers["Cache-Control"] = "public, max-age=3600"
        return resp

app.mount("/", SmartCacheStaticFiles(directory=str(FRONTEND)), name="frontend")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)