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

BASE_MODEL = "unsloth/Qwen2.5-Coder-7B-Instruct"

LANG_META = {
    "Python": {"repo": "AmareshHebbar/leetcode-python-qwen25-coder-7b", "icon": "🐍", "code_lang": "python"},
    "Java": {"repo": "AmareshHebbar/leetcode-java-qwen25-coder-7b", "icon": "☕", "code_lang": "java"},
    "C++": {"repo": "AmareshHebbar/leetcode-cpp-qwen25-coder-7b", "icon": "⚙️", "code_lang": "cpp"},
    "JavaScript": {"repo": "AmareshHebbar/leetcode-javascript-qwen25-coder-7b", "icon": "🟨", "code_lang": "javascript"},
}

tokenizer = None
model = None
_on_gpu = False


def _ensure_loaded():
    global tokenizer, model
    if model is not None:
        return
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
    m = PeftModel.from_pretrained(base_model, LANG_META["Python"]["repo"], adapter_name="Python")
    for lang, meta in LANG_META.items():
        if lang != "Python":
            m.load_adapter(meta["repo"], adapter_name=lang)
    m.eval()
    model = m


def _generate(inputs, use_cuda):
    device = "cuda" if use_cuda else "cpu"
    model.to(device)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    return model.generate(
        **inputs, max_new_tokens=512, temperature=0.2, do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )

EXAMPLES = [
    ["Merge k sorted linked lists into one sorted list.", "Divide and conquer / heap", "Python"],
    ["Given a set of non-overlapping intervals, insert a new interval and merge as needed.", "Sorting / interval merge", "JavaScript"],
    ["Find the length of the longest increasing path in a matrix.", "DFS + memoization", "Java"],
    ["Given the root of a binary tree, return the maximum path sum between any two nodes.", "Tree DFS / post-order", "C++"],
]

CSS = """
:root {
    --lc-bg: #0b0f14;
    --lc-panel: #121820;
    --lc-border: #1f2833;
    --lc-accent: #34d399;
    --lc-accent-dim: #34d39933;
    --lc-text: #e6edf3;
    --lc-text-dim: #8b98a5;
}
.gradio-container {
    background: var(--lc-bg) !important;
    font-family: 'Inter', -apple-system, sans-serif !important;
}
#lc-header {
    text-align: center;
    padding: 28px 0 8px 0;
}
#lc-header h1 {
    font-size: 2.1rem;
    font-weight: 700;
    background: linear-gradient(90deg, #34d399, #60a5fa);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 4px;
}
#lc-header p {
    color: var(--lc-text-dim);
    font-size: 0.95rem;
}
#lc-badges {
    display: flex;
    justify-content: center;
    gap: 8px;
    margin-top: 10px;
    flex-wrap: wrap;
}
.lc-badge {
    background: var(--lc-panel);
    border: 1px solid var(--lc-border);
    color: var(--lc-text-dim);
    padding: 4px 12px;
    border-radius: 999px;
    font-size: 0.78rem;
}
#lc-panel-left, #lc-panel-right {
    background: var(--lc-panel) !important;
    border: 1px solid var(--lc-border) !important;
    border-radius: 14px !important;
    padding: 18px !important;
}
#lc-generate {
    background: linear-gradient(90deg, #34d399, #22c55e) !important;
    border: none !important;
    color: #04120a !important;
    font-weight: 600 !important;
    border-radius: 10px !important;
}
#lc-lang-radio label {
    border-radius: 10px !important;
}
#lc-output-code {
    border-radius: 10px !important;
}
footer { display: none !important; }
"""

THEME = gr.themes.Base(
    primary_hue="emerald",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "sans-serif"],
).set(
    body_background_fill="#0b0f14",
    block_background_fill="#121820",
    block_border_color="#1f2833",
    body_text_color="#e6edf3",
    input_background_fill="#0b0f14",
    button_primary_background_fill="#34d399",
    button_primary_text_color="#04120a",
)


@spaces.GPU(duration=120)
def solve(problem, algorithm_tag, language):
    global _on_gpu
    if not problem.strip():
        return "", "Enter a problem statement first."

    _ensure_loaded()
    model.set_adapter(language)
    lang_name = language
    system_prompt = (
        f"You are an expert {lang_name} competitive programmer. Given a "
        f"LeetCode-style problem statement and an algorithm tag, write a "
        f"correct, efficient {lang_name} solution."
    )
    user_msg = f"Problem: {problem}"
    if algorithm_tag.strip():
        user_msg += f"\nAlgorithm: {algorithm_tag}"

    messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_msg}]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt")

    try:
        outputs = _generate(inputs, use_cuda=True)
        _on_gpu = True
        engine_note = "GPU"
    except RuntimeError as e:
        if "CUDA" not in str(e) and "cuda" not in str(e):
            raise
        outputs = _generate(inputs, use_cuda=False)
        engine_note = "CPU fallback"

    input_len = inputs["input_ids"].shape[1]
    code = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
    return code, f"{LANG_META[language]['icon']} generated with the {language} QDoRA adapter · {engine_note}"


def on_lang_change(language):
    return gr.Code(language=LANG_META[language]["code_lang"])


with gr.Blocks(title="LeetCode Multi-Language Coder Suite") as demo:
    with gr.Column(elem_id="lc-header"):
        gr.Markdown("# LeetCode Multi-Language Coder Suite")
        gr.Markdown("Qwen2.5-Coder-7B · QDoRA fine-tuned per language · execution-verified training data")
        gr.HTML(
            """
            <div id="lc-badges">
                <span class="lc-badge">🐍 Python</span>
                <span class="lc-badge">☕ Java</span>
                <span class="lc-badge">⚙️ C++</span>
                <span class="lc-badge">🟨 JavaScript</span>
                <span class="lc-badge">🔗 4 QDoRA adapters, 1 base model</span>
            </div>
            """
        )

    with gr.Row(equal_height=True):
        with gr.Column(scale=5, elem_id="lc-panel-left"):
            language = gr.Radio(
                choices=list(LANG_META.keys()),
                value="Python",
                label="Language",
                elem_id="lc-lang-radio",
            )
            problem = gr.Textbox(
                label="Problem statement",
                lines=5,
                placeholder="Given an array of integers nums and an integer target, return indices of the two numbers such that they add up to target.",
            )
            tag = gr.Textbox(label="Algorithm tag (optional)", placeholder="Hash Map")
            run = gr.Button("Generate solution", elem_id="lc-generate", size="lg")
            gr.Examples(
                examples=EXAMPLES,
                inputs=[problem, tag, language],
                label="Try an example",
            )

        with gr.Column(scale=6, elem_id="lc-panel-right"):
            status = gr.Markdown("")
            output = gr.Code(
                label="Generated solution",
                language="python",
                elem_id="lc-output-code",
                lines=22,
            )

    gr.HTML(
        """
        <div style="text-align:center; color:#8b98a5; font-size:0.82rem; margin-top:18px;">
            <a href="https://huggingface.co/collections/AmareshHebbar/leetcode-multi-language-coder-suite" style="color:#34d399;">Models & benchmarks</a>
            &nbsp;·&nbsp;
            <a href="https://github.com/amareshhebbar" style="color:#34d399;">GitHub</a>
        </div>
        """
    )

    language.change(on_lang_change, inputs=language, outputs=output)
    run.click(solve, inputs=[problem, tag, language], outputs=[output, status])

demo.launch(css=CSS, theme=THEME)