from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import mediapipe as mp
import cv2
import numpy as np
from PIL import Image
import io
import requests
import base64
import json
import os
from pathlib import Path
from typing import Optional, List, Dict
from motion_handler import MotionHandler
from joint_mapper import JointMapper, create_mapper_from_motion

app = FastAPI()

# Setup cache directory
CACHE_DIR = Path("./cache")
CACHE_DIR.mkdir(exist_ok=True)

# Enable CORS for Next.js frontend

# Allow localhost and production frontend
allowed_origins = [
    "http://localhost:3000",
    "https://dancify.suno.run",
]

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

# Initialize MediaPipe Pose
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(
    static_image_mode=True,
    model_complexity=2,
    enable_segmentation=False,
    min_detection_confidence=0.5,
)

# Store reference image landmarks
reference_landmarks = None

# Store motion data
motion_handler = MotionHandler(fps=30)
joint_mapper: Optional[JointMapper] = None


def extract_landmarks(image_data: bytes) -> Optional[List[Dict[str, float]]]:
    """Extract pose landmarks from image bytes using MediaPipe."""
    try:
        # Convert bytes to image
        image = Image.open(io.BytesIO(image_data))
        image_np = np.array(image)

        # Convert to RGB if needed
        if len(image_np.shape) == 2:  # Grayscale
            image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
        elif image_np.shape[2] == 4:  # RGBA
            image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)

        # Process with MediaPipe
        results = pose.process(image_np)

        if results.pose_landmarks:
            landmarks = []
            for landmark in results.pose_landmarks.landmark:
                landmarks.append(
                    {
                        "x": landmark.x,
                        "y": landmark.y,
                        "z": landmark.z,
                        "visibility": landmark.visibility,
                    }
                )
            return landmarks
        return None
    except Exception as e:
        print(f"Error extracting landmarks: {e}")
        return None


def calculate_similarity_score(
    ref_landmarks: List[Dict[str, float]], curr_landmarks: List[Dict[str, float]]
) -> float:
    """
    Calculate similarity score between two sets of landmarks.
    Returns a score from 0-100, where 100 is a perfect match.
    """
    if not ref_landmarks or not curr_landmarks:
        return 0.0

    # Focus on key body points (excluding face landmarks for better body pose comparison)
    # MediaPipe Pose has 33 landmarks, we'll focus on body landmarks (indices 11-28)
    key_indices = list(range(11, 29))  # Shoulders, elbows, wrists, hips, knees, ankles

    total_distance = 0.0
    valid_points = 0

    for idx in key_indices:
        if idx >= len(ref_landmarks) or idx >= len(curr_landmarks):
            continue

        ref_lm = ref_landmarks[idx]
        curr_lm = curr_landmarks[idx]

        # Only compare if both landmarks are visible enough
        if ref_lm["visibility"] < 0.5 or curr_lm["visibility"] < 0.5:
            continue

        # Calculate Euclidean distance (using x, y coordinates)
        # We normalize by using only x and y since z can be less reliable
        distance = np.sqrt(
            (ref_lm["x"] - curr_lm["x"]) ** 2 + (ref_lm["y"] - curr_lm["y"]) ** 2
        )

        total_distance += distance
        valid_points += 1

    if valid_points == 0:
        return 0.0

    # Calculate average distance
    avg_distance = total_distance / valid_points

    # Convert distance to similarity score
    # Distance of 0 = 100 score, distance of 0.5 or more = 0 score
    # This is a heuristic that can be tuned
    similarity = max(0, 100 * (1 - (avg_distance / 0.5)))

    return round(similarity, 2)


def get_cache_key(song_id: str, start_time: float, end_time: float) -> str:
    """Generate a cache key for the given parameters."""
    return f"{song_id}_{start_time}_{end_time}"


def get_cached_motion(song_id: str, start_time: float, end_time: float) -> Optional[Dict]:
    """
    Check if cached motion data exists and return it.
    Returns None if cache doesn't exist.
    """
    cache_key = get_cache_key(song_id, start_time, end_time)
    pkl_path = CACHE_DIR / f"{cache_key}.pkl"
    metadata_path = CACHE_DIR / f"{cache_key}_metadata.json"
    
    if not pkl_path.exists() or not metadata_path.exists():
        return None
    
    try:
        # Load pkl file
        with open(pkl_path, 'rb') as f:
            pkl_contents = f.read()
        
        # Load metadata
        with open(metadata_path, 'r') as f:
            cached_data = json.load(f)
        
        print(f"✅ Found cached motion data for {cache_key}")
        return {
            "pkl_contents": pkl_contents,
            "cached_data": cached_data
        }
    except Exception as e:
        print(f"⚠️  Error loading cache: {e}")
        return None


