import logging
import os
import io
import asyncio
import numpy as np
import torch
import soundfile as sf
import uuid
import time
import httpx
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional

from omnivoice import OmniVoice, OmniVoiceGenerationConfig
from text_preprocessor import chunk_text

logging.basicConfig(
    level=logging.WARNING,
    format="%(asctime)s %(name)s %(levelname)s: %(message)s",
)
logger = logging.getLogger(__name__)

# FastAPI app
app = FastAPI(title="Arabic TTS Server (OmniVoice)")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://arabic-tts-frontend.web.app",
        "https://arabic-tts-frontend.firebaseapp.com",
        "http://localhost:3000",
        "http://localhost:8000"
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables for model
CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
model = None
sampling_rate = 24000

# Simple In-Memory Database for State Tracking (Step 4 preview)
tasks_db = {}

# ── Distributed Worker Registry ──────────────────────────────────────────────
# Worker HuggingFace Space URLs for parallel audio generation.
WORKER_URLS = [
    "https://bilalrhch-arabic-tts-worker-1.hf.space",
    "https://bilalrhch-arabic-tts-worker-2.hf.space",
    "https://bilalrhch-arabic-tts-worker-3.hf.space",
    "https://bilalrhch-arabic-tts-worker-4.hf.space",
]

@app.on_event("startup")
async def startup_event():
    global model, sampling_rate
    print(f"Loading OmniVoice model from {CHECKPOINT} ...")
    model = OmniVoice.from_pretrained(
        CHECKPOINT,
        load_asr=True,
        device_map="cpu",  # Using CPU by default, adjust if a GPU is available.
    )
    sampling_rate = model.sampling_rate
    print("Model loaded successfully!")

class SynthesizeRequest(BaseModel):
    text: str
    voice: Optional[str] = "Auto"
    speed: Optional[float] = 1.0

import os
import shutil

def process_audio_task(task_id: str, chunks: list[str], speed: float, voice_id: str):
    """
    Background worker that iterivately generates audio, saves chunks to disk,
    concatenates them, and handles cleanup. Validates consistent voice!
    """
    try:
        total_chunks = len(chunks)
        tasks_db[task_id]["status"] = "processing"
        tasks_db[task_id]["total_chunks"] = total_chunks
        
        chunk_dir = os.path.join("audio_chunks", task_id)
        os.makedirs(chunk_dir, exist_ok=True)
        
        gen_config = OmniVoiceGenerationConfig(
            num_step=32,
            guidance_scale=2.0,
            denoise=True,
            preprocess_prompt=False,
            postprocess_output=False,
        )

        master_voice_prompt = None

        # Check if user requested a specific built-in voice
        if voice_id and voice_id != "Auto":
            voice_path = os.path.join("voices", f"{voice_id}.wav")
            text_path = os.path.join("voices", f"{voice_id}.txt")
            if os.path.exists(voice_path):
                ref_text = None
                if os.path.exists(text_path):
                    with open(text_path, "r", encoding="utf-8") as f:
                        ref_text = f.read().strip()
                try:
                    master_voice_prompt = model.create_voice_clone_prompt(ref_audio=voice_path, ref_text=ref_text)
                except Exception as e:
                    logger.warning(f"Voice clone setup failed: {e}")

        for i, chunk in enumerate(chunks):
            # Update state tracker
            tasks_db[task_id]["current_chunk"] = i + 1
            tasks_db[task_id]["progress"] = int(((i) / total_chunks) * 100)
            
            chunk_path = os.path.join(chunk_dir, f"chunk_{i}.wav")
            
            # Check if already generated for resume ability
            if not os.path.exists(chunk_path):
                kw = dict(
                    text=chunk,
                    language="Auto",
                    generation_config=gen_config
                )
                if speed is not None and speed != 1.0:
                    kw["speed"] = speed

                # Apply consistent voice cloning (prevents mid-book gender switching)
                if master_voice_prompt is not None:
                    kw["voice_clone_prompt"] = master_voice_prompt

                # Generate Audio via OmniVoice
                audio = model.generate(**kw)
                waveform = audio[0].squeeze(0).numpy()
                
                # Save chunk to disk incrementally
                sf.write(chunk_path, waveform, sampling_rate, format='wav', subtype='PCM_16')

                # If Auto mode, use the FIRST successfully generated chunk as the reference 
                # voice for ALL subsequent chunks. This locks the randomly chosen voice!
                if master_voice_prompt is None and i == 0:
                    try:
                        master_voice_prompt = model.create_voice_clone_prompt(ref_audio=chunk_path, ref_text=chunk)
                    except Exception as e:
                        logger.warning(f"Could not extract voice clone from chunk 0: {e}")
            
        # All chunks generated, now concatenate
        tasks_db[task_id]["status"] = "stitching"
        
        all_data = []
        sr = sampling_rate
        for i in range(total_chunks):
            chunk_path = os.path.join(chunk_dir, f"chunk_{i}.wav")
            data, sr = sf.read(chunk_path)
            all_data.append(data)
            
        final_waveform = np.concatenate(all_data)
        
        final_dir = os.path.join("static", "audio")
        os.makedirs(final_dir, exist_ok=True)
        final_path = os.path.join(final_dir, f"{task_id}.wav")
        
        sf.write(final_path, final_waveform, sr, format='wav', subtype='PCM_16')
        
        # Cleanup temporary chunks
        shutil.rmtree(chunk_dir)
            
        tasks_db[task_id]["status"] = "completed"
        tasks_db[task_id]["progress"] = 100
        tasks_db[task_id]["download_url"] = f"audio/{task_id}.wav"
        
    except Exception as e:
        logger.error(f"Background task failed: {str(e)}")
        tasks_db[task_id]["status"] = "failed"
        tasks_db[task_id]["error"] = str(e)


