agents/flowchart_agent.py

"""
agents/flowchart_agent.py
--------------------------
Flowchart Agent for AutoDevAgent.

Generates a Mermaid.js flowchart diagram describing the logic flow
of the final working code. Only triggered when:
  - The code is longer than settings.flowchart_min_lines (default 20), OR
  - The user explicitly requested a flowchart

Mermaid.js diagrams are plain text that render as visual flowcharts
in the Gradio UI via gr.HTML() with the Mermaid CDN library.

Design:
  - Uses the fast model (8B) — flowchart generation is a simpler task
    than code writing and 8B is sufficient for diagram syntax.
  - Instructs the LLM to output only valid Mermaid flowchart syntax.
  - Validates that the output starts with "flowchart" or "graph" to
    catch cases where the LLM returned prose instead of a diagram.
  - Flowchart quality degrades for complex nested logic — documented
    in README as a known limitation. Best for functions and short scripts.

Usage:
    from agents.flowchart_agent import FlowchartAgent, should_generate_flowchart
    from pipeline.state import PipelineState

    if should_generate_flowchart(state, user_requested=False):
        agent = FlowchartAgent()
        updated = agent.run(state)
        print(updated["flowchart_mermaid"])
"""

import logging
from typing import Any

from langchain_groq import ChatGroq
from langchain_core.messages import SystemMessage, HumanMessage

from config import settings
from pipeline.state import (
    PipelineState,
    PipelineStatus,
    Language,
)

logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Prompt                                                              #
# ------------------------------------------------------------------ #

FLOWCHART_SYSTEM = """
You are a diagram generator. Describe the logical steps of the given code as a simple Mermaid flowchart.

Output ONLY raw Mermaid syntax — no explanation, no code fences, no markdown.

STRICT RULES — follow every rule exactly or the diagram will fail to render:

1. First line: graph TD     (exactly this)
2. Node IDs: single capital letters only — A, B, C, D ... (max 10 nodes)
3. Node shape: RECTANGLES ONLY — A["label text"]
   - NEVER use diamonds {} or round nodes ()
4. Labels: 2 to 8 plain English words only
   - FORBIDDEN characters: " ' ; : & < > | ( ) , . / * % = [ ] { }
   - No SQL keywords, no code syntax, no punctuation
   - Use words like "Check condition" or "Repeat for each item" to imply logic
5. Edges: plain arrows only — A --> B
   - NEVER add edge labels (no -->|Yes|)
   - NEVER loop back to an earlier node
   - Use forward arrows only (top to bottom)
6. Max 10 nodes total
7. If the code contains branches or loops, describe them as sequential steps (e.g., "Check base case", "Handle base case", "Recursive case", "Combine results")
8. CONNECTIVITY — the graph must be a single connected component:
   - Every node except A must have at least one incoming arrow
   - Every node except the final node must have at least one outgoing arrow
   - Every node must be reachable by following arrows from A
   - NEVER define a node that is not connected to the main flow from A

CORRECT example for "fibonacci with memoisation":
graph TD
    A["Start"]
    B["Base case check"]
    C["Return base value"]
    D["Check memo cache"]
    E["Return cached value"]
    F["Compute recursively"]
    G["Store in memo"]
    H["Return result"]
    A --> B
    B --> C
    B --> D
    D --> E
    D --> F
    F --> G
    G --> H
""".strip()

SQL_FLOWCHART_SYSTEM = """
You are a diagram generator. Describe the logical steps of the given SQL query as a simple Mermaid flowchart.

Output ONLY raw Mermaid syntax — no explanation, no code fences, no markdown.

STRICT RULES — follow every rule exactly or the diagram will fail to render:

1. First line: graph TD     (exactly this)
2. Node IDs: single capital letters only — A, B, C, D ... (max 10 nodes)
3. Node shape: RECTANGLES ONLY — A["label text"]
   - NEVER use diamonds {} or round nodes ()
4. Labels: 2 to 8 plain English words only
   - FORBIDDEN characters: " ' ; : & < > | ( ) , . / * % = [ ] { }
   - FORBIDDEN SQL keywords: SELECT FROM WHERE JOIN GROUP ORDER COUNT SUM AVG MAX MIN
   - Describe what the step DOES in plain English (e.g. "Filter recent orders", "Calculate running total")
   - For subqueries, use a step like "Prepare subquery results"
5. Edges: plain arrows only — A --> B
   - NEVER add edge labels
   - No duplicate edges, no loops
6. Max 10 nodes total
7. CONNECTIVITY — the graph must be a single connected component:
   - Every node except A must have at least one incoming arrow
   - Every node except the final node must have at least one outgoing arrow
   - Every node must be reachable by following arrows from A
   - NEVER define a node that is not connected to the main flow from A

CORRECT example for "top 3 products by sales":
graph TD
    A["Start"]
    B["Load products table"]
    C["Load sales table"]
    D["Join tables together"]
    E["Sum sales per product"]
    F["Sort by total descending"]
    G["Keep top three rows"]
    H["Return results"]
    A --> B
    B --> C
    C --> D
    D --> E
    E --> F
    F --> G
    G --> H
""".strip()

