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

To do
- create the custom gradio look

"""
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="Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
    filename="qwen2.5-coder-7b-instruct-q4_k_m.gguf"
)

# Initialize llama.cpp with the local cached path
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2
)

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

# ----- Functions ----- #
@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: Keep text inside the flowchart nodes descriptive but concise, under 5 words.
    </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()

    # 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)
    # Mermaid can't parse double quotes inside [node labels] (e.g. C[Return "Active"]);
    # source-code string literals leak them, so downgrade to single quotes which parse fine.
    cleaned = cleaned.replace('"', "'")
    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; }
        #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..."></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');
            
            if (!codeValue.trim()) {
                targetDiv.innerHTML = "<p style='color:red;'>Please input code first.</p>";
                return;
            }

            targetDiv.innerHTML = "Generating diagram...";

            try {
                // Call the @app.api function registered in python (name + param must match)
                const result = await client.predict("/generate_flowchart", { src_code: codeValue });
                const 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) {
                targetDiv.innerHTML = `<p style='color:red;'>Error during generation: ${error.message}</p>`;
            }
        });
    </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)