app.py

from fastapi import FastAPI, HTTPException, Form, BackgroundTasks
from fastapi.responses import FileResponse
import gradio as gr
from kokoro_onnx import Kokoro
import tempfile
import os
import bcrypt
from datetime import datetime, timedelta
from supabase import create_client, Client
import soundfile as sf

# ============== CONFIG ==============
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

if not SUPABASE_URL or not SUPABASE_KEY:
    raise ValueError("SUPABASE_URL and SUPABASE_KEY environment variables must be set")

DAILY_QUOTA = 50  # Increased since Kokoro is much faster
MAX_CHARS = 4500  # ~5 minutes of audio (speaking rate: ~900 chars/min)
MIN_CHARS = 5
MAX_AUDIO_DURATION = 300  # 5 minutes of audio

# Admin credentials
ADMIN_USERNAME = "madhab"
ADMIN_PASSWORD = "Madhab@Studify2024!"

# ============== SUPABASE ==============
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

def init_admin():
    """Create admin user if not exists"""
    try:
        result = supabase.table("tts_users").select("username").eq("username", ADMIN_USERNAME).execute()
        
        if not result.data:
            password_hash = bcrypt.hashpw(ADMIN_PASSWORD.encode(), bcrypt.gensalt()).decode()
            supabase.table("tts_users").insert({
                "username": ADMIN_USERNAME,
                "password_hash": password_hash,
                "role": "admin",
                "daily_limit": -1,
                "is_active": True
            }).execute()
            print(f"✅ Admin user created: {ADMIN_USERNAME}")
        else:
            print(f"✅ Admin user exists: {ADMIN_USERNAME}")
    except Exception as e:
        print(f"⚠️ Error creating admin: {e}")

# ============== KOKORO TTS MODEL ==============
print("🎤 Loading Kokoro TTS model...")
try:
    kokoro = Kokoro("kokoro-v0_19.onnx", "voices")
    print("✅ Kokoro TTS loaded successfully!")
except Exception as e:
    print(f"⚠️  Kokoro not found locally. Will download on first use.")
    kokoro = None

app = FastAPI(title="Kokoro TTS API - Professional & Fast")

@app.on_event("startup")
def startup():
    global kokoro
    if kokoro is None:
        print("📥 Downloading Kokoro TTS model...")
        kokoro = Kokoro("kokoro-v0_19.onnx", "voices")
        print("✅ Kokoro TTS loaded!")
    init_admin()

# ============== AUTH ==============
def verify_password(plain_password: str, hashed_password: str) -> bool:
    try:
        return bcrypt.checkpw(plain_password.encode(), hashed_password.encode())
    except:
        return False

def authenticate_user(username: str, password: str) -> dict:
    result = supabase.table("tts_users").select("*").eq("username", username).execute()
    
    if not result.data or len(result.data) == 0:
        raise HTTPException(status_code=401, detail="Access denied. User not found in database.")
    
    user = result.data[0]
    
    if not user.get('is_active', True):
        raise HTTPException(status_code=403, detail="Account is disabled. Contact admin.")
    
    if not verify_password(password, user['password_hash']):
        raise HTTPException(status_code=401, detail="Invalid credentials.")
    
    return user

def check_quota(username: str, daily_limit: int, role: str) -> dict:
    if role == 'admin' or daily_limit == -1:
        return {"used": 0, "remaining": -1, "is_unlimited": True}
    
    since = (datetime.utcnow() - timedelta(hours=24)).isoformat()
    result = supabase.table("tts_usage_logs").select("id", count="exact").eq("username", username).gte("created_at", since).execute()
    
    used = result.count or 0
    remaining = daily_limit - used
    
    if remaining <= 0:
        raise HTTPException(status_code=429, detail=f"Daily quota exceeded. Used {used}/{daily_limit}. Resets in 24h.")
    
    return {"used": used, "remaining": remaining, "is_unlimited": False}

def log_usage(username: str, text_length: int, language: str):
    supabase.table("tts_usage_logs").insert({
        "username": username,
        "text_length": text_length,
        "language": language,
        "created_at": datetime.utcnow().isoformat()
    }).execute()

# ============== HELPERS ==============
def cleanup_file(path: str):
    try:
        if os.path.exists(path):
            os.unlink(path)
    except:
        pass

def generate_speech(text: str, voice: str = "af_heart", speed: float = 1.0) -> str:
    """
    Generate speech using Kokoro TTS
    Available voices: af (American Female), am (American Male), bf (British Female), etc.
    """
    if len(text) < MIN_CHARS:
        raise ValueError(f"Text too short. Minimum {MIN_CHARS} characters.")
    if len(text) > MAX_CHARS:
        raise ValueError(f"Text too long. Maximum {MAX_CHARS} characters.")
    
    # Generate audio samples
    samples, sample_rate = kokoro.create(
        text=text,
        voice=voice,
        speed=speed,
        lang="en-us"  # Kokoro supports: en-us, en-gb, ja, etc.
    )
    
    # Save to temporary file
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        sf.write(tmp.name, samples, sample_rate)
        return tmp.name

# ============== API ENDPOINTS ==============
@app.get("/health")
def health():
    return {
        "status": "healthy",
        "model": "Kokoro TTS 82M",
        "speed": "10x faster than XTTS",
        "authentication": "required",
        "default_quota": DAILY_QUOTA
    }

