import gradio as gr
import torch
from transformers import pipeline, AutoTokenizer, AutoModel
import soundfile as sf
import io
import os
from fastapi import FastAPI, HTTPException, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from typing import Optional
import uvicorn
import secrets
from cryptography.fernet import Fernet
import json

# Custom API Key Setup (tera control – keys file mein store, encrypted)
API_KEYS_FILE = "api_keys.json"
ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY", Fernet.generate_key().decode())  # Prod mein env set kar
cipher_suite = Fernet(ENCRYPTION_KEY.encode())

# Load or init API keys
def load_api_keys():
    if os.path.exists(API_KEYS_FILE):
        with open(API_KEYS_FILE, 'r') as f:
            return json.load(f)
    return {}

def save_api_keys(keys):
    with open(API_KEYS_FILE, 'w') as f:
        json.dump(keys, f)

api_keys = load_api_keys()  # Dict like {"user1": "encrypted_key"}

# Generate new API key (Gradio UI se call kar)
def generate_api_key(username: str):
    if username in api_keys:
        return "Key already exists for this user!"
    raw_key = secrets.token_urlsafe(32)
    encrypted_key = cipher_suite.encrypt(raw_key.encode()).decode()
    api_keys[username] = encrypted_key
    save_api_keys(api_keys)
    return f"Your API Key: {raw_key} (Save it safely! Use in headers: Authorization: Bearer {raw_key})"

# TTS Model Load (High Quality Indic Parler-TTS)
device = "cuda" if torch.cuda.is_available() else "cpu"
tts_pipeline = pipeline(
    "text-to-speech",
    model="ai4bharat/indic-parler-tts",
    tokenizer="ai4bharat/indic-parler-tts",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    device=device
)

# Gradio TTS Function
def generate_speech(text: str, voice_desc: str = "A neutral male voice speaking clearly and calmly.", language: str = "Auto (Urdu/Hindi)"):
    if not text.strip():
        return None, "Enter some text!"
    
    # Caption for control: voice, emotion, speed etc.
    caption = f"{voice_desc} High quality, natural Urdu/Hindi speech with proper intonation."
    prompt = f"[{caption}]{text}"
    
    # Generate audio
    output = tts_pipeline(prompt)
    audio_array = output["audio"]
    
    # Save to buffer
    buffer = io.BytesIO()
    sf.write(buffer, audio_array, output["sampling_rate"])
    buffer.seek(0)
    
    return buffer.getvalue(), f"Generated! Voice: {voice_desc}, Lang: {language}"

# API Models
class TTSRequest(BaseModel):
    text: str
    voice_desc: Optional[str] = "A neutral male voice speaking clearly and calmly."
    language: Optional[str] = "Auto (Urdu/Hindi)"
    api_key: str  # Will be in headers, but for body too

class APIResponse(BaseModel):
    audio_base64: str  # Or URL if save files
    message: str

# FastAPI Setup
app = FastAPI(title="Custom Urdu/Hindi TTS API")

security = HTTPBearer()

async def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    raw_key = credentials.credentials
    for user_key_enc in api_keys.values():
        try:
            decrypted = cipher_suite.decrypt(user_key_enc.encode()).decode()
            if decrypted == raw_key:
                return True
        except:
            pass
    raise HTTPException(status_code=401, detail="Invalid API Key")

@app.post("/generate-speech", response_model=APIResponse)
async def api_generate_speech(request: TTSRequest, verified: bool = Depends(verify_api_key)):
    # Same logic as Gradio
    caption = f"{request.voice_desc} High quality, natural Urdu/Hindi speech with proper intonation."
    prompt = f"[{caption}]{request.text}"
    
    output = tts_pipeline(prompt)
    audio_array = output["audio"]
    
    buffer = io.BytesIO()
    sf.write(buffer, audio_array, output["sampling_rate"])
    audio_base64 = base64.b64encode(buffer.getvalue()).decode()  # Need import base64
    
    return APIResponse(audio_base64=audio_base64, message="Success!")

# Gradio Interface
with gr.Blocks(title="Urdu/Hindi TTS Generator") as demo:
    gr.Markdown("# High Quality Urdu/Hindi TTS with Custom API Control 🔥")
    gr.Markdown("Enter text in Urdu/Hindi. Control voice/emotion/speed via description. Generate your own API Key below!")
    
    with gr.Tab("TTS Generator"):
        text_input = gr.Textbox(label="Text (Urdu/Hindi)", placeholder="اسلام آباد کا موسم آج بہت اچھا ہے۔ یا हिंदी: आज का मौसम बहुत अच्छा है।", lines=3)
        voice_input = gr.Textbox(label="Voice Description (e.g., 'Young female, excited and fast')", value="A neutral male voice speaking clearly and calmly.")
        lang_input = gr.Dropdown(["Auto (Urdu/Hindi)", "Hindi", "Urdu"], value="Auto (Urdu/Hindi)", label="Language")
        audio_output = gr.Audio(label="Generated Speech")
        generate_btn = gr.Button("Generate Speech")
        
        generate_btn.click(
            generate_speech,
            inputs=[text_input, voice_input, lang_input],
            outputs=[audio_output, gr.Textbox(label="Status")]
        )
    
    with gr.Tab("Custom API Key Generator"):
        username_input = gr.Textbox(label="Your Username (for key association)", placeholder="e.g., myapp_user")
        key_output = gr.Textbox(label="Your New API Key", interactive=False)
        gen_key_btn = gr.Button("Generate My API Key")
        
        gen_key_btn.click(
            generate_api_key,
            inputs=[username_input],
            outputs=[key_output]
        )
        gr.Markdown("""
        ### How to Use Custom API:
        - POST to `/generate-speech` with JSON: `{"text": "your text", "voice_desc": "description", "language": "Auto"}`
        - Header: `Authorization: Bearer YOUR_KEY`
        - Response: JSON with `audio_base64` (decode to play).
        
        Example cURL:
        ```
        curl -X POST "https://yourspace.hf.space/generate-speech" \
        -H "Authorization: Bearer YOUR_API_KEY" \
        -H "Content-Type: application/json" \
        -d '{"text": "ہیلو دنیا", "voice_desc": "Cheerful female voice"}'
        ```
        Integrate in any app (Python/JS) – full control, no HF limits directly!
        """)

# Embed FastAPI in Gradio (run on /run for API)
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)
    # For API: uvicorn.run(app, host="0.0.0.0", port=8000)  # Separate if needed