FLOWCHART_FIX_SYSTEM = """
You are fixing broken Mermaid.js syntax. Rewrite the diagram following these exact rules.

RULES (all mandatory — no exceptions):
1. First line must be exactly: graph TD
2. Node IDs: single capital letters A, B, C, D ...
3. RECTANGLES ONLY: A["plain label"]
   - Remove ALL diamond nodes {} — replace with rectangles
   - Remove ALL round nodes () — replace with rectangles
4. Labels: 2 to 6 plain English words. No special characters at all.
   - Remove: " ' ; : & < > | ( ) , . / * % = [ ] { }
5. Edges: A --> B only — no edge labels, no loops, no duplicate edges
6. Max 10 nodes
7. CONNECTIVITY — single connected component from A:
   - Every node except A must have at least one incoming arrow
   - Every node except the final node must have at least one outgoing arrow
   - Remove or reconnect any node that is not reachable from A

Output ONLY the corrected raw Mermaid syntax. No explanation. No code fences.
""".strip()


# ------------------------------------------------------------------ #
#  Gate function                                                       #
# ------------------------------------------------------------------ #

def should_generate_flowchart(
    state: PipelineState,
    user_requested: bool = False,
) -> bool:
    """
    Decide whether to generate a flowchart for the current code.

    SQL tasks: always generate — even a short query has a meaningful
    logical flow (filter → aggregate → sort) that benefits from a diagram.

    Python tasks: only generate when the code exceeds the configured
    minimum line threshold, or when the user explicitly requested one.

    Args:
        state:          Current PipelineState.
        user_requested: True if the user clicked "Generate flowchart".

    Returns:
        True if a flowchart should be generated, False otherwise.
    """
    if user_requested:
        return True

    if state.language == Language.SQL:
        logger.info("FlowchartAgent: triggering — SQL task always gets a flowchart")
        return True

    code = state.final_code()
    line_count = len(code.strip().splitlines())

    if line_count >= settings.flowchart_min_lines:
        logger.info(
            "FlowchartAgent: triggering — code is %d lines (threshold: %d)",
            line_count,
            settings.flowchart_min_lines,
        )
        return True

    logger.debug(
        "FlowchartAgent: skipping — code is %d lines (threshold: %d)",
        line_count,
        settings.flowchart_min_lines,
    )
    return False


# ------------------------------------------------------------------ #
#  Agent                                                               #
# ------------------------------------------------------------------ #

