import os
import uuid
import asyncio
import psutil
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, Depends, Security
from fastapi.responses import JSONResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import cv2
import json

from pipeline import init_models, process_image_pipeline

app = FastAPI(title="AuraLens API")

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

# Authentication setup
security = HTTPBearer()

def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
    expected_token = os.environ.get("HF_READ_TOKEN")
    # If no token is set in env, we allow it (for local testing), otherwise enforce
    if expected_token and credentials.credentials != expected_token:
        raise HTTPException(status_code=401, detail="Invalid authorization token")
    return credentials.credentials

# In-memory storage for jobs and their statuses
jobs = {}
job_results = {}

# Max RAM threshold (15 GB) in bytes
MAX_RAM_BYTES = 15.0 * 1024 * 1024 * 1024

# Setup Queue
task_queue = asyncio.Queue()

async def worker():
    while True:
        job = await task_queue.get()
        job_id = job['job_id']
        image_bytes = job['image_bytes']
        mode = job['mode']
        
        try:
            # Memory check before processing
            mem = psutil.virtual_memory()
            if mem.used > MAX_RAM_BYTES:
                await update_job_status(job_id, 0, "Server Busy: Memory near limits, waiting...")
                await asyncio.sleep(5)
                # Re-queue
                await task_queue.put(job)
                task_queue.task_done()
                continue
                
            # Process
            result_img = await process_image_pipeline(job_id, image_bytes, update_job_status, mode)
            
            # Encode result
            is_success, buffer = cv2.imencode(".jpg", result_img, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
            if is_success:
                job_results[job_id] = buffer.tobytes()
                await update_job_status(job_id, 100, "Done", status="completed")
            else:
                await update_job_status(job_id, 100, "Failed to encode image", status="failed")
                
        except Exception as e:
            print(f"Job {job_id} failed: {e}")
            await update_job_status(job_id, 100, f"Error: {str(e)}", status="failed")
        finally:
            task_queue.task_done()

@app.on_event("startup")
async def startup_event():
    # Initialize models
    init_models()
    # Start background worker
    asyncio.create_task(worker())
    
    # Mount static files here to avoid path issues during setup
    os.makedirs("static", exist_ok=True)
    app.mount("/", StaticFiles(directory="static", html=True), name="static")

async def update_job_status(job_id: str, progress: int, message: str, status: str = "processing"):
    if job_id not in jobs:
        jobs[job_id] = {"queues": []}
        
    jobs[job_id]["progress"] = progress
    jobs[job_id]["message"] = message
    jobs[job_id]["status"] = status
    
    # Notify all listeners
    data = json.dumps({"progress": progress, "message": message, "status": status})
    for q in jobs[job_id]["queues"]:
        await q.put(data)

@app.post("/api/process")
async def process_request(
    file: UploadFile = File(...), 
    mode: str = Form("Portrait HD"),
    token: str = Depends(verify_token)
):
    mem = psutil.virtual_memory()
    if mem.used > MAX_RAM_BYTES:
        return JSONResponse(status_code=503, content={"error": "Server Busy: Processing queued, please wait 10 seconds..."})
        
    image_bytes = await file.read()
    job_id = str(uuid.uuid4())
    
    jobs[job_id] = {
        "progress": 0,
        "message": "Queued...",
        "status": "queued",
        "queues": []
    }
    
    # Enqueue task
    await task_queue.put({
        "job_id": job_id,
        "image_bytes": image_bytes,
        "mode": mode
    })
    
    return {"job_id": job_id}

@app.get("/api/status/{job_id}")
async def get_status(job_id: str, request: Request, token: str = Depends(verify_token)):
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")
        
    async def event_generator():
        q = asyncio.Queue()
        jobs[job_id]["queues"].append(q)
        try:
            # Send current state first
            initial_data = json.dumps({
                "progress": jobs[job_id]["progress"],
                "message": jobs[job_id]["message"],
                "status": jobs[job_id]["status"]
            })
            yield f"data: {initial_data}\n\n"
            
            while True:
                if await request.is_disconnected():
                    break
                data = await q.get()
                yield f"data: {data}\n\n"
                state = json.loads(data)
                if state["status"] in ["completed", "failed"]:
                    break
        finally:
            if q in jobs[job_id]["queues"]:
                jobs[job_id]["queues"].remove(q)
                
    return StreamingResponse(event_generator(), media_type="text/event-stream")

@app.get("/api/result/{job_id}")
async def get_result(job_id: str, token: str = Depends(verify_token)):
    if job_id not in job_results:
        raise HTTPException(status_code=404, detail="Result not found or not ready")
    return Response(content=job_results[job_id], media_type="image/jpeg")