# ── Distributed Orchestrator: Parallel dispatch to workers ────────────────────
async def dispatch_to_workers(task_id: str, chunks: list[str], speed: float, voice: str):
    """
    Distributes text chunks across available workers in parallel.
    Each chunk is assigned to a worker in a round-robin fashion.
    All workers run simultaneously.
    """
    tasks_db[task_id]["status"] = "processing"
    total_chunks = len(chunks)
    tasks_db[task_id]["total_chunks"] = total_chunks

    # Assign each chunk to a worker (round-robin)
    assignments = []
    for i, chunk in enumerate(chunks):
        worker_url = WORKER_URLS[i % len(WORKER_URLS)]
        assignments.append({
            "worker_url": worker_url,
            "chunk": chunk,
            "chunk_index": i,
        })

    # Define a coroutine that sends ONE chunk to ONE worker
    async def send_chunk(assignment: dict, client: httpx.AsyncClient):
        chunk_index = assignment["chunk_index"]
        url = f"{assignment['worker_url']}/synthesize_chunk"
        payload = {
            "text": assignment["chunk"],
            "voice": voice,
            "speed": speed,
            "chunk_index": chunk_index,
        }
        try:
            # Long timeout for slow CPU workers (10 minutes per chunk)
            response = await client.post(url, json=payload, timeout=600.0)
            response.raise_for_status()
            return chunk_index, response.content  # Return chunk index + audio bytes
        except Exception as e:
            logger.error(f"Worker failed for chunk {chunk_index}: {e}")
            return chunk_index, None  # Signal failure

    # Fire ALL chunk requests simultaneously using asyncio
    audio_results = {}  # {chunk_index: audio_bytes}

    async with httpx.AsyncClient() as client:
        # Create all tasks at once — this is the parallelism!
        coroutines = [send_chunk(a, client) for a in assignments]
        results = await asyncio.gather(*coroutines)

    # Collect results
    for chunk_index, audio_bytes in results:
        if audio_bytes:
            audio_results[chunk_index] = audio_bytes
        tasks_db[task_id]["progress"] = int((len(audio_results) / total_chunks) * 100)

    # Check for failures
    if len(audio_results) != total_chunks:
        tasks_db[task_id]["status"] = "failed"
        tasks_db[task_id]["error"] = "One or more workers failed to generate audio."
        return

    # ── Stitch all audio chunks in ORDER ─────────────────────────────────────
    tasks_db[task_id]["status"] = "stitching"

    all_data = []
    sr = sampling_rate
    for i in range(total_chunks):
        buffer = io.BytesIO(audio_results[i])
        data, sr = sf.read(buffer)
        all_data.append(data)

    final_waveform = np.concatenate(all_data)

    final_dir = os.path.join("static", "audio")
    os.makedirs(final_dir, exist_ok=True)
    final_path = os.path.join(final_dir, f"{task_id}.wav")
    sf.write(final_path, final_waveform, sr, format='wav', subtype='PCM_16')

    tasks_db[task_id]["status"] = "completed"
    tasks_db[task_id]["progress"] = 100
    tasks_db[task_id]["download_url"] = f"audio/{task_id}.wav"

