import asyncio
import json
import os
import io
from pathlib import Path
from typing import List

from loguru import logger
from pydantic import BaseModel
from kokoro import KPipeline
import soundfile as sf

# --- Configuration ---
INPUT_FILE = "data/scripts.json"
OUTPUT_DIR = Path("data/audio")
VOICE = "af_heart"
CONCURRENCY_LIMIT = 1

class ScriptItem(BaseModel):
    headline: str
    script: str

class KokoroTTS:
    def __init__(self, lang_code: str = 'a'):
        self.pipeline = KPipeline(lang_code=lang_code)

    async def text_to_speech(self, text: str, output_wav: Path):
        """
        Generates speech locally using Kokoro and saves the resulting audio to a WAV file.
        """
        try:
            logger.debug(f"Generating TTS locally for: {text[:50]}...")

            # The pipeline returns a generator
            generator = self.pipeline(text, voice=VOICE)

            # We take the first audio segment produced
            # For short scripts, it's usually one segment.
            # For longer ones, we concatenate.
            all_audio = []
            for _, _, audio in generator:
                all_audio.append(audio)

            if not all_audio:
                raise RuntimeError("No audio data generated by Kokoro")

            import numpy as np
            final_audio = np.concatenate(all_audio)

            # Kokoro outputs at 24000 Hz
            sf.write(output_wav, final_audio, 24000)

        except Exception as e:
            logger.exception(f"Unexpected error during Kokoro TTS: {e}")
            raise

async def convert_wav_to_mp3(wav_path: Path, mp3_path: Path):
    """
    Converts a WAV file to MP3 using ffmpeg.
    """
    try:
        cmd = ["ffmpeg", "-y", "-i", str(wav_path), "-q:a", "2", str(mp3_path)]
        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            logger.error(f"ffmpeg error: {stderr.decode()}")
            return False
        return True
    except Exception as e:
        logger.exception(f"Unexpected error during conversion: {e}")
        return False

async def process_article(semaphore: asyncio.Semaphore, client: KokoroTTS, item: ScriptItem):
    """
    Processes a single article: combines text, generates audio, and converts to MP3.
    """
    async with semaphore:
        if not item.headline:
            logger.warning("Skipping article processing: headline is empty.")
            return
        safe_headline = "".join([c if c.isalnum() else "_" for c in item.headline])[:100]
        wav_path = OUTPUT_DIR / f"{safe_headline}.wav"
        mp3_path = OUTPUT_DIR / f"{safe_headline}.mp3"

        full_text = f"{item.headline}. {item.script}"

        try:
            await client.text_to_speech(full_text, wav_path)

            if await convert_wav_to_mp3(wav_path, mp3_path):
                logger.info(f"Successfully created audio for: {item.headline}")
            else:
                logger.error(f"Failed to convert audio for: {item.headline}")

            if wav_path.exists():
                wav_path.unlink()

        except Exception as e:
            logger.error(f"Error processing '{item.headline}': {e}")

async def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    try:
        with open(INPUT_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        scripts = [ScriptItem(**item) for item in data]
    except Exception as e:
        logger.error(f"Failed to load input file {INPUT_FILE}: {e}")
        return

    client = KokoroTTS()

    logger.info(f"Loaded {len(scripts)} scripts. Starting local TTS generation via Kokoro...")
    semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)

    tasks = [process_article(semaphore, client, item) for item in scripts]

    total = len(scripts)
    for i, coro in enumerate(asyncio.as_completed(tasks), 1):
        await coro
        logger.info(f"🚀 Progress: 🎙️ {i}/{total} articles processed! ✨")

    logger.info("TTS generation complete.")

if __name__ == "__main__":
    logger.add("tts_generator.log", rotation="1 MB")
    asyncio.run(main())