#!/usr/bin/env python3.11
"""
AI Video Generation Swarm - Main Application
Lazy loading version - no model loading at startup
"""

import os
import sys
import json
import uuid
import time
import asyncio
import hashlib
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI, HTTPException, Depends, WebSocket, WebSocketDisconnect, UploadFile, File, Form
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from jose import JWTError, jwt
from passlib.context import CryptContext
import boto3
from botocore.config import Config

# ───────────────────────────────────────────────────────────────
# CONFIGURATION
# ───────────────────────────────────────────────────────────────

HF_TOKEN = os.getenv("HF_TOKEN", "")
PASS_KEY = os.getenv("PASS_KEY", "default_secret_change_me")
R2_ACCOUNT_ID = os.getenv("R2_ACCOUNT_ID", "")
R2_BUCKET_NAME = os.getenv("R2_BUCKET_NAME", "video-ai-swarm")
R2_ACCESS_KEY = os.getenv("R2_ACCESS_KEY_ID", "")
R2_SECRET_KEY = os.getenv("R2_SECRET_ACCESS_KEY", "")
R2_PUBLIC_DOMAIN = os.getenv("R2_PUBLIC_DOMAIN", "")

SECRET_KEY = hashlib.sha256(PASS_KEY.encode()).hexdigest()
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_HOURS = 24

GPU_MAP = {
    "orchestrator": 7,
    "video_primary": 0,
    "video_backup": 1,
    "face_restore": 2,
    "frame_interpolate": 3,
    "video_upscale": 4,
    "image_generator": 5,
    "audio_generator": 6,
    "quality_evaluator": 6,
}

# ───────────────────────────────────────────────────────────────
# AUTHENTICATION
# ───────────────────────────────────────────────────────────────

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
security = HTTPBearer()

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    token = credentials.credentials
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

def verify_pass_key(key: str) -> bool:
    return key == PASS_KEY

# ───────────────────────────────────────────────────────────────
# R2 CLIENT
# ───────────────────────────────────────────────────────────────

class R2Client:
    def __init__(self):
        try:
            self.s3 = boto3.client(
                "s3",
                endpoint_url=f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com",
                region_name="auto",
                aws_access_key_id=R2_ACCESS_KEY,
                aws_secret_access_key=R2_SECRET_KEY,
                config=Config(signature_version="s3v4")
            )
        except Exception as e:
            print(f"[R2] Warning: R2 not configured: {e}")
            self.s3 = None
        self.bucket = R2_BUCKET_NAME
        self.public_domain = R2_PUBLIC_DOMAIN

    def upload_file(self, local_path: str, key: str):
        if not self.s3:
            return None
        self.s3.upload_file(local_path, self.bucket, key)
        return f"{self.public_domain}/{key}"

    def list_videos(self, prefix: str = "outputs/") -> List[dict]:
        if not self.s3:
            return []
        try:
            response = self.s3.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
            items = []
            for obj in response.get("Contents", []):
                key = obj["Key"]
                if key.endswith(".mp4"):
                    items.append({
                        "key": key,
                        "url": f"{self.public_domain}/{key}",
                        "size": obj["Size"],
                        "last_modified": obj["LastModified"].isoformat()
                    })
            return sorted(items, key=lambda x: x["last_modified"], reverse=True)
        except Exception as e:
            print(f"[R2] List error: {e}")
            return []

r2_client = R2Client()

# ───────────────────────────────────────────────────────────────
# JOB MANAGER & STATE
# ───────────────────────────────────────────────────────────────