def save_motion_to_cache(
    song_id: str, 
    start_time: float, 
    end_time: float,
    pkl_contents: bytes,
    metadata: Dict,
    frames: List
):
    """Save motion data to cache."""
    cache_key = get_cache_key(song_id, start_time, end_time)
    pkl_path = CACHE_DIR / f"{cache_key}.pkl"
    metadata_path = CACHE_DIR / f"{cache_key}_metadata.json"
    
    try:
        # Save pkl file
        with open(pkl_path, 'wb') as f:
            f.write(pkl_contents)
        
        # Save metadata and basic info (not all frames, too large)
        cache_data = {
            "metadata": metadata,
            "num_frames": len(frames)
        }
        with open(metadata_path, 'w') as f:
            json.dump(cache_data, f)
        
        print(f"💾 Saved motion data to cache: {cache_key}")
    except Exception as e:
        print(f"⚠️  Error saving to cache: {e}")


@app.get("/")
async def root():
    return {"message": "Dancify Pose Comparison API", "status": "running"}


@app.post("/api/set-reference")
async def set_reference(file: UploadFile = File(...)):
    """Set the reference pose image."""
    global reference_landmarks

    try:
        # Read the uploaded file
        contents = await file.read()

        # Extract landmarks
        landmarks = extract_landmarks(contents)

        if landmarks is None:
            raise HTTPException(
                status_code=400,
                detail="Could not detect pose in the reference image. Please ensure a person is visible.",
            )

        reference_landmarks = landmarks

        return {
            "message": "Reference image set successfully",
            "landmarks_detected": len(landmarks),
            "landmarks": landmarks,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/compare-pose")
async def compare_pose(file: UploadFile = File(...)):
    """Compare current pose to reference pose."""
    global reference_landmarks

    if reference_landmarks is None:
        raise HTTPException(
            status_code=400,
            detail="No reference image set. Please upload a reference image first.",
        )

    try:
        # Read the uploaded file
        contents = await file.read()

        # Extract landmarks from current frame
        current_landmarks = extract_landmarks(contents)

        if current_landmarks is None:
            return {
                "score": 0,
                "message": "No pose detected in current frame",
                "landmarks": None,
            }

        # Calculate similarity score
        score = calculate_similarity_score(reference_landmarks, current_landmarks)

        return {
            "score": score,
            "message": "Pose compared successfully!",
            "landmarks_detected": len(current_landmarks),
            "current_landmarks": current_landmarks,
            "reference_landmarks": reference_landmarks,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/load-motion")
async def load_motion(file: UploadFile = File(...)):
    """
    Load motion data from pickle file and return all frames as JSON.
    Frontend will handle animation and pose comparison locally.
    """
    global motion_handler, joint_mapper

    try:
        # Read the uploaded file
        contents = await file.read()

        # Load motion data
        metadata = motion_handler.load_pickle(contents)

        # Create and calibrate joint mapper
        if motion_handler.full_pose is not None:
            joint_mapper = create_mapper_from_motion(motion_handler.full_pose)

            # Convert ALL frames to MediaPipe format and return them
            all_frames = []
            for frame_idx in range(motion_handler.num_frames):
                smpl_joints = motion_handler.get_frame_by_index(frame_idx)
                if smpl_joints is not None:
                    landmarks = joint_mapper.smpl_to_mediapipe_landmarks(smpl_joints)
                    all_frames.append(landmarks)

        return {
            "message": "Motion data loaded successfully",
            "metadata": metadata,
            "frames": all_frames,  # All frames for client-side animation
        }

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error loading motion: {str(e)}")


@app.post("/api/generate-motion")
async def generate_motion(
    song_id: str, start_time: float = 0.0, end_time: float = 30.0
):
    """
    Generate motion data from Modal API using song ID.
    This fetches the dance choreography for a given song.
    Uses local cache to avoid regenerating existing choreography.

    Args:
        song_id: The Suno upload ID (song ID)
        start_time: Start time in seconds (default: 0.0)
        end_time: End time in seconds (default: 30.0)

    Returns:
        Motion metadata and all frames in MediaPipe format
    """
    global motion_handler, joint_mapper

    try:
        print(f"Requesting motion generation for song_id: {song_id}")
        print(f"Time range: {start_time}s to {end_time}s")

        # Check cache first
        cached_motion = get_cached_motion(song_id, start_time, end_time)
        
        if cached_motion:
            print("📦 Using cached motion data!")
            pkl_contents = cached_motion["pkl_contents"]
            
            # Load cached pkl file
            metadata = motion_handler.load_pickle(pkl_contents)

            # Create and calibrate joint mapper
            if motion_handler.full_pose is not None:
                joint_mapper = create_mapper_from_motion(motion_handler.full_pose)

                # Convert ALL frames to MediaPipe format
                all_frames = []
                for frame_idx in range(motion_handler.num_frames):
                    smpl_joints = motion_handler.get_frame_by_index(frame_idx)
                    if smpl_joints is not None:
                        landmarks = joint_mapper.smpl_to_mediapipe_landmarks(smpl_joints)
                        all_frames.append(landmarks)

            return {
                "message": "Motion data loaded from cache",
                "metadata": metadata,
                "frames": all_frames,
                "cached": True,
            }

        # Cache miss - call Modal API
        print("🌐 Cache miss, calling Modal API...")
        modal_response = requests.post(
            "https://suno-ai--popdg-dance-generator-web-generate-dance.modal.run/generate",
            json={
                "upload_id": song_id,
                "start_time": start_time,
                "end_time": end_time,
                "no_render": True,  # We don't need the video, just the pkl
            },
            timeout=600,  # 10 minute timeout for generation
        )
        modal_response.raise_for_status()

        result = modal_response.json()
        print(f"Modal API response status: {result.get('status')}")

        # Extract the pkl file
        motion_files = result["result"]["motion_files"]

        if not motion_files:
            raise HTTPException(
                status_code=500, detail="No motion file generated by Modal API"
            )

        # Get the first (and should be only) motion file
        pkl_base64 = list(motion_files.values())[0]
        pkl_contents = base64.b64decode(pkl_base64)

        print(f"Received pkl file: {len(pkl_contents) / (1024**2):.2f} MB")

        # Load motion data using existing handler
        metadata = motion_handler.load_pickle(pkl_contents)

        # Create and calibrate joint mapper
        if motion_handler.full_pose is not None:
            joint_mapper = create_mapper_from_motion(motion_handler.full_pose)

            # Convert ALL frames to MediaPipe format
            all_frames = []
            for frame_idx in range(motion_handler.num_frames):
                smpl_joints = motion_handler.get_frame_by_index(frame_idx)
                if smpl_joints is not None:
                    landmarks = joint_mapper.smpl_to_mediapipe_landmarks(smpl_joints)
                    all_frames.append(landmarks)
        
        # Save to cache for future use
        save_motion_to_cache(song_id, start_time, end_time, pkl_contents, metadata, all_frames)

        return {
            "message": "Motion data generated successfully",
            "metadata": metadata,
            "frames": all_frames,
            "cached": False,
        }

    except requests.exceptions.RequestException as e:
        raise HTTPException(
            status_code=503,
            detail=f"Error calling Modal API: {str(e)}",
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error generating motion: {str(e)}"
        )


@app.get("/api/get-reference-frame")
async def get_reference_frame(timestamp: float):
    """
    Get reference pose at a specific timestamp.

    Args:
        timestamp: Time in seconds

    Returns:
        MediaPipe-formatted landmarks for the frame at that timestamp
    """
    global motion_handler, joint_mapper

    if not motion_handler.is_loaded():
        raise HTTPException(
            status_code=400,
            detail="No motion data loaded. Please upload a motion file first.",
        )

    if joint_mapper is None:
        raise HTTPException(status_code=500, detail="Joint mapper not initialized")

    try:
        # Get SMPL joint positions at timestamp
        smpl_joints = motion_handler.get_frame_at_timestamp(timestamp)

        if smpl_joints is None:
            raise HTTPException(
                status_code=400, detail=f"Invalid timestamp: {timestamp}"
            )

        # Convert to MediaPipe landmarks
        landmarks = joint_mapper.smpl_to_mediapipe_landmarks(smpl_joints)

        return {
            "timestamp": timestamp,
            "landmarks": landmarks,
            "frame_index": int(timestamp * motion_handler.fps),
        }

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error getting reference frame: {str(e)}"
        )


@app.post("/api/compare-pose-motion")
async def compare_pose_motion(file: UploadFile = File(...), timestamp: float = 0.0):
    """
    Compare current pose to motion data at specific timestamp.
    Used for Just Dance mode.
    """
    global motion_handler, joint_mapper

    if not motion_handler.is_loaded():
        raise HTTPException(
            status_code=400,
            detail="No motion data loaded. Please upload a motion file first.",
        )

    if joint_mapper is None:
        raise HTTPException(status_code=500, detail="Joint mapper not initialized")

    try:
        # Get reference pose from motion data at timestamp
        smpl_joints = motion_handler.get_frame_at_timestamp(timestamp)
        if smpl_joints is None:
            raise HTTPException(
                status_code=400, detail=f"Invalid timestamp: {timestamp}"
            )

        # Convert to MediaPipe landmarks
        reference_landmarks = joint_mapper.smpl_to_mediapipe_landmarks(smpl_joints)

        # Read the uploaded frame
        contents = await file.read()

        # Extract landmarks from current frame
        current_landmarks = extract_landmarks(contents)

        if current_landmarks is None:
            return {
                "score": 0,
                "message": "No pose detected in current frame",
                "landmarks": None,
            }

        # Calculate similarity score
        score = calculate_similarity_score(reference_landmarks, current_landmarks)

        return {
            "score": score,
            "message": "Pose compared successfully",
            "landmarks_detected": len(current_landmarks),
            "current_landmarks": current_landmarks,
            "reference_landmarks": reference_landmarks,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error comparing pose: {str(e)}")


@app.get("/api/motion-status")
async def motion_status():
    """Check if motion data is loaded and get metadata."""
    global motion_handler

    if not motion_handler.is_loaded():
        return {"loaded": False, "message": "No motion data loaded"}

    return {
        "loaded": True,
        "num_frames": motion_handler.num_frames,
        "duration": motion_handler.duration,
        "fps": motion_handler.fps,
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)