import json
import asyncio
import aiohttp
import string
from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type

class FlakyAPIError(Exception):
    pass

def extract_story_beats(text: str) -> list[dict]:
    paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
    return [{"beat_id": i + 1, "text": p} for i, p in enumerate(paragraphs)]

def translate_to_prompts(beats: list[dict], lexicon: dict) -> list[dict]:
    prompts = []
    
    for beat in beats:
        raw_text = beat["text"]
        has_exclamation_at_end = raw_text.strip().endswith('!')
        
        # Replace hyphens and em-dashes with spaces
        import re; clean_text = re.sub(r'[-\—\–]', ' ', raw_text.lower())
        
        # Strip remaining punctuation
        translator = str.maketrans('', '', string.punctuation)
        words = clean_text.translate(translator).split()
        
        score = 0.0
        
        for i, word in enumerate(words):
            if word in lexicon:
                val = lexicon[word]
                if i > 0 and words[i-1] in ["absolutely", "pure"]:
                    val += 2
                score += val
                
        if has_exclamation_at_end:
            score *= 1.5
            
        glitch = 0.0
        motion = 2.0
        
        if score >= 4.0:
            glitch = 0.8
            motion = 3.0
        elif score <= 0.0:
            glitch = 0.0
            motion = 1.0
            
        prompts.append({
            "prompt": raw_text,
            "glitch_factor": glitch,
            "motion": motion,
            "calculated_score": score
        })
        
    return prompts

@retry(wait=wait_exponential(multiplier=1, min=1, max=10), stop=stop_after_attempt(5), retry=retry_if_exception_type(FlakyAPIError))
async def fetch_video(session: aiohttp.ClientSession, prompt_data: dict) -> dict:
    async with session.post("http://127.0.0.1:8000/generate", json=prompt_data, timeout=10.0) as response:
        if response.status in [429, 503]:
            raise FlakyAPIError(f"API Error: {response.status}")
        
        data = await response.json()
        return {
            "prompt": prompt_data["prompt"],
            "video_url": data.get("video_url", ""),
            "applied_glitch": data.get("applied_glitch", 0.0),
            "applied_motion": data.get("applied_motion", 0.0),
            "calculated_score": prompt_data.get("calculated_score", 0.0)
        }

async def compile_trailer(prompts: list[dict]) -> dict:
    async with aiohttp.ClientSession() as session:
        tasks = [fetch_video(session, p) for p in prompts]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
    # Filter out exceptions if retries were exhausted
    clean_results = [r for r in results if not isinstance(r, Exception)]
    return {"status": "success", "total_shots": len(clean_results), "shots": clean_results}

if __name__ == "__main__":
    with open("input_synopsis.txt", "r") as f:
        text = f.read()
    with open("arousal_lexicon.json", "r") as f:
        lexicon = json.load(f)
        
    beats = extract_story_beats(text)
    prompts = translate_to_prompts(beats, lexicon)
    manifest = asyncio.run(compile_trailer(prompts))
    print(json.dumps(manifest, indent=2))