import os
import json
import uuid

import spaces
from gradio import Server
from fastapi.responses import Response, FileResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import hf_hub_download
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

import torch
import torchaudio
from omnivoice import OmniVoice
from llama_cpp import Llama

import boto3
from typing import Generator

from prompts import get_system_prompt
from helpers import parse_dialogue_tags, create_episode_zip, postprocess_script, expected_schema, therapistCostumes, maleCostumes, femaleCostumes
from forced_alignment import forced_align

import json_repair
import re

# GLOBAL VARIABLE TO CHANGE TO TOGGLE OFF-GRID MODE
# True will use Cloudflare R2, meant for the ZeroGPU space.
# False will serve Unity build files from this Gradio Server, and store episodes here
# We are not defaulting to serving game from HF space, because outbound network bandwidth is very slow.
# As Unity build size is around 130 MB, and addressable assets are ~10MB per character (Each episode needs 3)

online = True

llm_model_path = hf_hub_download(
    repo_id="unsloth/gemma-4-26B-A4B-it-GGUF",
    filename="gemma-4-26B-A4B-it-UD-Q4_K_M.gguf"
)

tts_model = OmniVoice.from_pretrained(
    "k2-fsa/OmniVoice", dtype=torch.float16, device_map="cuda:0"
)
TTS_SAMPLE_RATE = 24000

align_processor = Wav2Vec2Processor.from_pretrained("facebook/mms-1b-all")
align_model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-1b-all", dtype=torch.float16).to("cuda:0")

language = "hin"
align_processor.tokenizer.set_target_lang(language)
align_model.load_adapter(language)

s3 = boto3.client(
    "s3",
    endpoint_url=os.getenv("R2_ENDPOINT_URL"),
    aws_access_key_id=os.getenv("R2_ACCESS_KEY_ID"),
    aws_secret_access_key=os.getenv("R2_SECRET_ACCESS_KEY"),
)

EPISODES_DIR = "./episodes"
os.makedirs(EPISODES_DIR, exist_ok=True)

characters = {}
ref_texts = {
    "emily": "Hey buddy! Ready to make something interesting? Let's go.",
    "priya": "Hey buddy! Ready to make something interesting? Let's go.",
    "mike": "I like horror stories, like a lot. And also romantic comedy.",
    "rahul": "I like horror stories, like a lot. And also romantic comedy.",
    "mia": "Hey there! I'm Hazel. What's on your mind? Let's create an awesome episode.",
    "neha": "Hey there! I'm Hazel. What's on your mind? Let's create an awesome episode.",
}

for name in ["emily", "mike", "mia", "priya", "rahul", "neha"]:
    waveform, sr = torchaudio.load(f"./characters/{name}.wav")
    characters[name] = {"ref_audio": (waveform, sr), "ref_text": ref_texts[name]}

@spaces.GPU(duration=120)
def warmup():
    print("Warming up .gguf cache...")
    llm = Llama(model_path=llm_model_path, n_gpu_layers=-1, n_ctx=1024, flash_attn=True, verbose=False)
    llm.create_chat_completion(messages=[{"role": "user", "content": "Hey!"}], max_tokens=8)
    print("Success! .gguf cached!")


@spaces.GPU()
def run_episode_generation_pipeline(character_names: dict[str, str], problem: str, extras: str = "", language="English"):
    llm = Llama(model_path=llm_model_path, n_gpu_layers=-1, n_ctx=8192, flash_attn=True, verbose=False)
    user_prompt = f"Their problem is that {problem}"
    if extras != "":
        user_prompt += f"\nTry to integrate these into the therapy at some point: {extras}"

    yield "Writing Script..."
    print("Generating Script...")

    response = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": get_system_prompt(character_names, language)},
            {"role": "user", "content": user_prompt}    
        ],
        max_tokens=4096,
    )
    
    script_raw = response['choices'][0]['message']['content']

    print("RAW LLM OUTPUT")
    print(script_raw)

    # Episode JSON
    try:
        # Try fixing the script if broken
        script = json_repair.loads(
        script_raw,
        schema=expected_schema,
        schema_repair_mode="salvage"
        )
        # Fix starting look_at tags
        for item in script.get("dialogues", []):
            original_text = item.get("dialogue", "")
            fixed_text = re.sub(r'(?<!<)look_at:', '<look_at:', original_text)
            item["dialogue"] = fixed_text
    except json.JSONDecodeError as e:
        print(f"JSON Parse Error: {e}")
        yield "Error: LLM generated an invalid script format. Please try again."
        return
    
    yield "Generating Audio..."
    print("Generating Audio...")
    all_audio = {}
    formatted_transcript = ""
    dialogues = script.get("dialogues", [])

    for i, line in enumerate(dialogues):
        char_name = line.get("character", "").lower()
        dialogue = line.get("dialogue", "")

        clean_text, tags = parse_dialogue_tags(dialogue)
        formatted_transcript += f"{char_name.capitalize()}: {clean_text}\n"

        if char_name in characters and clean_text.strip():
            char_info = characters[char_name]

            audios = tts_model.generate(
                text=clean_text,
                language=language,
                ref_audio=char_info["ref_audio"],
                ref_text=char_info["ref_text"]
            )
            audio_chunk = audios[0].squeeze()

            audio_file_name = f"dialogue_{i}.bin"
            script["dialogues"][i]["audioFileName"] = audio_file_name
            all_audio[audio_file_name] = audio_chunk

            align_result = forced_align(align_model, align_processor, audio_chunk, TTS_SAMPLE_RATE, clean_text)
            words_align = align_result["words"]
            visemes = align_result["visemes"]

            look_at = []
            for tag in tags:
                # Skip if look_at tag says to look at self, or character isn't valid
                if tag["character"] == char_name or tag["character"] not in character_names.values(): 
                    continue
                idx = min(tag["word_idx"], len(words_align) - 1) if words_align else 0
                start_t = words_align[idx]["start"] if words_align else 0.0
                look_at.append({"start": start_t, "character": tag["character"]})
            
            script["dialogues"][i]["lookAt"] = look_at
            script["dialogues"][i]["visemes"] = visemes

    script = postprocess_script(script, character_names)
    yield (script, all_audio)

