import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_ID = "arinbalyan/code-translation-lora"

theme = (
    gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")
    .set(button_primary_background_fill_hover="#4f46e5")
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
)

LANGUAGES = [
    "Python", "JavaScript", "TypeScript", "Java", "Go",
    "Rust", "C++", "Ruby", "PHP", "C#", "Swift", "Kotlin",
]


def generate(
    task: str,
    source_lang: str,
    target_lang: str,
    source_code: str,
    description: str,
) -> str:
    if task == "Code Translation":
        if not source_code.strip():
            return "# Enter source code to translate."
        instruction = f"Translate the following {source_lang} code to {target_lang}:"
        prompt = (
            f"{instruction}\n\n"
            f"Source code ({source_lang}):\n"
            f"```{source_lang.lower()}\n"
            f"{source_code.strip()}\n"
            f"```\n\n"
            f"Translated code ({target_lang}):\n"
            f"```{target_lang.lower()}\n"
        )
    else:
        if not description.strip():
            return "# Describe what you want to code."
        prompt = (
            f"### Instruction:\n{description.strip()}\n\n"
            f"### Code:\n"
        )

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.3,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    code = result[len(prompt):]
    if code.endswith("```"):
        code = code[:-3].strip()
    return code.strip()


def update_visibility(task: str):
    return {
        translation_box: gr.update(visible=task == "Code Translation"),
        translation_source_lang: gr.update(visible=task == "Code Translation"),
        translation_target_lang: gr.update(visible=task == "Code Translation"),
        translation_target_label: gr.update(visible=task == "Code Translation"),
        generation_box: gr.update(visible=task == "Code Generation"),
        generation_label: gr.update(visible=task == "Code Generation"),
    }


with gr.Blocks(theme=theme, title="Code Translation") as demo:
    gr.Markdown(
        """
        # 🔄 Code Translation Studio
        Fine-tuned Qwen2.5-Coder-1.5B on code translation and generation tasks.
        Translate code between languages or generate code from descriptions.
        """
    )

    task = gr.Radio(
        choices=["Code Translation", "Code Generation"],
        value="Code Translation",
        label="Task",
    )

    with gr.Row():
        with gr.Column(scale=1):
            translation_box = gr.Textbox(
                label="Source Code",
                placeholder="Paste your source code here...",
                lines=10,
                visible=True,
            )
            translation_source_lang = gr.Dropdown(
                choices=LANGUAGES,
                value="Python",
                label="Source Language",
                visible=True,
            )
            translation_target_lang = gr.Dropdown(
                choices=LANGUAGES,
                value="JavaScript",
                label="Target Language",
                visible=True,
            )

        with gr.Column(scale=1):
            translation_target_label = gr.Code(
                label="Translated Code",
                language="javascript",
                lines=10,
                visible=True,
            )

    generation_box = gr.Textbox(
        label="Description",
        placeholder="e.g., Write a function to merge two sorted lists in Python...",
        lines=4,
        visible=False,
    )
    generation_label = gr.Code(
        label="Generated Code",
        language="python",
        lines=10,
        visible=False,
    )

    run_btn = gr.Button("Generate", variant="primary", size="lg")

    gr.Examples(
        examples=[
            [
                "Code Translation",
                "def add(a, b):\n    return a + b",
                "Python",
                "JavaScript",
                "",
            ],
            [
                "Code Translation",
                "function isEven(n) { return n % 2 === 0; }",
                "JavaScript",
                "Python",
                "",
            ],
            [
                "Code Generation",
                "",
                "Python",
                "JavaScript",
                "Write a Python function to reverse a string.",
            ],
            [
                "Code Generation",
                "",
                "Python",
                "JavaScript",
                "Write a function to check if a number is prime in Python.",
            ],
        ],
        inputs=[
            task,
            translation_box,
            translation_source_lang,
            translation_target_lang,
            generation_box,
        ],
        label="Try one of these",
    )

    task.change(
        fn=update_visibility,
        inputs=task,
        outputs=[
            translation_box,
            translation_source_lang,
            translation_target_lang,
            translation_target_label,
            generation_box,
            generation_label,
        ],
    )

    run_btn.click(
        fn=generate,
        inputs=[
            task,
            translation_source_lang,
            translation_target_lang,
            translation_box,
            generation_box,
        ],
        outputs=[translation_target_label, generation_label],
    )

    gr.Markdown(
        """
        <div style="text-align:center; margin-top:1rem; color:gray; font-size:0.9rem;">
        Base: Qwen2.5-Coder-1.5B &nbsp;•&nbsp; Fine-tune: LoRA (r=8) &nbsp;•&nbsp; Trained on Kaggle P100
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch()