"""
HarmoniFind Gradio Backend
Deploy this to Hugging Face Spaces or run locally

To deploy to Hugging Face Spaces:
1. Create account at https://huggingface.co
2. Create new Space with Gradio SDK
3. Upload this file as app.py
4. Upload your embeddings and index files
5. Copy the Space URL to your frontend config
"""

import gradio as gr
import numpy as np
from sentence_transformers import SentenceTransformer
import json

# ============= CONFIGURATION =============
MODEL_NAME = 'all-mpnet-base-v2'
MIN_SIMILARITY_THRESHOLD = 0.5
TOP_K_RESULTS = 10

# ============= LOAD YOUR DATA HERE =============
# TODO: Replace these with your actual data loading
# Load your FAISS index, song metadata, and embeddings

# Example structure for your song metadata:
# songs_metadata = [
#     {
#         "title": "Song Title",
#         "artist": "Artist Name", 
#         "lyrics": "Full lyrics text...",
#         "spotify_url": "https://open.spotify.com/track/..."
#     },
#     ...
# ]

# For demo purposes, using mock data:
songs_metadata = [
    {
        "title": "Rise Up",
        "artist": "Andra Day",
        "lyrics": "You're broken down and tired, Of living life on a merry go round...",
        "spotify_url": "https://open.spotify.com/track/4fSlpKIGm3xa5Q0h7r0qVL"
    },
    {
        "title": "Stronger",
        "artist": "Kelly Clarkson",
        "lyrics": "What doesn't kill you makes you stronger...",
        "spotify_url": "https://open.spotify.com/track/0WqIKmW4BTrj3eJFmnCKMv"
    },
    {
        "title": "Fight Song",
        "artist": "Rachel Platten",
        "lyrics": "This is my fight song, Take back my life song...",
        "spotify_url": "https://open.spotify.com/track/4PXLm6TfjCvQsn9PkW78eN"
    }
]

# ============= INITIALIZE MODEL =============
print("Loading sentence transformer model...")
model = SentenceTransformer(MODEL_NAME)

# TODO: Load your FAISS index
# import faiss
# index = faiss.read_index('path_to_your_index.faiss')

# TODO: Load your precomputed embeddings
# with open('song_embeddings.npy', 'rb') as f:
#     song_embeddings = np.load(f)

# For demo: compute embeddings on the fly (replace with your pre-computed ones)
print("Computing embeddings for demo songs...")
demo_lyrics = [song['lyrics'] for song in songs_metadata]
song_embeddings = model.encode(demo_lyrics, convert_to_numpy=True)
song_embeddings = song_embeddings / np.linalg.norm(song_embeddings, axis=1, keepdims=True)

# ============= SEARCH FUNCTION =============
def search_songs(query: str) -> list:
    """
    Perform semantic search on song lyrics
    
    Args:
        query: Text description of desired song characteristics
        
    Returns:
        List of matching songs with similarity scores
    """
    if not query or not query.strip():
        return []
    
    try:
        # Encode the query
        query_embedding = model.encode([query], convert_to_numpy=True)
        query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
        
        # Compute similarities (cosine similarity since embeddings are normalized)
        similarities = np.dot(song_embeddings, query_embedding.T).flatten()
        
        # Get top K indices sorted by similarity
        top_indices = np.argsort(similarities)[::-1][:TOP_K_RESULTS]
        
        # Format results and filter by threshold
        results = []
        for idx in top_indices:
            similarity = float(similarities[idx])
            
            # Only include results above threshold
            if similarity >= MIN_SIMILARITY_THRESHOLD:
                song = songs_metadata[idx]
                results.append({
                    "title": song['title'],
                    "artist": song['artist'],
                    "similarity": similarity,
                    "spotifyUrl": song.get('spotify_url', '')
                })
        
        return results
        
    except Exception as e:
        print(f"Error during search: {str(e)}")
        return []

# ============= GRADIO INTERFACE =============
def gradio_search(query: str) -> str:
    """Wrapper function for Gradio that returns JSON string"""
    results = search_songs(query)
    return json.dumps(results, indent=2)

# Create Gradio interface
demo = gr.Interface(
    fn=gradio_search,
    inputs=gr.Textbox(
        label="Search Query",
        placeholder="Describe the song you're looking for...",
        lines=3
    ),
    outputs=gr.JSON(label="Search Results"),
    title="🎵 HarmoniFind Backend",
    description=f"""
    Semantic music search powered by {MODEL_NAME} embeddings.
    
    **Current Settings:**
    - Minimum Similarity: {MIN_SIMILARITY_THRESHOLD * 100}%
    - Top Results: {TOP_K_RESULTS}
    - Dataset: {len(songs_metadata)} songs
    
    Enter a description of lyrical themes, emotions, or narratives to find matching songs.
    """,
    examples=[
        ["Uplifting song about overcoming personal challenges"],
        ["Melancholic love song with introspective narrative"],
        ["Energetic anthem about friendship and loyalty"],
        ["Reflective song about life transitions and growth"],
    ],
    api_name="predict"  # This creates the /api/predict endpoint
)

# ============= LAUNCH =============
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Set to True for temporary public link
    )