import os
import time
import shutil
import uuid
import json
import asyncio
import base64
from typing import List, Optional
from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
import cv2
import numpy as np

# Configuration
GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=GEMINI_API_KEY)

app = FastAPI(title="BJJ AI Coach - Progressive Analysis")

# Enable CORS for frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- MODELS ---
class TimestampedEvent(BaseModel):
    time: str
    title: str
    description: str

class Drill(BaseModel):
    name: str
    focus_area: str
    reason: str

class SkillBreakdown(BaseModel):
    offense_score: int
    defense_score: int
    positioning_score: int
    stamina_score: int

class AnalysisResult(BaseModel):
    skill_breakdown: SkillBreakdown
    strengths: List[str]
    weaknesses: List[str]
    missed_opportunities: List[TimestampedEvent]
    key_moments: List[TimestampedEvent]
    coach_notes: str
    recommended_drills: List[Drill]

db_storage = {}

# --- DYNAMIC FRAME EXTRACTION ---

def calculate_optimal_frames(duration: float) -> int:
    """
    Dynamically calculate frame count based on video length
    
    Logic:
    - 0-15s: 6 frames (every 2.5s)
    - 15-30s: 8 frames (every 3.75s)
    - 30-60s: 12 frames (every 5s)
    - 60-120s: 16 frames (every 7.5s)
    - 120s+: 20 frames (every 6s)
    """
    if duration <= 15:
        return 6
    elif duration <= 30:
        return 8
    elif duration <= 60:
        return 12
    elif duration <= 120:
        return 16
    else:
        return 20

def extract_key_frames_with_metadata(video_path: str, max_frames: int = None) -> tuple:
    """
    Extract frames with full metadata for frontend display
    
    Returns:
    - frames: List of (frame_bytes, timestamp, frame_number)
    - metadata: Dict with video info
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # Dynamic frame calculation
    if max_frames is None:
        max_frames = calculate_optimal_frames(duration)
    
    metadata = {
        "duration": round(duration, 2),
        "fps": round(fps, 2),
        "total_frames": total_frames,
        "resolution": f"{width}x{height}",
        "frames_to_extract": max_frames
    }
    
    frames = []
    interval = max(1, total_frames // max_frames)
    
    prev_frame = None
    frame_idx = 0
    extracted_count = 0
    
    while cap.isOpened() and extracted_count < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        
        if frame_idx % interval == 0:
            should_capture = False
            
            if prev_frame is not None:
                diff = cv2.absdiff(frame, prev_frame)
                score = np.mean(diff)
                should_capture = score > 15 or extracted_count == 0
            else:
                should_capture = True  
            
            if should_capture:
                small_frame = cv2.resize(frame, (640, 480))
                _, buffer = cv2.imencode('.jpg', small_frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
                frame_bytes = buffer.tobytes()
                
                timestamp_sec = frame_idx / fps
                timestamp_str = f"{int(timestamp_sec // 60):02d}:{int(timestamp_sec % 60):02d}"
                
                frames.append({
                    "bytes": frame_bytes,
                    "timestamp": timestamp_str,
                    "frame_number": frame_idx,
                    "second": round(timestamp_sec, 2)
                })
                
                extracted_count += 1
            
            prev_frame = frame.copy()
        
        frame_idx += 1
    
    cap.release()
    
    return frames, metadata

# --- ANALYSIS PROMPT ---

ANALYSIS_PROMPT = """
You are analyzing {num_frames} key frames from a BJJ roll ({duration}s total).

## CONTEXT
- User: {user_desc}
- Opponent: {opp_desc}
- Activity: {activity_type}

## FRAMES PROVIDED
{frame_timestamps}

## TASK
Analyze these snapshots to reconstruct the tactical story.

### VALIDATION
Does this show {activity_type}?
- If NO → Return: {{"error": "MISMATCH", "detected_activity": "..."}}
- If YES → Continue

### ANALYSIS
For each frame:
1. Identify position (Guard/Mount/Side Control/Back/Standing)
2. Determine who has advantage
3. Note immediate threats (chokes, armbars, etc.)

Based on position changes:
- What transitions occurred?
- Aggression level of each player
- Submission attempts
- Defensive responses

### SCORING (0-100)
- **Offense**: Aggression + submission attempts + sweeps
- **Defense**: Escape quality + threat recognition
- **Positioning**: Dominant position time % (estimate)
- **Stamina**: Movement quality (early vs late frames)

