import os
import gradio as gr
from translator.agent import TranslationAgent
import re

agent = TranslationAgent()

title = "Unity → Unreal Code Translator"
description = "AI-assisted developer tool that automatically converts Unity C# scripts into Unreal Engine-ready C++ and Blueprint Python scripts. It's designed to help developers migrate gameplay logic, components, and class structures from Unity to Unreal with minimal manual rewriting. Powered by large language models Hugging Face Mistral and OpenAI"


def extract_class_name(csharp_code: str) -> str:
    """Extracts Unity class name for file naming."""
    match = re.search(r"class\s+(\w+)", csharp_code)
    return match.group(1) if match else "MyActor"


def split_cpp_code(full_cpp_code: str):
    """
    Split model output into .h and .cpp parts using file markers like:
    ---FILE: MyActor.h---
    ---FILE: MyActor.cpp---
    """
    if not full_cpp_code:
        return "", ""

    header_match = re.search(r"---FILE:\s*.*?\.h---([\s\S]*?)(?=---FILE:|$)", full_cpp_code)
    cpp_match = re.search(r"---FILE:\s*.*?\.cpp---([\s\S]*?)(?=---FILE:|$)", full_cpp_code)

    header_code = header_match.group(1).strip() if header_match else ""
    cpp_code = cpp_match.group(1).strip() if cpp_match else full_cpp_code.strip()

    # Cleanup any ``` fences if the model wrapped code blocks
    header_code = re.sub(r"^```[a-zA-Z]*|```$", "", header_code, flags=re.MULTILINE).strip()
    cpp_code = re.sub(r"^```[a-zA-Z]*|```$", "", cpp_code, flags=re.MULTILINE).strip()

    return header_code, cpp_code


def do_translate(csharp_code, target_format, model_choice):
    """Run translation and return all outputs."""
    class_name = extract_class_name(csharp_code)
    cpp_code, blueprint_code = agent.translate(
        csharp_code, target_format=target_format, model_choice=model_choice
    )

    # Split .h / .cpp parts if present
    header_code, cpp_code_only = split_cpp_code(cpp_code)

    # Save to temporary files for download
    h_filename = f"{class_name}.h"
    cpp_filename = f"{class_name}.cpp"

    with open(h_filename, "w", encoding="utf-8") as h:
        h.write(header_code or "// No header generated")

    with open(cpp_filename, "w", encoding="utf-8") as c:
        c.write(cpp_code_only or "// No cpp generated")

    # Split blueprint output into main and extra sections if available
    blueprint_parts = blueprint_code.split("###")
    main_script = blueprint_parts[0].strip()
    extra_info = blueprint_parts[1].strip() if len(blueprint_parts) > 1 else ""

    return header_code, cpp_code_only, main_script, extra_info, h_filename, cpp_filename


# Gradio UI
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"### {title}")
    gr.Markdown(description)

    with gr.Row():
        csharp_input = gr.Textbox(
            lines=20,
            label="Unity C# Code",
            placeholder="Paste your Unity MonoBehaviour or script here..."
        )

    with gr.Row():
        target_format = gr.Radio(
            ["cpp", "blueprint", "both"],
            value="both",
            label="Target Format"
        )
        model_choice = gr.Radio(
            ["hf", "openai"],
            value="hf",
            label="Model Backend"
        )

    translate_btn = gr.Button("🔁 Translate")

    # C++ output tab
    with gr.Tab("C++ Output"):
        with gr.Row():
            header_box = gr.Code(label=".h (Header File)", language="cpp")
            cpp_box = gr.Code(label=".cpp (Implementation File)", language="cpp")
        with gr.Row():
            download_h = gr.File(label="Download .h")
            download_cpp = gr.File(label="Download .cpp")

    # Blueprint output tab
    with gr.Tab("Blueprint Output"):
        with gr.Row():
            blueprint_main = gr.Textbox(label="Blueprint Python Script", lines=15)
        with gr.Row():
            blueprint_notes = gr.Textbox(
                label="Blueprint Hints / Logs",
                lines=10,
                placeholder="Any additional output, logs, or notes will appear here."
            )

    translate_btn.click(
        do_translate,
        inputs=[csharp_input, target_format, model_choice],
        outputs=[
            header_box, cpp_box,
            blueprint_main, blueprint_notes,
            download_h, download_cpp
        ]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True, debug=True)