"""
3. Graph. Capture the resulting mermaid string and visualize it

To do
- create the custom gradio look
- explore making it look better
- get a better model — Qwen 30b coder
- use zerogpu

"""
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import gradio as gr
from gradio import Server
from fastapi.responses import HTMLResponse # serve the custom frontend from a route
from typing import Any, cast # to resolve PyLance freaking out over llama-cpp-python in the generate_flowchart function
from textwrap import dedent
import re # remove thinking tag from response 

# ----- Get Model ----- #
# Download Q4_K_M GGUF file from the repo
model_path = hf_hub_download(
    repo_id="unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF",
    filename="Qwen3-Coder-30B-A3B-Instruct-UD-Q3_K_XL.gguf" # fallback: Q2_K_XL
)

# Initialize llama.cpp with the local cached path


# ----- Init App ----- #
app = gr.Server(title="Code-to-Flowchart Generator")

# ----- Functions ----- #

# This is a cleaning function to resolve common syntax errors.
def quote_labels(text: str) -> str:
    # Mermaid node labels can't hold raw code characters, so quote-wrap each label body
    # A label's real closing bracket is followed by a Mermaid connector, edge-label, pipe, statement end, or EOL
    # operators after a subscript (== < <= > >= != %) are never mistaken for a close.
    END = r'(?=\s*(?:[-<][-.>xo]|==[>=xo]|\||;|$))'

    def esc(body: str) -> str:
        return (body.replace('"', "'")
                    .replace('[', '&#91;').replace(']', '&#93;')
                    .replace('{', '&#123;').replace('}', '&#125;'))

    out = []
    for line in text.split('\n'):
        line = re.sub(r'(?<=\w)\[(.*?)\]' + END, lambda m: '["' + esc(m.group(1)) + '"]', line)
        line = re.sub(r'(?<=\w)\{(.*?)\}' + END, lambda m: '{"' + esc(m.group(1)) + '"}', line)
        out.append(line)
    return '\n'.join(out)

@app.api(name="generate_flowchart")
def generate_flowchart(src_code: str) -> str:
    # check if src_code is empty
    if not src_code.strip(): return ""

    # Set system prompt
    system_prompt = dedent("""
    ## Role/Persona
    You are a senior staff software architect and compiler engineer specializing in visual control-flow mapping. Your philosophy is pure utility: you translate raw execution logic into highly accurate, scannable, structural diagrams without any conversational filler, meta-commentary, or stylistic fluff.

    ## Context/Objective
    The user will provide source code files or logic snippets. Your sole objective is to parse the syntax and output a corresponding, valid Mermaid.js flowchart graph. This graph will be rendered natively in a production UI to help developers audit execution paths at a glance.

    ## Strict Constraints
    <constraints>
    1. OUTPUT FORMAT: Output ONLY valid, raw Mermaid.js syntax.
    2. NO MARKDOWN FENCING: Do not wrap the output in ```mermaid or ``` blocks. Start directly with the Mermaid graph definition, for example: graph TD.
    3. NO PROSE: Do not include introductory text, explanations, or concluding remarks. If the code cannot be parsed, output an isolated error node.
    4. NODE NAMING: Paraphrase conditions into plain words — never put raw code, operators, quotes, parentheses, or square brackets/subscripts inside labels (write Index in bounds?, not i < len(nums); write Element is even?, not nums[i] % 2 == 0)
    </constraints>

    <banned_vocabulary>
    - Here is the flowchart
    - ```mermaid
    - ```
    - Note:
    - Explanation:
    - In this diagram
    - As requested
    </banned_vocabulary>

    ## Response Workflow
    Before outputting the final diagram syntax, perform structural parsing inside a hidden <thinking> tag according to these steps:
    1. Identify all conditional branches, including if/else, loops, including for/while, and termination points, including return/throw.
    2. Map out the execution flow nodes chronologically.
    3. Verify that every opening bracket and node label matching syntax, including [ ], ( ), and { }, is perfectly balanced and closed according to Mermaid specifications.
    4. Ensure no markdown formatting tags leak past the closing </thinking> tag.

    ## Few-Shot Examples

    Input:
    def check_status(val):
        if val > 10:
            return "Active"
        else:
            return "Inactive"

    Output:
    <thinking>
    1. Control structures: One conditional check, two return branches.
    2. Nodes: A Start, B Conditional, C Active return, D Inactive return.
    3. Syntax verification: B uses curly braces for decisions. Edges use standard arrows.
    </thinking>
    graph TD
        A[Start: check_status] --> B{val > 10}
        B -- True --> C[Return 'Active']
        B -- False --> D[Return 'Inactive']
    """).strip()

    # Reset the cache per request so no cross-request bleeding
    llm.reset()

    # Casting else PyLance gets mad
    response = cast(Any, llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": src_code}
        ],
        temperature=0.1, # Keep it quite deterministic for now
        max_tokens=1024,
        stream=False
    ))

    content = response["choices"][0]["message"]["content"]

    # remove the thinking tags from the response
    cleaned = re.sub(r'<thinking>.*?</thinking>', '', content, flags=re.DOTALL)

    # Quote-wrap each node label and escape any leaked code characters
    cleaned = quote_labels(cleaned)

    return cleaned.strip() # and remove excess whitespace

