import os
import pandas as pd
import base64
import random
import unicodedata
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import JSONResponse
from typing import List, Dict, Optional

app = FastAPI(title="TTS Arena Audio Server")

# --- CONFIGURATION ---
AUDIO_BASE_PATH = os.getenv("AUDIO_BASE_PATH", "/root/tts_arena/audios")
METADATA_FILE = "metadata.csv"  # The single consolidated CSV

# 1. ELEVENLABS Speaker Mapping
ELEVENLABS_GENDER_MAP = {
    "male": ["Adam", "Bill"],
    "female": ["Alice"]
}

# 2. PARLER Speaker Mapping
PARLER_GENDER_MAP = {
    "male": ["Rohit"],
    "female": ["Divya"]
}

# Global Data Cache
# We now only need one dataframe since filenames are consistent
db = {
    "data": pd.DataFrame() 
}

def normalize_text(text):
    """Normalize text to ensure cache hits."""
    return unicodedata.normalize('NFC', str(text).strip().strip('"').strip("'"))

@app.on_event("startup")
def load_data():
    """Loads the single unified metadata CSV into memory."""
    try:
        if os.path.exists(METADATA_FILE):
            df = pd.read_csv(METADATA_FILE)
            
            # Validation
            if 'text' in df.columns and 'filename' in df.columns:
                df['text_clean'] = df['text'].apply(normalize_text)
                db["data"] = df
                print(f"✅ Metadata: Loaded {len(df)} sentence mappings.")
            else:
                print("⚠️ Metadata CSV missing 'text' or 'filename' columns.")
        else:
            print(f"⚠️ '{METADATA_FILE}' not found. Audio lookups will fail.")
            
    except Exception as e:
        print(f"❌ Critical Error Loading CSV: {e}")

def _get_audio_response(
    model_name: str, 
    gender_map: Dict[str, List[str]], 
    sentence: str, 
    gender: Optional[str] = None, 
    name: Optional[str] = None
):
    # 0. Validate Inputs
    if not gender and not name:
        raise HTTPException(status_code=400, detail="You must provide either 'gender' or 'name'.")

    # 1. Lookup Sentence in the Unified DB
    df = db.get("data")
    if df is None or df.empty:
        raise HTTPException(status_code=503, detail="Metadata database not loaded.")

    clean_sent = normalize_text(sentence)
    match = df[df['text_clean'] == clean_sent]
    
    if match.empty:
        raise HTTPException(status_code=404, detail="Sentence not found in index.")
    
    # Extract the consistent filename (e.g., '0_female.wav')
    target_filename = match.iloc[0]['filename']
    
    # 2. Identify Candidate Speakers
    candidate_speakers = []

    if name:
        # Priority: Specific Name
        candidate_speakers = [name]
    elif gender:
        # Fallback: Random speaker from Gender Map
        candidate_speakers = gender_map.get(gender, [])
        if not candidate_speakers:
            raise HTTPException(status_code=400, detail=f"No speakers defined for gender: {gender}")
        random.shuffle(candidate_speakers)
    
    # 3. Search for the file on disk
    # Path pattern: audios/{model_name}/{Speaker}/{target_filename}
    found_path = None
    found_speaker = None
    
    for speaker in candidate_speakers:
        potential_path = os.path.join(AUDIO_BASE_PATH, model_name, speaker, target_filename)
        if os.path.exists(potential_path):
            found_path = potential_path
            found_speaker = speaker
            break
    
    if not found_path:
        search_criteria = f"speaker '{name}'" if name else f"any {gender} speaker"
        raise HTTPException(
            status_code=404, 
            detail=f"Audio '{target_filename}' not found for {search_criteria} in {model_name}."
        )

    # 4. Read and Encode
    try:
        with open(found_path, "rb") as audio_file:
            audio_bytes = audio_file.read()
            base64_string = base64.b64encode(audio_bytes).decode('utf-8')
            
        return JSONResponse(content={
            "model": model_name,
            "filename": target_filename,
            "speaker_found": found_speaker,
            "audio_base64": base64_string
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"File processing error: {str(e)}")

# --- ENDPOINTS ---

@app.get("/elevenlabs")
def get_elevenlabs_audio(
    sentence: str = Query(..., description="Text content"),
    gender: Optional[str] = Query(None, regex="^(male|female)$"),
    name: Optional[str] = Query(None, description="Specific speaker name")
):
    return _get_audio_response(
        model_name="elevenlabs",
        gender_map=ELEVENLABS_GENDER_MAP,
        sentence=sentence,
        gender=gender,
        name=name
    )

@app.get("/parler")
def get_parler_audio(
    sentence: str = Query(..., description="Text content"),
    gender: Optional[str] = Query(None, regex="^(male|female)$"),
    name: Optional[str] = Query(None, description="Specific speaker name")
):
    return _get_audio_response(
        model_name="parler",
        gender_map=PARLER_GENDER_MAP,
        sentence=sentence,
        gender=gender,
        name=name
    )