import os
import re
import threading

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama


MODEL_REPO = "TaruniSwathi/Qwen2.5-Coder-1.5B-Java-CSharp-GGUF"
MODEL_FILE = "Qwen2.5-Coder-1.5B-Java-CSharp_V2.Q4_K_M.gguf"

STAGE1_RESPONSE_MARKER = "### Response:\n\n"
STAGE2_RESPONSE_MARKER = "### Response\n"

# Prevent two users from running CPU inference simultaneously.
generation_lock = threading.Lock()


print("Downloading GGUF model...")

model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
)

print("Loading GGUF model...")

llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=max(1, os.cpu_count() or 2),
    n_threads_batch=max(1, os.cpu_count() or 2),
    n_batch=128,
    n_gpu_layers=0,
    verbose=False,
)

print("Model loaded successfully.")


def add_java_hint(instruction: str) -> str:
    instruction = instruction.strip()

    if "java" in instruction.lower():
        return instruction

    return f"{instruction} Write the solution in Java."


def build_nl_to_java_prompt(instruction: str) -> str:
    return (
        "### Instruction:\n\n"
        f"{add_java_hint(instruction)}\n\n"
        "### Response:\n\n"
    )


def build_java_to_csharp_prompt(java_code: str) -> str:
    return (
        "### Instruction\n"
        "Translate the following Java code into equivalent C#. "
        "Write the solution in C#.\n\n"
        "### Java\n"
        f"{java_code.strip()}\n\n"
        "### Response\n"
    )


def clean_generated_code(text: str) -> str:
    text = text.strip()

    # Remove Markdown code fences if the model adds them.
    fenced = re.search(
        r"```(?:java|csharp|cs|c#)?\s*(.*?)```",
        text,
        flags=re.DOTALL | re.IGNORECASE,
    )

    if fenced:
        text = fenced.group(1).strip()

    stop_markers = [
        "### Instruction:",
        "### Instruction\n",
        "### Java:",
        "### Java\n",
        "### Response:",
        "### Response\n",
        "<|im_start|>",
        "<|im_end|>",
    ]

    for marker in stop_markers:
        if marker in text:
            text = text.split(marker, 1)[0].strip()

    return text


def run_generation(prompt: str, max_tokens: int) -> str:
    with generation_lock:
        response = llm(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=0.0,
            top_p=1.0,
            repeat_penalty=1.0,
            echo=False,
            stop=[
                "</s>",
                "<|endoftext|>",
                "<|im_end|>",
                "### Instruction:",
                "### Instruction\n",
            ],
        )

    return response["choices"][0]["text"]


def generate_code(task: str, user_input: str) -> str:
    if not user_input or not user_input.strip():
        return "Please enter a requirement or Java code."

    try:
        if task == "Natural Language → Java":
            prompt = build_nl_to_java_prompt(user_input)
            generated = run_generation(prompt, max_tokens=300)
            language = "java"

        else:
            prompt = build_java_to_csharp_prompt(user_input)
            generated = run_generation(prompt, max_tokens=400)
            language = "csharp"

        code = clean_generated_code(generated)

        if not code:
            return "The model returned an empty response. Please try again."

        return f"```{language}\n{code}\n```"

    except Exception as error:
        return (
            "Generation failed.\n\n"
            f"Error: {type(error).__name__}: {error}"
        )


def update_input(task: str):
    if task == "Natural Language → Java":
        return gr.update(
            label="Natural-language requirement",
            placeholder=(
                "Example: Write a Java method to check whether "
                "a number is prime."
            ),
            value="",
        )

    return gr.update(
        label="Java code",
        placeholder=(
            "Example:\n"
            "public static int factorial(int n) {\n"
            "    int result = 1;\n"
            "    for (int i = 2; i <= n; i++) {\n"
            "        result *= i;\n"
            "    }\n"
            "    return result;\n"
            "}"
        ),
        value="",
    )


with gr.Blocks(title="Java and C# CodeGen") as demo:
    gr.Markdown(
        """
# Java and C# CodeGen

Generate Java code from natural-language requirements or translate Java code
into equivalent C# using a fine-tuned Qwen2.5-Coder model.
"""
    )

    task = gr.Dropdown(
        choices=[
            "Natural Language → Java",
            "Java → C#",
        ],
        value="Natural Language → Java",
        label="Select task",
    )

    user_input = gr.Textbox(
        label="Natural-language requirement",
        placeholder=(
            "Example: Write a Java method to check whether "
            "a number is prime."
        ),
        lines=14,
    )

    generate_button = gr.Button(
        "Generate Code",
        variant="primary",
    )

    output = gr.Markdown()

    task.change(
        fn=update_input,
        inputs=task,
        outputs=user_input,
    )

    generate_button.click(
        fn=generate_code,
        inputs=[
            task,
            user_input,
        ],
        outputs=output,
    )

    gr.Examples(
        examples=[
            [
                "Natural Language → Java",
                "Write a Java method to calculate factorial of a number using a loop.",
            ],
            [
                "Natural Language → Java",
                "Write a Java method to reverse a string.",
            ],
            [
                "Java → C#",
                """public static int factorial(int n) {
    int result = 1;
    for (int i = 2; i <= n; i++) {
        result *= i;
    }
    return result;
}""",
            ],
        ],
        inputs=[
            task,
            user_input,
        ],
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
        max_size=10,
    ).launch()