class JobManager:
    def __init__(self):
        self.jobs: Dict[str, dict] = {}
        self.active_websockets: List[WebSocket] = []
        self.log_websockets: List[WebSocket] = []
        self.logs: List[dict] = []

    def create_job(self, params: dict) -> str:
        job_id = str(uuid.uuid4())[:8]
        self.jobs[job_id] = {
            "id": job_id,
            "status": "queued",
            "progress": 0,
            "stage": "pending",
            "params": params,
            "created_at": datetime.utcnow().isoformat(),
            "completed_at": None,
            "output_url": None,
            "logs": [],
            "retries": 0,
        }
        return job_id

    def update_job(self, job_id: str, **kwargs):
        if job_id in self.jobs:
            self.jobs[job_id].update(kwargs)

    def add_log(self, level: str, agent: str, message: str):
        entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "level": level,
            "agent": agent,
            "message": message
        }
        self.logs.append(entry)
        if len(self.logs) > 10000:
            self.logs = self.logs[-10000:]
        return entry

job_manager = JobManager()

# ───────────────────────────────────────────────────────────────
# WEBSOCKET MANAGER
# ───────────────────────────────────────────────────────────────

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        disconnected = []
        for conn in self.active_connections:
            try:
                await conn.send_json(message)
            except:
                disconnected.append(conn)
        for conn in disconnected:
            self.disconnect(conn)

ws_manager = ConnectionManager()

# ───────────────────────────────────────────────────────────────
# LAZY AGENT LOADING (no loading at startup!)
# ───────────────────────────────────────────────────────────────

_agents_loaded = False
_orchestrator = None
_agent_lock = asyncio.Lock()

async def load_agents():
    """Lazy load all agents - called on first generation request"""
    global _agents_loaded, _orchestrator

    async with _agent_lock:
        if _agents_loaded:
            return

        print("[INFO] Starting lazy agent loading...")
        print(f"[INFO] CUDA available: {torch.cuda.is_available()}")
        print(f"[INFO] GPU count: {torch.cuda.device_count()}")

        sys.path.insert(0, "/workspace")
        try:
            from agents.orchestrator import OrchestratorAgent
            _orchestrator = OrchestratorAgent(job_manager, ws_manager, r2_client, GPU_MAP)
            _agents_loaded = True
            print("[INFO] All agents loaded successfully")
        except Exception as e:
            print(f"[ERROR] Failed to load agents: {e}")
            import traceback
            traceback.print_exc()
            _agents_loaded = False

# ───────────────────────────────────────────────────────────────
# PIPELINE EXECUTOR
# ───────────────────────────────────────────────────────────────

async def run_pipeline(job_id: str):
    """Execute the full video generation pipeline"""
    job = job_manager.jobs[job_id]

    try:
        await load_agents()
        if _orchestrator is None:
            raise RuntimeError("Orchestrator not available")

        await _orchestrator.execute(job_id, job["params"])

    except Exception as e:
        job_manager.update_job(job_id, status="failed", error=str(e))
        job_manager.add_log("ERROR", "ORCHESTRATOR", f"Job {job_id} failed: {e}")
        await ws_manager.broadcast({
            "type": "job_failed",
            "job_id": job_id,
            "error": str(e)
        })

# ───────────────────────────────────────────────────────────────
# FASTAPI APP
# ───────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("="*60)
    print("[INFO] AI Video Swarm starting up...")
    print("[INFO] Models will be loaded LAZILY on first request")
    print("[INFO] This prevents startup timeout")
    print("="*60)
    yield
    print("[INFO] AI Video Swarm shutting down...")

app = FastAPI(title="AI Video Swarm", version="1.0.0", lifespan=lifespan)

