import re

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


# ============================================================
# MODEL CONFIGURATION
# ============================================================

MODEL_ID = "shibsankardhara2/Qwen2.5-Coder-1.5B-Java-CSharp_V5"


# ============================================================
# LOAD TOKENIZER
# ============================================================

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id


# ============================================================
# LOAD MODEL ON CPU
# ============================================================

print("Loading model on CPU...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)

model.eval()

print("Model loaded successfully.")


# ============================================================
# PROMPT FUNCTIONS
# ============================================================

def add_java_hint(instruction: str) -> str:
    """
    Explicitly mention Java when it is missing from the
    natural-language instruction.
    """
    instruction = instruction.strip()

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

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


def build_prompt(task: str, user_input: str) -> str:
    """
    Construct the prompt expected by the fine-tuned model.
    """
    user_input = user_input.strip()

    if task == "Natural Language → Java":
        return (
            "### Instruction:\n\n"
            f"{add_java_hint(user_input)}\n\n"
            "### Response:\n\n"
        )

    return (
        "### Instruction:\n\n"
        "Translate the following Java code into equivalent C#. "
        "Write only the C# solution.\n\n"
        "### Java:\n\n"
        f"{user_input}\n\n"
        "### Response:\n\n"
    )


# ============================================================
# GENERAL OUTPUT CLEANING
# ============================================================

def clean_output(generated_text: str) -> str:
    """
    Remove Markdown fences, repeated prompt sections and
    model-specific special tokens.
    """
    if not generated_text:
        return ""

    cleaned_text = generated_text.strip()

    # Extract the contents of a Markdown code block, when present.
    fenced_match = re.search(
        r"```(?:java|csharp|cs|c#)?\s*(.*?)```",
        cleaned_text,
        flags=re.DOTALL | re.IGNORECASE,
    )

    if fenced_match:
        cleaned_text = fenced_match.group(1).strip()

    # Remove prompt sections that the model may generate again.
    stop_markers = [
        "### Instruction:",
        "### Instruction\n",
        "### Java:",
        "### Java\n",
        "### Response:",
        "### Response\n",
        "<|im_start|>",
        "<|im_end|>",
        "<|endoftext|>",
    ]

    for marker in stop_markers:
        marker_position = cleaned_text.find(marker)

        if marker_position != -1:
            cleaned_text = cleaned_text[:marker_position].strip()

    # Remove remaining opening or closing fences.
    cleaned_text = re.sub(
        r"^```(?:java|csharp|cs|c#)?\s*",
        "",
        cleaned_text,
        flags=re.IGNORECASE,
    )

    cleaned_text = re.sub(
        r"\s*```$",
        "",
        cleaned_text,
    )

    return cleaned_text.strip()


# ============================================================
# C# OUTPUT CLEANING
# ============================================================

def clean_csharp_output(text: str) -> str:
    """
    Remove invalid virtual or override modifiers from standalone
    C# method snippets.

    If the output contains a class, struct, interface, record or
    enum declaration, modifiers are preserved.
    """
    if not text:
        return ""

    # Explicitly initialise code before it is accessed.
    code = text.strip()

    has_type_declaration = re.search(
        r"\b(class|struct|interface|record|enum)\b",
        code,
        flags=re.IGNORECASE,
    )

    if has_type_declaration:
        return code

    # public virtual int Method() -> public int Method()
    # protected override void Method() -> protected void Method()
    code = re.sub(
        r"\b(public|private|protected|internal)\s+"
        r"(?:virtual|override)\s+",
        r"\1 ",
        code,
        flags=re.IGNORECASE,
    )

    # virtual int Method() -> int Method()
    # override void Method() -> void Method()
    code = re.sub(
        r"(^|\n)(\s*)(?:virtual|override)\s+",
        r"\1\2",
        code,
        flags=re.IGNORECASE,
    )

    return code.strip()


# ============================================================
# CODE GENERATION
# ============================================================

@spaces.GPU(duration=120)
def generate_code(task: str, user_input: str) -> str:
    """
    Generate Java from natural language or translate Java to C#.
    """
    if not user_input or not user_input.strip():
        return "Please enter a requirement or Java code."

    prompt = build_prompt(task, user_input)

    max_new_tokens = (
        300
        if task == "Natural Language → Java"
        else 400
    )

    try:
        if not torch.cuda.is_available():
            return "Generation failed: GPU is not available."

        # Move the model to the allocated ZeroGPU device.
        model.to("cuda")
        model.eval()

        tokenized_inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=2048,
        )

        tokenized_inputs = {
            key: value.to("cuda")
            for key, value in tokenized_inputs.items()
        }

        with torch.inference_mode():
            generated_ids = model.generate(
                **tokenized_inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                use_cache=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        prompt_length = tokenized_inputs["input_ids"].shape[1]

        new_token_ids = generated_ids[0][prompt_length:]

        generated_text = tokenizer.decode(
            new_token_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )

        # Always run the general output cleaner.
        output_code = clean_output(generated_text)

        # Run the C#-specific cleaner only for Java → C#.
        if task == "Java → C#":
            output_code = clean_csharp_output(output_code)
            markdown_language = "csharp"
        else:
            markdown_language = "java"

        if not output_code:
            return (
                "The model returned an empty response. "
                "Please try a more specific input."
            )

        return (
            f"```{markdown_language}\n"
            f"{output_code}\n"
            "```"
        )

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

    finally:
        # Return the model to CPU after the ZeroGPU request.
        try:
            model.to("cpu")
        except Exception as move_error:
            print(f"Could not move model to CPU: {move_error}")

        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ============================================================
# UPDATE TEXTBOX
# ============================================================

def update_input(task: str):
    """
    Update the input textbox for the selected task.
    """
    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="",
    )


# ============================================================
# GRADIO INTERFACE
# ============================================================

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(
        value="Generated code will appear here."
    )

    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.",
            ],
            [
                "Natural Language → Java",
                (
                    "Write a Java method to check whether "
                    "a number is prime."
                ),
            ],
            [
                "Java → C#",
                """public static int factorial(int n) {
    int result = 1;

    for (int i = 2; i <= n; i++) {
        result *= i;
    }

    return result;
}""",
            ],
            [
                "Java → C#",
                """public boolean isEven(int n) {
    return n % 2 == 0;
}""",
            ],
        ],
        inputs=[
            task,
            user_input,
        ],
    )


# ============================================================
# START APPLICATION
# ============================================================

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