### RULES
1. Use frame timestamps for events
2. Be direct in coach notes and give suggestions

## OUTPUT (JSON ONLY)
{{
  "skill_breakdown": {{
    "offense_score": <int>,
    "defense_score": <int>,
    "positioning_score": <int>,
    "stamina_score": <int>
  }},
  "strengths": ["<specific>", "<specific>", "<specific>"],
  "weaknesses": ["<actionable>", "<actionable>", "<actionable>"],
  "missed_opportunities": [
    {{"time": "MM:SS", "title": "...", "description": "..."}}
  ],
  "key_moments": [
    {{"time": "MM:SS", "title": "...", "description": "..."}}
  ],
  "coach_notes": "<Direct feedback>",
  "recommended_drills": [
    {{"name": "...", "focus_area": "...", "reason": "..."}}
  ]
}}
"""

async def analyze_with_frames(
    video_path: str,
    user_desc: str,
    opp_desc: str,
    activity_type: str,
    analysis_id: str = None
) -> AnalysisResult:
    """Analysis with progress tracking"""
    
    try:

        if analysis_id:
            db_storage[analysis_id]["status"] = "extracting_frames"
            db_storage[analysis_id]["progress"] = 0
        
        start = time.time()
        frames, metadata = await asyncio.get_event_loop().run_in_executor(
            None, extract_key_frames_with_metadata, video_path, None
        )
        
        print(f"Extracted {len(frames)} frames in {time.time() - start:.2f}s")
        
        if not frames:
            raise Exception("No frames extracted")
        
        if analysis_id:
            db_storage[analysis_id]["frames"] = [
                {
                    "timestamp": f["timestamp"],
                    "frame_number": f["frame_number"],
                    "second": f["second"],
                    "image": base64.b64encode(f["bytes"]).decode('utf-8')
                }
                for f in frames
            ]
            db_storage[analysis_id]["metadata"] = metadata
            db_storage[analysis_id]["status"] = "analyzing"
            db_storage[analysis_id]["progress"] = 50
        

        frame_timestamps = "\n".join([
            f"Frame {i+1}: {f['timestamp']} ({f['second']}s)"
            for i, f in enumerate(frames)
        ])
        
        prompt = ANALYSIS_PROMPT.format(
            num_frames=len(frames),
            duration=metadata['duration'],
            user_desc=user_desc,
            opp_desc=opp_desc,
            activity_type=activity_type,
            frame_timestamps=frame_timestamps
        )
        
        content = []
        for frame_data in frames:
            content.append({
                "mime_type": "image/jpeg",
                "data": base64.b64encode(frame_data["bytes"]).decode('utf-8')
            })
        content.append(prompt)
        
        # Call Gemini
        start = time.time()
        model = genai.GenerativeModel(
            model_name="gemini-2.5-flash",
            generation_config={
                "temperature": 0.4,
                "response_mime_type": "application/json"
            }
        )
        
        response = await asyncio.get_event_loop().run_in_executor(
            None,
            lambda: model.generate_content(
                content,
                safety_settings={
                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
                }
            )
        )
        
        print(f"Gemini inference: {time.time() - start:.2f}s")
        
        result_json = json.loads(response.text)
        
        if analysis_id:
            db_storage[analysis_id]["progress"] = 100
        
        return AnalysisResult(**result_json)
        
    except Exception as e:
        raise Exception(f"Analysis failed: {str(e)}")

# --- STREAMING ENDPOINT FOR PROGRESSIVE DISPLAY ---

@app.get("/stream-frames/{analysis_id}")
async def stream_frames(analysis_id: str):
    """
    Server-Sent Events stream for progressive frame display
    
    Frontend receives frames one by one, then final result
    """
    async def event_generator():
        max_wait = 30
        waited = 0
        while analysis_id not in db_storage or "frames" not in db_storage[analysis_id]:
            await asyncio.sleep(0.5)
            waited += 0.5
            if waited > max_wait:
                yield f"data: {json.dumps({'error': 'Timeout waiting for frames'})}\n\n"
                return
        
        frames = db_storage[analysis_id]["frames"]
        metadata = db_storage[analysis_id]["metadata"]
        
        yield f"data: {json.dumps({'type': 'metadata', 'data': metadata})}\n\n"
        
        for i, frame in enumerate(frames):
            frame_data = {
                "type": "frame",
                "index": i,
                "total": len(frames),
                "timestamp": frame["timestamp"],
                "image": frame["image"]
            }
            yield f"data: {json.dumps(frame_data)}\n\n"
            await asyncio.sleep(0.1)  # Small delay for smooth streaming
        
        # Wait for analysis to complete
        while db_storage[analysis_id]["status"] != "completed":
            if db_storage[analysis_id]["status"] == "failed":
                yield f"data: {json.dumps({'type': 'error', 'message': db_storage[analysis_id].get('error', 'Unknown error')})}\n\n"
                return
            await asyncio.sleep(0.5)
        
        # Send final result
        result_data = {
            "type": "result",
            "data": db_storage[analysis_id]["data"]
        }
        yield f"data: {json.dumps(result_data)}\n\n"
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

# --- BACKGROUND TASK ---

async def analyze_video_task(
    analysis_id: str,
    video_path: str,
    user_desc: str,
    opp_desc: str,
    activity_type: str
):
    """Background processing with progress tracking"""
    try:
        db_storage[analysis_id]["status"] = "processing"
        
        result = await analyze_with_frames(
            video_path, user_desc, opp_desc, activity_type, analysis_id
        )
        
        db_storage[analysis_id]["status"] = "completed"
        db_storage[analysis_id]["data"] = result.model_dump()
        
    except Exception as e:
        db_storage[analysis_id]["status"] = "failed"
        db_storage[analysis_id]["error"] = str(e)
    finally:
        try:
            os.remove(video_path)
        except:
            pass

# --- API ENDPOINTS ---

@app.post("/upload")
async def upload_video(file: UploadFile = File(...)):
    """Upload endpoint"""
    file_name = f"{uuid.uuid4()}_{file.filename}"
    file_path = f"temp_videos/{file_name}"
    os.makedirs("temp_videos", exist_ok=True)
    
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    return {"file_name": file_path}

@app.post("/analyze")
async def start_analysis(
    video_file_name: str,
    user_description: str = "User in blue gi",
    opponent_description: str = "Opponent in white gi",
    activity_type: str = "Brazilian Jiu-Jitsu",
    background_tasks: BackgroundTasks = None
):
    """
    Start analysis with progressive frame streaming
    
    Returns analysis_id for streaming endpoint
    """
    analysis_id = str(uuid.uuid4())
    db_storage[analysis_id] = {
        "status": "queued",
        "progress": 0
    }
    
    background_tasks.add_task(
        analyze_video_task,
        analysis_id,
        video_file_name,
        user_description,
        opponent_description,
        activity_type
    )
    
    return {
        "analysis_id": analysis_id,
        "stream_url": f"/stream-frames/{analysis_id}"
    }

@app.get("/status/{analysis_id}")
async def get_status(analysis_id: str):
    """Check progress"""
    if analysis_id not in db_storage:
        raise HTTPException(status_code=404, detail="Not found")
    return db_storage[analysis_id]

@app.post("/analyze-complete")
async def analyze_complete(
    file: UploadFile = File(...),
    user_description: str = "User",
    opponent_description: str = "Opponent",
    activity_type: str = "Brazilian Jiu-Jitsu"
):
    """Synchronous analysis with frames"""
    start_time = time.time()
    
   
    file_name = f"{uuid.uuid4()}_{file.filename}"
    file_path = f"temp_videos/{file_name}"
    os.makedirs("temp_videos", exist_ok=True)
    
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    analysis_id = str(uuid.uuid4())
    db_storage[analysis_id] = {"status": "processing"}
    
    try:
        result = await analyze_with_frames(
            file_path, user_description, opponent_description, activity_type, analysis_id
        )
        
        total_time = time.time() - start_time
        
        return {
            "status": "completed",
            "data": result.model_dump(),
            "frames": db_storage[analysis_id].get("frames", []),
            "metadata": db_storage[analysis_id].get("metadata", {}),
            "processing_time": f"{total_time:.2f}s"
        }
    
    except Exception as e:
        return {"status": "failed", "error": str(e)}
    
    finally:
        try:
            os.remove(file_path)
        except:
            pass

@app.get("/")
async def root():
    return {
        "message": "BJJ AI Coach - Progressive Frame Streaming",
        "version": "5.0.0",
        "features": [
            "Dynamic frame extraction (6-20 frames based on duration)",
            "Progressive frame display via SSE",
            "Real-time analysis status"
        ]
    }

# Run with: uvicorn main:app --reload