# ───────────────────────────────────────────────────────────────
# HTML ROUTES
# ───────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def root():
    with open("/workspace/pass.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

@app.get("/app", response_class=HTMLResponse)
async def main_app():
    with open("/workspace/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

@app.get("/logs", response_class=HTMLResponse)
async def logs_page():
    with open("/workspace/logs.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

# ───────────────────────────────────────────────────────────────
# AUTH ENDPOINTS
# ───────────────────────────────────────────────────────────────

@app.post("/api/auth/login")
async def login(data: dict):
    key = data.get("password", "")
    if not verify_pass_key(key):
        raise HTTPException(status_code=401, detail="Invalid password")

    token = create_access_token({"sub": "user", "auth": True})
    return {"access_token": token, "token_type": "bearer"}

@app.post("/api/auth/verify")
async def verify_auth(payload: dict = Depends(verify_token)):
    return {"valid": True, "user": payload.get("sub")}

# ───────────────────────────────────────────────────────────────
# SETTINGS ENDPOINTS
# ───────────────────────────────────────────────────────────────

_user_settings = {}

@app.get("/api/settings")
async def get_settings(payload: dict = Depends(verify_token)):
    return _user_settings

@app.post("/api/settings")
async def save_settings(data: dict, payload: dict = Depends(verify_token)):
    global _user_settings
    _user_settings.update(data)
    return {"status": "saved", "settings": _user_settings}

# ───────────────────────────────────────────────────────────────
# GENERATION ENDPOINTS
# ───────────────────────────────────────────────────────────────

@app.post("/api/generate")
async def generate_video(data: dict, payload: dict = Depends(verify_token)):
    job_id = job_manager.create_job(data)
    job_manager.add_log("TASK", "ORCHESTRATOR", f"Starting video generation job #{job_id}")

    # Start pipeline in background
    asyncio.create_task(run_pipeline(job_id))

    return {"job_id": job_id, "status": "queued"}

@app.get("/api/status/{job_id}")
async def get_status(job_id: str, payload: dict = Depends(verify_token)):
    if job_id not in job_manager.jobs:
        raise HTTPException(status_code=404, detail="Job not found")
    return job_manager.jobs[job_id]

# ───────────────────────────────────────────────────────────────
# GALLERY ENDPOINTS
# ───────────────────────────────────────────────────────────────

@app.get("/api/videos")
async def list_videos(payload: dict = Depends(verify_token)):
    try:
        return r2_client.list_videos()
    except Exception as e:
        return []

@app.get("/api/videos/{video_id}/download")
async def get_download_url(video_id: str, payload: dict = Depends(verify_token)):
    return {"url": f"{R2_PUBLIC_DOMAIN}/outputs/{video_id}"}

# ───────────────────────────────────────────────────────────────
# LOGS ENDPOINTS
# ───────────────────────────────────────────────────────────────

@app.get("/api/logs")
async def get_logs(limit: int = 100, offset: int = 0, payload: dict = Depends(verify_token)):
    logs = job_manager.logs[offset:offset+limit]
    return {"logs": logs, "total": len(job_manager.logs)}

@app.post("/api/logs/clear")
async def clear_logs(payload: dict = Depends(verify_token)):
    job_manager.logs = []
    return {"status": "cleared"}

# ───────────────────────────────────────────────────────────────
# WEBSOCKET ENDPOINTS
# ───────────────────────────────────────────────────────────────

@app.websocket("/ws/progress")
async def ws_progress(websocket: WebSocket):
    await ws_manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            pass
    except WebSocketDisconnect:
        ws_manager.disconnect(websocket)

@app.websocket("/ws/logs")
async def ws_logs(websocket: WebSocket):
    await websocket.accept()
    job_manager.log_websockets.append(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            if data == "ping":
                await websocket.send_text("pong")
    except WebSocketDisconnect:
        if websocket in job_manager.log_websockets:
            job_manager.log_websockets.remove(websocket)

# ───────────────────────────────────────────────────────────────
# HEALTH & INFO
# ───────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "cuda": torch.cuda.is_available(),
        "gpus": torch.cuda.device_count(),
        "agents_loaded": _agents_loaded
    }

@app.get("/api/info")
async def info():
    return {
        "name": "AI Video Swarm",
        "version": "1.0.0",
        "hardware": "8xL40S",
        "models": {
            "wan_2.6": "Wan-AI/Wan2.6-I2V-14B-720P",
            "ltx": "Lightricks/LTX-2.3",
            "flux": "black-forest-labs/FLUX.1-dev",
        },
        "agents_loaded": _agents_loaded
    }

# ───────────────────────────────────────────────────────────────
# MAIN
# ───────────────────────────────────────────────────────────────

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