# To upload episodes, for public gallery
BUCKET = "the-emergent-show"
R2_PREFIX = f"rizz-therapy-episodes"

app = Server()

dist_dir = "./rizz-therapy-frontend/dist"

app.mount("/assets", StaticFiles(directory=os.path.join(dist_dir, "assets")), name="assets")
app.mount("/costumes", StaticFiles(directory=os.path.join(dist_dir, "costumes")), name="costumes")

app.mount("/StreamingAssets", StaticFiles(directory="./StreamingAssets"), name="streaming_assets")
app.mount("/WebGL", StaticFiles(directory="./ServerData/WebGL"), name="webgl_bundles")

app.mount("/episodes", StaticFiles(directory=EPISODES_DIR), name="episodes")

BUILD_DIR = "./Build"


# In case of off-grid, it can serve the game build directly from this repo
# Not using in online version, because of network bandwidth
# Forcing Gradio to serve 140 MB of game build files is not good
@app.get("/Build/{filename}")
async def serve_unity_build(filename: str):
    filepath = os.path.join(BUILD_DIR, filename)

    if not os.path.exists(filepath):
        return Response(status_code=404, content="File not found")

    with open(filepath, "rb") as f:
        content = f.read()

    headers = {}

    if ".wasm" in filename:
        media_type = "application/wasm"
    elif ".js" in filename:
        media_type = "application/javascript"
    else:
        media_type = "application/octet-stream"
    
    # Inject brotli header
    if filename.endswith(".br"):
        headers["Content-Encoding"] = "br"
    
    return Response(content=content, media_type=media_type, headers=headers)


@app.get("/")
async def homepage():
    return FileResponse(os.path.join(dist_dir, "index.html"))

@app.api(name="generate-episode")
def generate_episode(problem: str, language: str, therapistCostume: str, maleCostume: str, femaleCostume: str) -> Generator[str, None, None]:
    if language not in ["English", "Hindi"]:
        language = "English"
    if language == "English":
        character_names = {
            "therapist": "emily",
            "male": "mike",
            "female": "mia"
        }
    else:
        character_names = {
            "therapist": "priya",
            "male": "rahul",
            "female": "neha"
        }

    yield "Starting Generation..."

    pipeline_generator = run_episode_generation_pipeline(character_names, problem, language=language)
    script = None
    all_audio = None

    for result in pipeline_generator:
        if isinstance(result, str):
            yield result
        else:
            script, all_audio = result

    # Add costume to the script
    script["therapistCostume"] = therapistCostume if therapistCostume in therapistCostumes else ""
    script["maleCostume"] = maleCostume if maleCostume in maleCostumes else ""
    script["femaleCostume"] = femaleCostume if femaleCostume in femaleCostumes else ""

    zip_bytes = create_episode_zip(script, all_audio)
    unique_id = uuid.uuid4().hex[:8]
    filename = f"episode_{unique_id}.zip"

    yield "Uploading episode..."
    # Online version
    if online:
        print("Uploading episode...")
        # This takes time as
        # HF Space is throttling outbound traffic, even though each episode is ~6 MB
        # So using Cloudflare R2.
        # Subsequent episode sharings would be blazing fast
        s3_key = f"{R2_PREFIX}/{filename}"

        s3.put_object(
            Bucket=BUCKET,
            Key=s3_key,
            Body=zip_bytes,
            ContentType="application/zip"
        )
        print(f"Successfully uploaded episode to R2: {BUCKET}/{s3_key}")
        yield f"[EPISODE]{filename}"
    
    # Off grid episode saving and access
    else:
        local_path = os.path.join(EPISODES_DIR, filename)
        with open(local_path, "wb") as f:
            f.write(zip_bytes)
        print(f"Successfully saved locally: {local_path}")
        yield f"[EPISODE]{filename}"

@app.get("/online-status")
async def get_online_status():
    return {"online": online}

if __name__ == "__main__":
    warmup()
    app.launch()