import os
import io
import tempfile
import numpy as np
import librosa
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import tensorflow as tf

app = FastAPI(title="AI Music Mood Categorizer")

# 1. Define the 8 target moods matching our Flutter app
MOODS = ["gym", "travel", "chill", "sleep", "love", "sad", "party", "study"]

# 2. Load your TensorFlow Model (Placeholder)
# In production, you would uncomment this and load your trained model:
# model = tf.keras.models.load_model("mood_classifier.h5")

def extract_features(audio_path):
    """
    Extracts audio features using librosa.
    This takes the 15-second audio chunk sent by the Flutter app.
    """
    # Load audio with librosa
    # sr=22050 is the default sample rate.
    y, sr = librosa.load(audio_path, sr=22050)
    
    # Extract features
    # MFCCs (Mel-frequency cepstral coefficients) describe the 'shape' of the sound
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    mfcc_mean = np.mean(mfcc.T, axis=0)
    
    # Chroma measures pitch/harmonic content (useful for 'sad' vs 'happy' chords)
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    chroma_mean = np.mean(chroma.T, axis=0)
    
    # Spectral contrast (useful for distinguishing energetic music)
    contrast = librosa.feature.spectral_contrast(y=y, sr=sr)
    contrast_mean = np.mean(contrast.T, axis=0)

    # Tempo (BPM)
    tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
    
    # Combine all features into a single numpy array
    # This is what your ML model will take as input for prediction
    features = np.hstack([mfcc_mean, chroma_mean, contrast_mean, tempo])
    return features, tempo[0]

@app.post("/analyze")
async def analyze_audio(file: UploadFile = File(...)):
    """
    Endpoint that receives the 15s audio chunk from Flutter, 
    extracts features, runs the AI model, and returns the mood.
    """
    try:
        # 1. Save the uploaded audio chunk to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_audio:
            content = await file.read()
            temp_audio.write(content)
            temp_audio_path = temp_audio.name

        # 2. Extract AI Features using Librosa
        features, bpm = extract_features(temp_audio_path)
        
        # 3. Predict Mood using TensorFlow (Mock logic)
        # In production, you would run:
        # prediction = model.predict(np.array([features]))
        # predicted_index = np.argmax(prediction)
        # predicted_mood = MOODS[predicted_index]
        
        # --- MOCK HEURISTIC FOR TESTING ---
        # Since you haven't trained the TFLite/TF model yet, 
        # this heuristic uses the BPM (tempo) to mock a prediction 
        # so you can test the Flutter app immediately!
        predicted_mood = "chill"
        if bpm > 140:
            predicted_mood = "gym" # Aggressive, fast
        elif bpm > 120:
            predicted_mood = "party" # Upbeat
        elif bpm > 100:
            predicted_mood = "travel" # Moving
        elif bpm < 70:
            predicted_mood = "sleep" # Very slow, ambient
        elif bpm < 90:
            predicted_mood = "sad" # Slow, melancholic
        # ----------------------------------

        # Clean up temp file
        os.remove(temp_audio_path)

        return JSONResponse(content={
            "mood": predicted_mood,
            "bpm": float(bpm),
            "status": "success"
        })

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": str(e)}
        )

# To run this server locally for testing:
# pip install -r requirements.txt
# uvicorn main:app --host 0.0.0.0 --port 8000