agents/debug_agent.py

"""
agents/debug_agent.py
---------------------
Debug Agent for AutoDevAgent.

Responsible for two things per debug iteration:
  1. Error Classification — categorises the error as syntax / runtime /
     logic / timeout using a fast cheap LLM call (Llama 3.1 8B).
  2. Self-Reflection + Rewrite — produces structured reasoning in the
     form "I saw X. I believe the cause is Y. I will change Z." then
     rewrites the code based on that reasoning.

Also handles the Human-in-the-Loop escalation path: when max retries
is hit, generates 2–3 targeted fix options for the user to choose from.

Design:
  - Error classification uses the fast model (8B) — it's a simple
    classification task that doesn't need 70B quality.
  - Rewrite uses the primary model (70B) — code quality matters here.
  - The error cache (utils/error_cache.py) is consulted before
    classification. If the same error fingerprint appeared in the
    previous iteration, a different fix strategy is forced.
  - Self-reflection is stored in state.reflections_history so the UI
    can display the full reasoning trail across all iterations.

Usage:
    from agents.debug_agent import DebugAgent
    from pipeline.state import PipelineState, Language

    agent = DebugAgent()
    # state must have generated_code and execution_result populated
    updated = agent.run(state)
    print(updated["self_reflection"].what_i_saw)
    print(updated["generated_code"])
"""

import json
import logging
from typing import Any

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

from config import settings
from pipeline.state import (
    PipelineState,
    PipelineStatus,
    ErrorClassification,
    ErrorType,
    SelfReflection,
    Language,
)
from utils.error_cache import ErrorCache

logger = logging.getLogger(__name__)

# Module-level cache shared across all DebugAgent instances in a session
_error_cache = ErrorCache()


# ------------------------------------------------------------------ #
#  Prompts                                                             #
# ------------------------------------------------------------------ #

CLASSIFY_SYSTEM = """
You are a code error classifier. Classify the error into exactly one category.

Categories:
  - syntax:  The code has invalid syntax that prevented it from running at all.
  - runtime: The code ran but crashed due to an exception (NameError, TypeError, etc.).
  - logic:   The code ran without exceptions but produced wrong output.
  - timeout: The code exceeded the time limit — likely an infinite loop.

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "error_type": "runtime",
  "root_cause": "One sentence describing the exact cause",
  "suggested_fix": "One concrete change to fix it",
  "cache_key": "normalised short fingerprint of the error, e.g. NameError:x_not_defined"
}
""".strip()

REFLECT_AND_REWRITE_SYSTEM = """
You are a senior software engineer debugging code.
You will be given:
  - The original task
  - The current (broken) code
  - The error classification and suggested fix
  - Your previous self-reflection (if any)

Step 1 — Self-reflection: Think out loud about what went wrong.
Step 2 — Rewrite: Write the corrected code.

CRITICAL — __main__ block rules:
- The fixed_code __main__ block MUST only use valid, realistic, non-None inputs.
- NEVER call functions with None, "", [], or invalid/edge-case values in __main__.
- Handle None/edge cases INSIDE the function body, but NEVER demonstrate them in __main__.
- WRONG:  print(are_anagrams(None, "listen"))   # crashes the executor
- CORRECT: print(are_anagrams("listen", "silent"))

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "what_i_saw": "The exact error or wrong output I observed",
  "what_i_think": "My hypothesis for the root cause",
  "what_i_will_do": "The specific change I will make",
  "fixed_code": "the full corrected code here — raw code, no fences"
}
""".strip()

REFLECT_AND_REWRITE_SQL_SYSTEM = """
You are a senior data engineer debugging a SQL query.
You will be given:
  - The original task
  - The current (broken) SQL query
  - The database schema (CREATE TABLE statements)
  - The dummy data currently in the database (INSERT statements)
  - The error or test failure details
  - Your previous self-reflection (if any)

IMPORTANT: The dummy data is what is actually in the database when the query runs.
If the query returns 0 rows, check whether the dummy data satisfies the WHERE conditions.
For example, if the query filters by date range "last 30 days" but all INSERT dates are
hardcoded old dates (e.g. '2023-01-15'), the query is correct but the data is wrong.
In that case, FIX THE QUERY to work with the available data
(e.g. remove the date filter or use a different date range that matches the data).

Step 1 — Self-reflection: Think out loud about what went wrong (data issue vs query issue).
Step 2 — Rewrite: Write the corrected SQL that will return correct results with the given data.

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "what_i_saw": "The exact error or wrong output I observed",
  "what_i_think": "My hypothesis — is this a query error or a data/filter mismatch?",
  "what_i_will_do": "The specific change I will make to the SQL query",
  "fixed_code": "the full corrected SQL query here — raw SQL, no fences"
}
""".strip()