class FlowchartAgent:
    """
    Generates Mermaid.js flowchart syntax for the final code.

    Uses the fast 8B model since flowchart generation is a simpler
    structured output task than code generation or debugging.

    Attributes:
        llm: ChatGroq instance using the fast model (8B).
    """

    def __init__(self) -> None:
        """LLM client is built lazily in run() once model assignments are known."""
        self.llm = None

    def _build_llm(self, model: str) -> "ChatGroq":
        return ChatGroq(
            api_key=settings.groq_api_key,
            model=model,
            temperature=0.1,
            max_tokens=600,
            request_timeout=settings.groq_request_timeout,
        )

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Generate a Mermaid.js flowchart for the code in state.

        Args:
            state: Current PipelineState. Reads: task, language,
                   final_code().

        Returns:
            Partial state dict with keys:
              - "flowchart_mermaid": str — Mermaid diagram syntax
              - "status":            PipelineStatus.EXPLAINING
        """
        ma = state.model_assignments or {}
        model = ma.get("flowchart", settings.groq_model_fast)
        self.llm = self._build_llm(model)
        logger.info("FlowchartAgent using model: %s", model)

        code = state.final_code()

        if not code.strip():
            logger.warning("FlowchartAgent: no code to diagram")
            return {
                "flowchart_mermaid": "",
                "status":            PipelineStatus.EXPLAINING,
            }

        logger.info(
            "FlowchartAgent generating diagram for %d lines of %s code",
            len(code.splitlines()),
            state.language.value,
        )

        system_prompt = (
            SQL_FLOWCHART_SYSTEM
            if state.language == Language.SQL
            else FLOWCHART_SYSTEM
        )

        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=(
                f"Task: {state.task}\n\n"
                f"Code:\n{code}"
            )),
        ]

        # ── LLM call ──────────────────────────────────────────────── #
        try:
            response = self.llm.invoke(messages)
            raw      = response.content.strip()
        except Exception as e:
            logger.error("FlowchartAgent LLM call failed: %s", e)
            raise RuntimeError(f"Flowchart agent LLM call failed: {e}") from e

        # ── Layer 2: Strip fences + validate ─────────────────────── #
        mermaid = _clean_mermaid(raw)

        # ── Track token usage ─────────────────────────────────────── #
        from observability.langsmith_tracer import extract_token_usage_from_response
        prompt_t, completion_t = extract_token_usage_from_response(response)
        updated_token_usage = state.token_usage.model_copy()
        updated_token_usage.add(prompt_t, completion_t)

        if not _is_valid_mermaid(mermaid):
            logger.warning(
                "FlowchartAgent: output does not look like valid Mermaid syntax — discarding"
            )
            return {
                "flowchart_mermaid": "",
                "status":            PipelineStatus.EXPLAINING,
                "token_usage":       updated_token_usage,
            }

        is_valid, issues = _validate_mermaid_syntax(mermaid)

        # ── Layer 3: LLM retry with error feedback ────────────────── #
        if not is_valid:
            error_summary = "; ".join(issues)
            logger.warning(
                "FlowchartAgent: Layer 2 detected %d issue(s): %s — retrying with fix prompt",
                len(issues), error_summary,
            )
            try:
                fix_messages = [
                    SystemMessage(content=FLOWCHART_FIX_SYSTEM),
                    HumanMessage(content=(
                        f"Broken diagram (errors: {error_summary}):\n\n{mermaid}"
                    )),
                ]
                fix_response = self.llm.invoke(fix_messages)
                fixed_raw    = fix_response.content.strip()
                fixed_mermaid = _clean_mermaid(fixed_raw)

                # Count fix attempt tokens
                fix_pt, fix_ct = extract_token_usage_from_response(fix_response)
                updated_token_usage.add(fix_pt, fix_ct)

                fixed_valid, fixed_issues = _validate_mermaid_syntax(fixed_mermaid)
                if fixed_valid and _is_valid_mermaid(fixed_mermaid):
                    logger.info("FlowchartAgent: Layer 3 fix succeeded")
                    mermaid = fixed_mermaid
                else:
                    logger.warning(
                        "FlowchartAgent: Layer 3 fix still invalid (%s) — passing to Layer 4 rebuild",
                        "; ".join(fixed_issues),
                    )
                    mermaid = fixed_mermaid  # Layer 4 (parse→rebuild in app.py) handles it
            except Exception as fix_err:
                logger.warning("FlowchartAgent: Layer 3 retry failed: %s — passing to Layer 4", fix_err)

        logger.info(
            "FlowchartAgent produced %d line diagram", len(mermaid.splitlines())
        )

        return {
            "flowchart_mermaid": mermaid,
            "status":            PipelineStatus.EXPLAINING,
            "token_usage":       updated_token_usage,
        }


# ------------------------------------------------------------------ #
#  Helpers                                                             #
# ------------------------------------------------------------------ #

def _validate_mermaid_syntax(text: str) -> tuple[bool, list[str]]:
    """
    Layer 2 — Syntax validation.

    Enforces the strict simple-flowchart rules:
      1. Must start with 'graph' or 'flowchart'
      2. No diamond {} nodes (replaced by strict prompt, but catch LLM drift)
      3. No edge labels (|...|) — pipe chars inside edges are a common error
      4. No unbalanced brackets on node-definition lines
      5. Forbidden characters inside quoted labels
      6. At least one edge (-->) present
      7. No more than 12 nodes (safety cap — prompt says max 10)
    """
    import re
    issues: list[str] = []

    if not text or not text.strip():
        return False, ["Empty output"]

    lines = [l.strip() for l in text.splitlines() if l.strip()]

    # 1. Must start with graph / flowchart
    if not lines[0].lower().startswith(("graph", "flowchart")):
        issues.append(f"Missing graph header — first line is: '{lines[0]}'")

    # 2. Diamond nodes {} are forbidden — they cause Mermaid v11 syntax errors
    diamond_re = re.compile(r'\b[A-Za-z][A-Za-z0-9_]*\s*\{')
    for line in lines:
        if "-->" in line:
            continue
        if diamond_re.search(line):
            issues.append(f"Diamond node detected (forbidden): {line}")
            break

    # 3. Edge labels with pipes e.g. A -->|Yes| B — forbidden in strict mode
    edge_label_re = re.compile(r'--+>\s*\|')
    for line in lines:
        if edge_label_re.search(line):
            issues.append(f"Edge label with pipe detected (forbidden): {line}")
            break

    # 4. Unbalanced brackets on node-definition lines
    for line in lines:
        if "-->" in line or "---" in line:
            continue
        opens  = line.count("[") + line.count("{") + line.count("(")
        closes = line.count("]") + line.count("}") + line.count(")")
        if opens != closes:
            issues.append(f"Unbalanced brackets: {line}")
            break

    # 5. Forbidden chars inside quoted labels
    forbidden_re = re.compile(r'"[^"]*[;:&<>|][^"]*"')
    for line in lines:
        if forbidden_re.search(line):
            issues.append(f"Forbidden characters in label: {line}")
            break

    # 6. Must have at least one edge
    has_edge = any("-->" in l for l in lines)
    if not has_edge:
        issues.append("No edges found — diagram has no connections")

    # 7. Node count sanity check (safety cap)
    node_def_re = re.compile(r'^\s*[A-Za-z][A-Za-z0-9_]*\s*[\[({]')
    node_count = sum(1 for l in lines if node_def_re.match(l) and "-->" not in l)
    if node_count > 12:
        issues.append(f"Too many nodes ({node_count}) — max 12 allowed")

    # 8. Connectivity — every non-root node must have an incoming edge;
    #    every non-terminal node must have an outgoing edge; graph must
    #    be a single connected component reachable from A.
    edge_re = re.compile(r'\b([A-Za-z][A-Za-z0-9_]*)\s*--+>\s*([A-Za-z][A-Za-z0-9_]*)\b')
    defined_nodes: set[str] = set()
    for l in lines:
        if "-->" not in l:
            m = re.match(r'^\s*([A-Za-z][A-Za-z0-9_]*)\s*[\[({]', l)
            if m:
                defined_nodes.add(m.group(1))
    parsed_edges: list[tuple[str, str]] = []
    for l in lines:
        m = edge_re.search(l)
        if m:
            parsed_edges.append((m.group(1), m.group(2)))
    if defined_nodes and parsed_edges:
        sources  = {s for s, _ in parsed_edges}
        targets  = {d for _, d in parsed_edges}
        # Nodes with no incoming edge (except A which is the root)
        no_incoming = defined_nodes - targets - {"A"}
        if no_incoming:
            issues.append(f"Orphaned nodes with no incoming edge: {sorted(no_incoming)}")
        # BFS reachability from A
        adj: dict[str, list[str]] = {}
        for s, d in parsed_edges:
            adj.setdefault(s, []).append(d)
        reachable: set[str] = set()
        queue = ["A"]
        while queue:
            cur = queue.pop(0)
            if cur in reachable:
                continue
            reachable.add(cur)
            for nb in adj.get(cur, []):
                queue.append(nb)
        unreachable = defined_nodes - reachable
        if unreachable:
            issues.append(f"Nodes not reachable from A: {sorted(unreachable)}")

    return len(issues) == 0, issues


def _clean_mermaid(raw: str) -> str:
    """
    Strip markdown fences from Mermaid output.

    LLMs frequently wrap Mermaid diagrams in ```mermaid fences even
    when instructed not to. The Gradio renderer needs raw syntax.

    Args:
        raw: Raw LLM response string.

    Returns:
        Clean Mermaid syntax string.
    """
    stripped = raw.strip()

    if stripped.startswith("```"):
        lines = stripped.splitlines()
        end   = len(lines) - 1 if lines[-1].strip() == "```" else len(lines)
        return "\n".join(lines[1:end]).strip()

    return stripped


def _is_valid_mermaid(text: str) -> bool:
    """
    Basic validation that the output looks like Mermaid flowchart syntax.

    Checks that the text starts with a known Mermaid diagram keyword.
    This catches cases where the LLM returned a prose explanation
    instead of diagram syntax.

    Args:
        text: Cleaned Mermaid output string.

    Returns:
        True if it looks like valid Mermaid syntax, False otherwise.
    """
    if not text:
        return False

    first_line = text.splitlines()[0].strip().lower()

    valid_starters = (
        "flowchart",
        "graph td",
        "graph lr",
        "graph tb",
        "graph rl",
        "graph bt",
        "sequencediagram",
        "classDiagram",
    )

    return any(first_line.startswith(s) for s in valid_starters)