import json
import asyncio
import httpx
from loguru import logger
from typing import List, Dict

OLLAMA_URL = "http://127.0.0.1:11434/api/generate"
MODEL_NAME = "gemma4:e2b"
INPUT_FILE = "data/output.json"
OUTPUT_FILE = "data/scripts.json"

async def generate_script(client: httpx.AsyncClient, headline: str, content: str, attempt: int = 1) -> str:
    system_prompt = (
        "You are an expert newscaster script writer. Your sole purpose is to convert articles into scripts for a TTS engine. "
        "CRITICAL RULE: Provide ONLY the literal words to be spoken by the news anchor. "
        "STRICTLY FORBIDDEN: Do not include peripheral annotations, stage directions, or non-spoken text. "
        "Specifically, do NOT include:\n"
        "- Bracketed text (e.g., [Music], [Pause], [Cut to B-roll])\n"
        "- Parenthetical notes (e.g., (smiling), (slowly))\n"
        "- Speaker labels or names (e.g., 'Anchor:', 'Reporter:', 'Host:')\n"
        "- Scene headings, timestamps, or markdown formatting.\n"
        "The output must be a clean stream of spoken text ready for a teleprompter."
    )
    prompt = f"Write a detailed, lengthy newscaster's script based on the following text:\n\n```text\n{content}```"
    payload = {
        "model": MODEL_NAME,
        "prompt": prompt,
        "system": system_prompt,
        "stream": False
    }

    try:
        logger.info(f"Generating script for: {headline} (Attempt {attempt})")
        # High timeout for slow local LLM
        response = await client.post(OLLAMA_URL, json=payload, timeout=600.0)
        response.raise_for_status()
        result = response.json()
        return result.get("response", "")
    except (httpx.HTTPStatusError, httpx.RequestError) as e:
        if attempt < 3:
            wait = attempt * 10
            logger.warning(f"Error occurred for {headline}: {e}. Retrying in {wait}s...")
            await asyncio.sleep(wait)
            return await generate_script(client, headline, content, attempt + 1)
        logger.error(f"Max retries reached for {headline}. Final error: {e}")
        return ""
    except Exception:
        logger.exception(f"Unexpected error occurred while requesting script for {headline}")
        return ""

async def main():
    try:
        with open(INPUT_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.error(f"Failed to read {INPUT_FILE}: {e}")
        return

    articles = data.get("articles", [])
    if not articles:
        logger.warning("No articles found in the input file.")
        return

    results = []
    async with httpx.AsyncClient() as client:
        for i, article in enumerate(articles):
            headline = article.get("headline", "Unknown Headline")
            body = article.get("body", "")

            if not body:
                logger.warning(f"Skipping article {headline} because body is empty.")
                continue

            script = await generate_script(client, headline, body)
            if script:
                results.append({
                    "headline": headline,
                    "script": script
                })
                # Incremental save
                with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
                    json.dump(results, f, indent=4, ensure_ascii=False)
            else:
                logger.warning(f"Failed to generate script for {headline}")

            # Avoid overloading Ollama
            await asyncio.sleep(2)

    logger.info(f"Completed processing. Total scripts written: {len(results)}")

if __name__ == "__main__":
    asyncio.run(main())