# ----- Custom Frontend ----- #
index_html = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Code-to-Flowchart Generator</title>
    <style>
        body { font-family: sans-serif; background: #111827; color: #f3f4f6; margin: 0; padding: 20px; }
        .container { display: flex; gap: 20px; height: 90vh; }
        .panel { flex: 1; display: flex; flex-direction: column; background: #1f2937; padding: 15px; border-radius: 8px; }
        textarea { flex: 1; background: #111827; color: #34d399; border: 1px solid #374151; padding: 10px; font-family: monospace; resize: none; border-radius: 4px; }
        button { background: #059669; color: white; border: none; padding: 12px; margin-top: 10px; cursor: pointer; font-weight: bold; border-radius: 4px; }
        button:hover { background: #10b981; }
        button:disabled { background: #374151; cursor: not-allowed; }
        #flowchart-target { flex: 1; background: #ffffff; padding: 10px; border-radius: 4px; overflow: auto; display: flex; justify-content: center; align-items: start; }
    </style>
</head>
<body>
    <h2>Flowchart Transpiler</h2>
    <div class="container">
        <div class="panel">
            <h3>Source Code Input</h3>
            <textarea id="code-input" placeholder="Paste your code here..." spellcheck="false"></textarea>
            <button id="submit-btn">Generate Flowchart</button>
        </div>
        <div class="panel">
            <h3>Mermaid Flowchart Visualizer</h3>
            <div id="flowchart-target">
                <pre class="mermaid" id="mermaid-string">
                    graph TD
                    A[Paste Code] --> B[Click Generate]
                </pre>
            </div>
        </div>
    </div>

    <script type="module">
        import { Client } from "https://cdn.jsdelivr.net/npm/@gradio/client@1/dist/index.min.js";
        import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
        
        mermaid.initialize({ startOnLoad: true, theme: 'neutral' });

        // Instantiate the local Gradio application client dynamically
        const client = await Client.connect(window.location.origin);

        document.getElementById('submit-btn').addEventListener('click', async () => {
            const codeValue = document.getElementById('code-input').value;
            const targetDiv = document.getElementById('flowchart-target');
            const submitBtn = document.getElementById('submit-btn');

            if (!codeValue.trim()) {
                targetDiv.innerHTML = "<p style='color:red;'>Please input code first.</p>";
                return;
            }

            // Disable the button while a request is in flight so a slow CPU
            // generation can't be double-fired into a concurrent request.
            submitBtn.disabled = true;
            submitBtn.textContent = "Generating...";
            targetDiv.innerHTML = "Generating diagram...";

            let mermaidSyntax = "";
            try {
                // Call the @app.api function registered in python (name + param must match)
                const result = await client.predict("/generate_flowchart", { src_code: codeValue });
                mermaidSyntax = result.data[0];

                // Inject the raw string into a clean layout block and re-trigger parsing
                targetDiv.innerHTML = `<pre class="mermaid">${mermaidSyntax}</pre>`;
                await mermaid.run();

            } catch (error) {
                // On failure show the error AND the exact raw Mermaid we tried to render,
                // so a parse error can be diagnosed from the real output. textContent is
                // used for the raw string so newlines/special chars can't break the page.
                targetDiv.innerHTML = "<p style='color:red;'>Error during generation: " + error.message + "</p><p style='color:#111;font-weight:bold;text-align:left;'>Raw Mermaid output:</p>";
                const dbg = document.createElement("pre");
                dbg.style.color = "#111";
                dbg.style.whiteSpace = "pre-wrap";
                dbg.style.textAlign = "left";
                dbg.textContent = mermaidSyntax;
                targetDiv.appendChild(dbg);
            } finally {
                submitBtn.disabled = false;
                submitBtn.textContent = "Generate Flowchart";
            }
        });
    </script>
</body>
</html>
"""

# Load the custom HTML
# / takes precedent over default Blocks UI
@app.get("/")
def index():
    return HTMLResponse(index_html)

app.launch(share=True)