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

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

# --- CONFIGURATION ---
AUDIO_BASE_PATH = "/root/tts_arena/audios"

# 1. ELEVENLABS Speaker Mapping (Gender -> List of Speaker Names)
ELEVENLABS_GENDER_MAP = {
    "male": [
        "Adam", "Bill"
    ],
    "female": [
        "Alice"
    ]
}

# 2. PARLER Speaker Mapping (Gender -> List of Speaker Names)
# UPDATE THIS with your actual Parler folder names/speakers
PARLER_GENDER_MAP = {
    "male": [
        "Rohit"
    ],
    "female": [
        "Divya"
    ]
}

# Global Data Cache
db = {
    "elevenlabs": pd.DataFrame(),
    "parler": pd.DataFrame()
}

def normalize_text(text):
    """Normalize text to ensure cache hits (trim whitespace, strict string conversion)."""
    return str(text).strip()

@app.on_event("startup")
def load_data():
    """Loads CSV mappings into memory on startup."""
    try:
        # Load ElevenLabs CSV
        if os.path.exists("elevenlabs.csv"):
            df = pd.read_csv("elevenlabs.csv")
            # Ensure text column exists and normalize it
            if 'text' in df.columns and 'filename' in df.columns:
                df['text_clean'] = df['text'].apply(normalize_text)
                db["elevenlabs"] = df
                print(f"✅ ElevenLabs: Loaded {len(df)} mappings.")
            else:
                print("⚠️ ElevenLabs CSV missing 'text' or 'filename' columns.")
        else:
            print("⚠️ 'elevenlabs.csv' not found.")
        
        # Load Parler CSV
        if os.path.exists("parler.csv"):
            df = pd.read_csv("parler.csv")
            if 'text' in df.columns and 'filename' in df.columns:
                df['text_clean'] = df['text'].apply(normalize_text)
                db["parler"] = df
                print(f"✅ Parler: Loaded {len(df)} mappings.")
            else:
                print("⚠️ Parler CSV missing 'text' or 'filename' columns.")
        else:
            print("⚠️ 'parler.csv' not found.")
            
    except Exception as e:
        print(f"❌ Critical Error Loading CSVs: {e}")

def _get_audio_response(model_name: str, gender_map: Dict[str, List[str]], sentence: str, gender: str):
    """
    Generic helper to find file -> encode base64 -> return JSON.
    """
    # 1. Lookup Sentence in DB
    df = db.get(model_name)
    if df is None or df.empty:
        raise HTTPException(status_code=503, detail=f"{model_name} 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.")
    
    # Get the expected filename from the CSV
    target_filename = match.iloc[0]['filename']
    
    # 2. Identify Candidate Speakers for the requested Gender
    candidate_speakers = gender_map.get(gender, [])
    if not candidate_speakers:
        raise HTTPException(status_code=400, detail=f"No speakers defined for gender: {gender}")

    # 3. Search for the file on disk
    # Path pattern: audios/{model_name}/{Speaker}/{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:
        # Debugging aid: return what we looked for
        raise HTTPException(
            status_code=404, 
            detail=f"Audio '{target_filename}' not found in any {gender} folder for {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: str = Query(..., regex="^(male|female)$")
):
    return _get_audio_response(
        model_name="elevenlabs",
        gender_map=ELEVENLABS_GENDER_MAP,
        sentence=sentence,
        gender=gender
    )

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