@app.post("/synthesize")
async def synthesize(req: SynthesizeRequest, background_tasks: BackgroundTasks):
    if not model:
        raise HTTPException(status_code=500, detail="Model not loaded yet.")

    chunks = chunk_text(req.text.strip())
    if not chunks:
        raise HTTPException(status_code=400, detail="Text is empty or invalid.")

    task_id = str(uuid.uuid4())
    tasks_db[task_id] = {
        "task_id": task_id,
        "status": "pending",
        "progress": 0,
        "current_chunk": 0,
        "total_chunks": len(chunks),
    }

    # Use distributed dispatch to workers for parallel generation
    background_tasks.add_task(
        asyncio.run,
        dispatch_to_workers(task_id, chunks, req.speed, req.voice)
    )

    return JSONResponse(content={
        "task_id": task_id,
        "message": f"Dispatching {len(chunks)} chunks to {len(WORKER_URLS)} workers."
    })

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    if task_id not in tasks_db:
         raise HTTPException(status_code=404, detail="Task not found.")
    return JSONResponse(content=tasks_db[task_id])
# ── Distributed Worker: /synthesize_chunk endpoint ───────────────────────────
class ChunkRequest(BaseModel):
    text: str
    voice: Optional[str] = "Auto"
    speed: Optional[float] = 1.0
    chunk_index: int = 0

@app.post("/synthesize_chunk")
async def synthesize_chunk(req: ChunkRequest):
    """
    Lightweight endpoint for distributed workers.
    Receives a single text chunk, generates audio, and returns it as bytes.
    The Orchestrator calls this on multiple Workers simultaneously.
    """
    if not model:
        raise HTTPException(status_code=500, detail="Model not loaded.")
    
    gen_config = OmniVoiceGenerationConfig(
        num_step=32,
        guidance_scale=2.0,
        denoise=True,
        preprocess_prompt=False,
        postprocess_output=False,
    )
    
    kw = dict(text=req.text.strip(), language="Auto", generation_config=gen_config)
    if req.speed != 1.0:
        kw["speed"] = req.speed
    
    # Voice selection
    if req.voice and req.voice != "Auto":
        voice_path = os.path.join("voices", f"{req.voice}.wav")
        if os.path.exists(voice_path):
            try:
                kw["voice_clone_prompt"] = model.create_voice_clone_prompt(ref_audio=voice_path)
            except Exception as e:
                logger.warning(f"Voice clone failed: {e}")
    
    try:
        audio = model.generate(**kw)
        waveform = audio[0].squeeze(0).numpy()
        
        # Return the waveform as bytes in a WAV container
        buffer = io.BytesIO()
        sf.write(buffer, waveform, sampling_rate, format='wav', subtype='PCM_16')
        buffer.seek(0)
        
        headers = {"X-Chunk-Index": str(req.chunk_index)}
        return StreamingResponse(buffer, media_type="audio/wav", headers=headers)
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Ensure voices directory is explicitly available
os.makedirs("voices", exist_ok=True)
app.mount("/voices", StaticFiles(directory="voices"), name="voices")

# Mount static files directly on root
app.mount("/", StaticFiles(directory="static", html=True), name="static")