@app.post("/api/generate")
async def generate_tts(
    background_tasks: BackgroundTasks,
    username: str = Form(...),
    password: str = Form(...),
    text: str = Form(...),
    voice: str = Form("af_heart"),  # American Female - Heart
    speed: float = Form(1.0)
):
    """
    Generate TTS with Kokoro (Fast & Emotional)
    
    Performance:
    - Max audio length: 5 minutes
    - Speaking rate: ~900 chars/minute
    - Max chars: 4500 (~5 min audio)
    - Generation time: ~20-30 seconds on CPU
    
    Available voices:
    - af_heart: American Female (warm)
    - af_bella: American Female (professional)
    - am_adam: American Male (confident)
    - am_michael: American Male (friendly) 
    - bf_emma: British Female (elegant)
    - bf_isabella: British Female (storytelling) ⭐
    
    Usage:
    curl -X POST https://your-service.hf.space/api/generate \
      -F "username=madhab" \
      -F "password=Madhab@Studify2024!" \
      -F "text=Hello world. This is much faster!" \
      -F "voice=bf_isabella" \
      -F "speed=1.0" \
      --output output.wav
    """
    user = authenticate_user(username, password)
    quota = check_quota(user['username'], user['daily_limit'], user['role'])
    
    try:
        output_path = generate_speech(text.strip(), voice, speed)
        
        if not quota['is_unlimited']:
            log_usage(user['username'], len(text), "en")
        
        background_tasks.add_task(cleanup_file, output_path)
        
        response = FileResponse(output_path, media_type="audio/wav", filename="kokoro_tts.wav")
        response.headers["X-Quota-Used"] = str(quota["used"] + (0 if quota["is_unlimited"] else 1))
        response.headers["X-Quota-Remaining"] = "unlimited" if quota["is_unlimited"] else str(quota["remaining"] - 1)
        response.headers["X-Model"] = "Kokoro-82M"
        return response
        
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"TTS generation failed: {str(e)}")

@app.post("/api/quota")
async def check_user_quota(username: str = Form(...), password: str = Form(...)):
    user = authenticate_user(username, password)
    quota = check_quota(user['username'], user['daily_limit'], user['role'])
    
    return {
        "username": user['username'],
        "role": user['role'],
        "used_today": quota["used"],
        "remaining": "unlimited" if quota["is_unlimited"] else quota["remaining"],
        "daily_limit": "unlimited" if quota["is_unlimited"] else user['daily_limit']
    }

@app.post("/api/admin/create-user")
async def create_user(
    admin_username: str = Form(...),
    admin_password: str = Form(...),
    new_username: str = Form(...),
    new_password: str = Form(...),
    role: str = Form("user"),
    daily_limit: int = Form(50)
):
    admin = authenticate_user(admin_username, admin_password)
    if admin['role'] != 'admin':
        raise HTTPException(status_code=403, detail="Admin access required")
    
    existing = supabase.table("tts_users").select("username").eq("username", new_username).execute()
    if existing.data:
        raise HTTPException(status_code=400, detail="Username already exists")
    
    password_hash = bcrypt.hashpw(new_password.encode(), bcrypt.gensalt()).decode()
    
    supabase.table("tts_users").insert({
        "username": new_username,
        "password_hash": password_hash,
        "role": role,
        "daily_limit": daily_limit,
        "is_active": True
    }).execute()
    
    return {"success": True, "username": new_username, "role": role, "daily_limit": daily_limit}

@app.post("/api/admin/list-users")
async def list_users(admin_username: str = Form(...), admin_password: str = Form(...)):
    admin = authenticate_user(admin_username, admin_password)
    if admin['role'] != 'admin':
        raise HTTPException(status_code=403, detail="Admin access required")
    
    result = supabase.table("tts_users").select("username, role, daily_limit, is_active, created_at").execute()
    return {"users": result.data}

# ============== GRADIO UI ==============
def generate_tts_gradio(username, password, text, voice="af_heart", speed=1.0):
    if not username or not password:
        raise gr.Error("Username and password required")
    if not text or len(text.strip()) < MIN_CHARS:
        raise gr.Error(f"Text must be at least {MIN_CHARS} characters")
    
    try:
        user = authenticate_user(username, password)
        quota = check_quota(user['username'], user['daily_limit'], user['role'])
        
        output_path = generate_speech(text.strip(), voice, speed)
        
        if not quota['is_unlimited']:
            log_usage(user['username'], len(text), "en")
        
        return output_path
    except HTTPException as e:
        raise gr.Error(e.detail)
    except Exception as e:
        raise gr.Error(str(e))

gradio_app = gr.Interface(
    fn=generate_tts_gradio,
    inputs=[
        gr.Textbox(label="Username", placeholder="Enter your username"),
        gr.Textbox(label="Password", type="password", placeholder="Enter your password"),
        gr.Textbox(label="Text", placeholder=f"Enter text ({MIN_CHARS}-{MAX_CHARS} chars)", lines=8),
        gr.Dropdown(
            choices=["af_heart", "af_bella", "am_adam", "am_michael", "bf_emma", "bf_isabella"],
            value="af_heart",
            label="Voice (Emotional & Expressive)"
        ),
        gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Speed")
    ],
    outputs=gr.Audio(label="Generated Speech (Kokoro TTS)", type="filepath"),
    title="🚀 Kokoro TTS - Professional & Lightning Fast",
    description=f"""
    **High-Speed Text-to-Speech with Emotional Expression**
    
    - ⚡ Lightning fast generation (~20-30 sec)
    - 🎭 Emotional & expressive voices
    - 🔒 Secure authentication required
    - 📊 Quota: {DAILY_QUOTA} generations/day
    - 🎵 Max audio: 5 minutes (4500 chars)
    - 💾 Runs smoothly on CPU
    
    Perfect for audiobooks, educational content, and storytelling!
    """,
)

app = gr.mount_gradio_app(app, gradio_app, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)