HITL_OPTIONS_SYSTEM = """
You are a debugging assistant. The automated debug loop has failed after multiple attempts.
You will be shown the task, the broken code, and the last error.

Generate 2–3 concrete, targeted options that a human developer could choose from to fix the issue.
Each option should be a specific, actionable suggestion — not generic advice.

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "options": [
    "Option 1: specific suggestion here",
    "Option 2: specific suggestion here",
    "Option 3: specific suggestion here"
  ]
}
""".strip()


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

class DebugAgent:
    """
    Classifies errors, reflects on them, and rewrites code.

    Each call to run() performs one full debug iteration:
      1. Classify the error (fast model)
      2. Reflect and rewrite (primary model)
      3. Increment debug_iterations
      4. If max retries is hit next, generate HITL options

    Attributes:
        llm_fast:    ChatGroq using Llama 3.1 8B for classification.
        llm_primary: ChatGroq using Llama 3.1 70B for rewriting.
    """

    def __init__(self) -> None:
        """LLM clients are built lazily in run() once model assignments are known."""
        self.llm_fast    = None
        self.llm_primary = None

    def _build_llms(self, fast_model: str, primary_model: str) -> None:
        self.llm_fast = ChatGroq(
            api_key=settings.groq_api_key,
            model=fast_model,
            temperature=0.0,
            max_tokens=400,
            request_timeout=settings.groq_request_timeout,
        )
        self.llm_primary = ChatGroq(
            api_key=settings.groq_api_key,
            model=primary_model,
            temperature=0.2,
            max_tokens=2500,
            request_timeout=settings.groq_request_timeout,
        )
        # Accumulated token counts across all LLM calls in one run()
        self._prompt_tokens:     int = 0
        self._completion_tokens: int = 0

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Run one debug iteration on the current pipeline state.

        Classifies the error, checks the cache for repeated errors,
        reflects on the cause, rewrites the code, and updates state.

        Args:
            state: Current PipelineState. Reads: task, language,
                   generated_code, execution_result, debug_iterations,
                   reflections_history, error_classification.

        Returns:
            Partial state dict with keys:
              - "error_classification": ErrorClassification
              - "self_reflection":      SelfReflection
              - "generated_code":       str  (rewritten code)
              - "code_history":         list (all versions)
              - "debug_iterations":     int  (incremented)
              - "reflections_history":  list (all reflections)
              - "status":               PipelineStatus
              - "hitl_options":         list (only if max retries hit)
        """
        ma = state.model_assignments or {}
        self._build_llms(
            fast_model    = ma.get("classifier", settings.groq_model_fast),
            primary_model = ma.get("debugger",   settings.groq_model_primary),
        )
        logger.info(
            "DebugAgent using classifier=%s  debugger=%s",
            ma.get("classifier", settings.groq_model_fast),
            ma.get("debugger",   settings.groq_model_primary),
        )

        iteration = state.debug_iterations + 1
        logger.info("DebugAgent starting iteration %d/%d", iteration, settings.max_debug_retries)

        # Reset per-run token accumulators
        self._prompt_tokens     = 0
        self._completion_tokens = 0

        error_msg   = state.latest_error()
        broken_code = state.generated_code

        # ── Step 1: Classify the error ────────────────────────────── #
        classification = self._classify_error(error_msg, broken_code)

        # ── Step 2: Check error cache ─────────────────────────────── #
        # If we've seen this exact error before, force a different strategy
        if _error_cache.is_repeated(classification.cache_key):
            logger.info(
                "Repeated error detected (%s) — forcing alternative fix strategy",
                classification.cache_key,
            )
            classification = _force_alternative_strategy(classification)
        _error_cache.record(classification.cache_key)

        # ── Step 3: Reflect and rewrite ───────────────────────────── #
        reflection, fixed_code = self._reflect_and_rewrite(state, classification)

        # ── Step 4: Update history ────────────────────────────────── #
        updated_history     = list(state.code_history) + [fixed_code]
        updated_reflections = list(state.reflections_history) + [reflection]
        new_iteration_count = iteration

        logger.info(
            "DebugAgent iteration %d complete — reflection: '%s'",
            iteration,
            reflection.what_i_will_do[:80],
        )

        # ── Step 5: HITL options if this was the last allowed retry ── #
        hitl_options: list[str] = []
        if new_iteration_count >= settings.max_debug_retries:
            logger.info("Max retries reached — generating HITL options")
            hitl_options = self._generate_hitl_options(state, fixed_code, error_msg)

        # ── Accumulate token usage across all sub-calls ───────────── #
        updated_token_usage = state.token_usage.model_copy()
        updated_token_usage.add(self._prompt_tokens, self._completion_tokens)

        return {
            "error_classification":  classification,
            "self_reflection":       reflection,
            "generated_code":        fixed_code,
            "code_history":          updated_history,
            "debug_iterations":      new_iteration_count,
            "reflections_history":   updated_reflections,
            "hitl_options":          hitl_options,
            "token_usage":           updated_token_usage,
            "status": (
                PipelineStatus.AWAITING_HUMAN
                if new_iteration_count >= settings.max_debug_retries
                else (
                    PipelineStatus.TEST_DEBUGGING
                    if getattr(state, "failed_at_test", False)
                    else PipelineStatus.DEBUGGING
                )
            ),
        }

    # ---------------------------------------------------------------- #
    #  Private helpers                                                  #
    # ---------------------------------------------------------------- #

    def _classify_error(
        self,
        error_msg: str,
        broken_code: str,
    ) -> ErrorClassification:
        """
        Classify the error using the fast LLM (Llama 3.1 8B).

        Args:
            error_msg:   The cleaned error string from the executor.
            broken_code: The code that produced the error.

        Returns:
            ErrorClassification with type, root cause, and cache key.
        """
        human_content = (
            f"Error output:\n{error_msg}\n\n"
            f"Code:\n{broken_code[:1500]}"
        )

        messages = [
            SystemMessage(content=CLASSIFY_SYSTEM),
            HumanMessage(content=human_content),
        ]

        try:
            response = self.llm_fast.invoke(messages)
            raw      = response.content.strip()
            from observability.langsmith_tracer import extract_token_usage_from_response
            p, c = extract_token_usage_from_response(response)
            self._prompt_tokens += p
            self._completion_tokens += c
            parsed   = _parse_json(raw, "DebugAgent.classify")

            return ErrorClassification(
                error_type    = ErrorType(parsed.get("error_type", "unknown")),
                root_cause    = parsed.get("root_cause", "Unknown root cause"),
                suggested_fix = parsed.get("suggested_fix", "Review the error and fix manually"),
                cache_key     = parsed.get("cache_key", error_msg[:80]),
            )

        except Exception as e:
            logger.warning("Error classification failed: %s — defaulting to UNKNOWN", e)
            return ErrorClassification(
                error_type    = ErrorType.UNKNOWN,
                root_cause    = "Classification failed",
                suggested_fix = "Review the error message and fix manually",
                cache_key     = error_msg[:80],
            )

    def _reflect_and_rewrite(
        self,
        state: PipelineState,
        classification: ErrorClassification,
    ) -> tuple[SelfReflection, str]:
        """
        Produce structured self-reflection and rewrite the code.

        Args:
            state:          Current PipelineState.
            classification: Output of the error classifier.

        Returns:
            Tuple of (SelfReflection, fixed_code_string).
        """
        # Build context including previous reflection if available
        prev_reflection = ""
        if state.self_reflection:
            prev_reflection = (
                f"\nPrevious reflection:\n"
                f"  Saw:      {state.self_reflection.what_i_saw}\n"
                f"  Thought:  {state.self_reflection.what_i_think}\n"
                f"  Did:      {state.self_reflection.what_i_will_do}\n"
                "(That fix did not work — try a different approach.)"
            )

        # For SQL: use specialised prompt and include schema + data so the
        # agent can distinguish "query is wrong" from "data doesn't match filter"
        is_sql = (state.language == Language.SQL)

        if is_sql and state.sql_schema:
            schema = state.sql_schema
            create_block  = "\n".join(schema.create_statements)
            insert_block  = "\n".join(schema.insert_statements[:15])  # cap at 15 rows
            desc_block    = "\n".join(schema.table_descriptions) if schema.table_descriptions else ""
            sql_schema_context = (
                f"\n\nDatabase schema (CREATE TABLE):\n{create_block}\n\n"
                f"Dummy data currently in DB (INSERT statements):\n{insert_block}\n"
                + (f"\nTable descriptions:\n{desc_block}" if desc_block else "")
            )
        else:
            sql_schema_context = ""

        system_prompt   = REFLECT_AND_REWRITE_SQL_SYSTEM if is_sql else REFLECT_AND_REWRITE_SYSTEM
        code_label      = "Broken SQL query" if is_sql else "Broken code"

        human_content = (
            f"Task: {state.task}\n\n"
            f"{code_label}:\n{state.generated_code}\n\n"
            f"Error type: {classification.error_type.value}\n"
            f"Root cause: {classification.root_cause}\n"
            f"Suggested fix: {classification.suggested_fix}\n"
            f"{sql_schema_context}"
            f"{prev_reflection}\n\n"
            f"Now reflect and rewrite the {'SQL query' if is_sql else 'code'}:"
        )

        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=human_content),
        ]

        try:
            response = self.llm_primary.invoke(messages)
            raw      = response.content.strip()
            from observability.langsmith_tracer import extract_token_usage_from_response
            p, c = extract_token_usage_from_response(response)
            self._prompt_tokens += p
            self._completion_tokens += c
            parsed   = _parse_json(raw, "DebugAgent.reflect")

            reflection = SelfReflection(
                what_i_saw     = parsed.get("what_i_saw", "Error observed"),
                what_i_think   = parsed.get("what_i_think", "Unknown cause"),
                what_i_will_do = parsed.get("what_i_will_do", "Rewrite the code"),
            )

            fixed_code = parsed.get("fixed_code", "").strip()
            if not fixed_code:
                raise ValueError("DebugAgent returned empty fixed_code")

            # Strip any accidental fences inside the JSON value
            if fixed_code.startswith("```"):
                lines = fixed_code.splitlines()
                fixed_code = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])

            return reflection, fixed_code.strip()

        except Exception as e:
            logger.error("DebugAgent reflect+rewrite failed: %s", e)
            # Fallback: return a safe reflection and the original code
            fallback_reflection = SelfReflection(
                what_i_saw     = state.latest_error()[:200],
                what_i_think   = "The rewrite agent encountered an error itself",
                what_i_will_do = "Return original code — manual review needed",
            )
            return fallback_reflection, state.generated_code

    def _generate_hitl_options(
        self,
        state: PipelineState,
        last_code: str,
        last_error: str,
    ) -> list[str]:
        """
        Generate 2–3 targeted fix options for the human-in-the-loop panel.

        Called only when max_debug_retries is reached.

        Args:
            state:      Current PipelineState.
            last_code:  Most recent code after all debug attempts.
            last_error: Most recent error message.

        Returns:
            List of 2–3 option strings for the UI to display.
        """
        human_content = (
            f"Task: {state.task}\n\n"
            f"Code after {settings.max_debug_retries} debug attempts:\n{last_code}\n\n"
            f"Last error:\n{last_error}"
        )

        messages = [
            SystemMessage(content=HITL_OPTIONS_SYSTEM),
            HumanMessage(content=human_content),
        ]

        try:
            response = self.llm_primary.invoke(messages)
            raw      = response.content.strip()
            from observability.langsmith_tracer import extract_token_usage_from_response
            p, c = extract_token_usage_from_response(response)
            self._prompt_tokens += p
            self._completion_tokens += c
            parsed   = _parse_json(raw, "DebugAgent.hitl")
            options  = parsed.get("options", [])
            if not options:
                raise ValueError("No options returned")
            return [str(o) for o in options[:3]]

        except Exception as e:
            logger.warning("HITL option generation failed: %s", e)
            return [
                "Option 1: Review the error message and edit the code manually below.",
                "Option 2: Simplify the task description and resubmit.",
                "Option 3: Switch to a different approach entirely.",
            ]


# ------------------------------------------------------------------ #
#  Module-level helpers                                                #
# ------------------------------------------------------------------ #

def _parse_json(raw: str, agent: str) -> dict:
    """
    Safely parse a JSON string from the LLM, stripping markdown fences.

    Args:
        raw:   Raw LLM response string.
        agent: Caller name for logging.

    Returns:
        Parsed dict.

    Raises:
        ValueError: If JSON cannot be parsed.
    """
    cleaned = raw
    if cleaned.startswith("```"):
        cleaned = cleaned.split("\n", 1)[-1]
    if cleaned.endswith("```"):
        cleaned = cleaned.rsplit("```", 1)[0]
    cleaned = cleaned.strip()

    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as e:
        logger.error("%s could not parse JSON: %s", agent, e)
        raise ValueError(f"{agent} returned invalid JSON: {e}") from e


def _force_alternative_strategy(
    classification: ErrorClassification,
) -> ErrorClassification:
    """
    When the same error is repeated, escalate the suggested fix strategy.

    Instead of the original suggestion, instruct the debug agent to take
    a more radical approach — restructure the logic rather than patch it.

    Args:
        classification: The original ErrorClassification.

    Returns:
        Updated ErrorClassification with a stronger suggested_fix.
    """
    escalated_fix = (
        f"Previous fix attempt did not resolve this error. "
        f"Original suggestion was: '{classification.suggested_fix}'. "
        f"Now try a fundamentally different approach — consider restructuring "
        f"the logic, using a different algorithm, or simplifying the code significantly."
    )
    return ErrorClassification(
        error_type    = classification.error_type,
        root_cause    = classification.root_cause,
        suggested_fix = escalated_fix,
        cache_key     = classification.cache_key,
    )