agents/code_generator.py

"""
agents/code_generator.py
------------------------
Code Generator Agent for AutoDevAgent.

Takes the plan produced by the Planning Agent and generates executable
code from it. This is a pure code-writing step — no execution happens here.

Design:
  - Uses Groq's primary model (Llama 3.1 70B) for code quality.
  - Receives the full plan as context so the LLM writes code
    that matches the reasoning, not just the task description.
  - For SQL: also receives the inferred schema so the generated
    query targets the exact tables and columns that exist in SQLite.
  - Strips markdown fences from the output — executors need raw code.
  - Appends each generated version to code_history for full traceability.

Usage:
    from agents.code_generator import CodeGeneratorAgent
    from pipeline.state import PipelineState, Language

    agent = CodeGeneratorAgent()
    state = PipelineState(task="reverse a string", language=Language.PYTHON)
    # state.plan must already be populated by PlanningAgent
    updated = agent.run(state)
    print(updated["generated_code"])
"""

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,
    PlanStep,
    SQLSchema,
    SelfReflection,
    ErrorClassification,
)

logger = logging.getLogger(__name__)


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

PYTHON_CODE_SYSTEM = """
You are an expert Python developer.
You will be given a task and a step-by-step plan.
Write clean, correct Python code that implements the plan exactly.

Rules:
- Output ONLY the raw Python code. No markdown, no explanation, no code fences.
- Follow the plan steps in order.
- Include a main guard (if __name__ == "__main__":) with 1-2 simple, realistic example calls.
- Add a brief docstring to every function.
- Use descriptive variable names.
- Handle obvious edge cases (empty input, None, etc.) INSIDE the function body.

CRITICAL — __main__ block rules:
- ONLY use valid, non-None real-world example inputs in the __main__ block.
- NEVER call the function with None, empty string "", [], or any invalid/edge-case value in __main__.
- The __main__ block is for demonstration only — show the function working correctly with typical inputs.
- WRONG:  print(are_anagrams(None, "listen"))   # crashes executor
- CORRECT: print(are_anagrams("listen", "silent"))  # works fine
""".strip()

SQL_CODE_SYSTEM = """
You are an expert SQL developer working with SQLite.
You will be given a task, a step-by-step plan, and the exact schema with tables and columns.

Write a single SQL query that solves the task using ONLY the tables and columns defined in the schema.

Rules:
- Output ONLY the raw SQL query. No markdown, no explanation, no code fences.
- Use only SQLite-compatible SQL syntax.
- Reference only the exact table and column names from the schema.
- End the query with a semicolon.
""".strip()

PYTHON_CODE_REGEN_SYSTEM = """
You are an expert Python developer doing a complete rewrite.
Previous attempts to fix this code have all failed. You must start from scratch with a fresh approach.

You will be given:
  - The task description
  - The step-by-step plan
  - A history of what was tried and what errors occurred

Your job is to write completely new code that avoids ALL the errors listed in the history.
Do NOT patch or extend the previous code — write a different implementation from scratch.

Rules:
- Output ONLY the raw Python code. No markdown, no explanation, no code fences.
- Choose a different algorithm or implementation strategy than what was tried before.
- Include a main guard (if __name__ == "__main__":) with 1-2 simple, realistic example calls.
- ONLY use valid, non-None, real-world inputs in the __main__ block.
- Handle all edge cases mentioned in the error history INSIDE the function body.
- Add a brief docstring to every function.
""".strip()

SQL_CODE_REGEN_SYSTEM = """
You are an expert SQL developer doing a complete rewrite.
Previous attempts to fix this query have all failed. You must start from scratch with a fresh approach.

You will be given:
  - The task description
  - The schema (CREATE TABLE + INSERT statements)
  - A history of what was tried and what errors occurred

Your job is to write a completely new SQL query that avoids ALL the errors listed in the history.
Do NOT patch the previous query — write a different approach from scratch.

Rules:
- Output ONLY the raw SQL query. No markdown, no explanation, no code fences.
- Use only SQLite-compatible SQL syntax.
- Reference only the exact table and column names from the schema.
- End the query with a semicolon.
""".strip()


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

class CodeGeneratorAgent:
    """
    Generates executable code from a structured plan.

    Reads the plan (and SQL schema if applicable) from state and
    produces clean, runnable code. Appends the result to code_history
    so every version is traceable across debug iterations.

    Attributes:
        llm: ChatGroq instance using the primary model (70B).
    """

    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.2,
            max_tokens=2000,
            request_timeout=settings.groq_request_timeout,
        )

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Generate code from the plan in the current pipeline state.

        Builds a human message containing the task, the full plan,
        and (for SQL) the schema. Sends to the LLM, strips any
        markdown fences from the response, and returns the clean code.

        Args:
            state: Current PipelineState. Reads: task, language,
                   plan, sql_schema.

        Returns:
            Partial state dict with keys:
              - "generated_code": str  — raw, executable code
              - "code_history":   list — all versions including this one
              - "status":         PipelineStatus.GENERATING
        """
        ma = state.model_assignments or {}
        model = ma.get("generator", settings.groq_model_primary)
        self.llm = self._build_llm(model)
        logger.info("CodeGeneratorAgent using model: %s", model)

        logger.info(
            "CodeGeneratorAgent running (iteration %d)", state.debug_iterations
        )

        system_prompt = (
            SQL_CODE_SYSTEM
            if state.language == Language.SQL
            else PYTHON_CODE_SYSTEM
        )

        human_content = _build_human_message(
            task=state.task,
            plan=state.plan,
            language=state.language,
            sql_schema=state.sql_schema,
        )

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

        # ── LLM call ──────────────────────────────────────────────── #
        try:
            response = self.llm.invoke(messages)
            raw = response.content.strip()
            logger.debug("CodeGeneratorAgent raw response length: %d chars", len(raw))
        except Exception as e:
            logger.error("CodeGeneratorAgent LLM call failed: %s", e)
            raise RuntimeError(f"Code generator LLM call failed: {e}") from e

        # ── 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)

        # ── Strip markdown fences ─────────────────────────────────── #
        code = _strip_code_fences(raw)

        if not code.strip():
            raise ValueError("CodeGeneratorAgent returned empty code.")

        logger.info(
            "CodeGeneratorAgent produced %d lines of %s code",
            len(code.splitlines()),
            state.language.value,
        )

        # ── Append to history ─────────────────────────────────────── #
        # Preserve all previous versions for traceability
        updated_history = list(state.code_history) + [code]

        return {
            "generated_code": code,
            "code_history":   updated_history,
            "status":         PipelineStatus.GENERATING,
            "token_usage":    updated_token_usage,
        }


    def regen_run(self, state: PipelineState) -> dict[str, Any]:
        """
        Fresh regeneration run — called when the debug loop has exhausted
        all retries without fixing the code.

        Unlike run(), which generates from the plan alone, regen_run()
        passes the full error history to the LLM and asks it to write a
        completely different implementation that avoids all known failures.

        Args:
            state: Current PipelineState. Reads: task, language, plan,
                   sql_schema, reflections_history, error_classification.

        Returns:
            Partial state dict with keys:
              - "generated_code": str  — fresh code, different approach
              - "code_history":   list — all versions including this one
              - "status":         PipelineStatus.REGENERATING
        """
        ma = state.model_assignments or {}
        model = ma.get("generator", settings.groq_model_primary)
        self.llm = self._build_llm(model)
        logger.info(
            "CodeGeneratorAgent.regen_run — regen #%d, model: %s",
            state.regen_count + 1, model,
        )

        system_prompt = (
            SQL_CODE_REGEN_SYSTEM
            if state.language == Language.SQL
            else PYTHON_CODE_REGEN_SYSTEM
        )

        human_content = _build_regen_human_message(
            task=state.task,
            plan=state.plan,
            language=state.language,
            sql_schema=state.sql_schema,
            reflections_history=state.reflections_history,
            error_classification=state.error_classification,
        )

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

        try:
            response = self.llm.invoke(messages)
            raw = response.content.strip()
        except Exception as e:
            logger.error("CodeGeneratorAgent.regen_run LLM call failed: %s", e)
            raise RuntimeError(f"Code regeneration LLM call failed: {e}") from e

        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)

        code = _strip_code_fences(raw)
        if not code.strip():
            raise ValueError("CodeGeneratorAgent.regen_run returned empty code.")

        logger.info(
            "CodeGeneratorAgent.regen_run produced %d lines of %s code",
            len(code.splitlines()), state.language.value,
        )

        updated_history = list(state.code_history) + [code]

        return {
            "generated_code": code,
            "code_history":   updated_history,
            "status":         PipelineStatus.REGENERATING,
            "token_usage":    updated_token_usage,
        }


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

def _build_human_message(
    task: str,
    plan: list[PlanStep],
    language: Language,
    sql_schema: SQLSchema | None,
) -> str:
    """
    Construct the human message sent to the LLM.

    Combines the task description, formatted plan steps, and (for SQL)
    the full schema into a single prompt string.

    Args:
        task:       The user's original task description.
        plan:       List of PlanStep objects from the Planning Agent.
        language:   The target language (PYTHON or SQL).
        sql_schema: Inferred schema for SQL tasks, else None.

    Returns:
        Formatted string ready to send as a HumanMessage.
    """
    lines: list[str] = [f"Task: {task}", "", "Plan:"]

    for step in plan:
        lines.append(f"  Step {step.step_number}: {step.description}")
        lines.append(f"    Reason: {step.reasoning}")

    # Append schema context for SQL tasks so the LLM uses correct table names
    if language == Language.SQL and sql_schema:
        lines.append("")
        lines.append("Database schema (SQLite):")
        for stmt in sql_schema.create_statements:
            lines.append(f"  {stmt}")
        if sql_schema.table_descriptions:
            lines.append("")
            lines.append("Table descriptions:")
            for desc in sql_schema.table_descriptions:
                lines.append(f"  - {desc}")

    lines.append("")
    lines.append("Write the code now:")
    return "\n".join(lines)


def _build_regen_human_message(
    task: str,
    plan: list[PlanStep],
    language: Language,
    sql_schema: SQLSchema | None,
    reflections_history: list,
    error_classification,
) -> str:
    """
    Build the human message for a fresh regeneration call.

    Includes the task, plan, schema (SQL), and the full history of
    what was tried and what went wrong so the LLM can deliberately
    avoid repeating the same mistakes.

    Args:
        task:                 The user's original task description.
        plan:                 Plan steps from PlanningAgent.
        language:             PYTHON or SQL.
        sql_schema:           Schema for SQL tasks, else None.
        reflections_history:  All SelfReflection objects across all debug iterations.
        error_classification: The last ErrorClassification (may be None).

    Returns:
        Formatted string ready to send as a HumanMessage.
    """
    lines: list[str] = [f"Task: {task}", "", "Plan:"]
    for step in plan:
        lines.append(f"  Step {step.step_number}: {step.description}")
        lines.append(f"    Reason: {step.reasoning}")

    # Schema for SQL
    if language == Language.SQL and sql_schema:
        lines.append("")
        lines.append("Database schema (SQLite):")
        for stmt in sql_schema.create_statements:
            lines.append(f"  {stmt}")
        lines.append("")
        lines.append("Dummy data:")
        for stmt in sql_schema.insert_statements:
            lines.append(f"  {stmt}")

    # Error history — most important context
    lines.append("")
    lines.append("Previous attempts all failed. Here is what was tried and what went wrong:")
    lines.append("(DO NOT repeat any of these approaches — use a completely different strategy)")

    if reflections_history:
        for i, r in enumerate(reflections_history, 1):
            lines.append(f"\n  Attempt {i}:")
            lines.append(f"    Error observed:  {r.what_i_saw}")
            lines.append(f"    Root cause:      {r.what_i_think}")
            lines.append(f"    Fix attempted:   {r.what_i_will_do}")
    else:
        lines.append("  (no reflection history available — write a robust implementation from scratch)")

    if error_classification:
        lines.append("")
        lines.append(f"Last error type: {error_classification.error_type}")
        lines.append(f"Last error hint: {error_classification.suggested_fix}")

    lines.append("")
    lines.append("Write completely new code now (different approach, no patches):")
    return "\n".join(lines)


def _strip_code_fences(raw: str) -> str:
    """
    Remove markdown code fences from LLM output.

    LLMs frequently wrap code in triple backtick fences (```python ...```)
    even when the prompt says not to. Raw executors need clean code.

    Handles these patterns:
      - ```python\\n...\\n```
      - ```sql\\n...\\n```
      - ```\\n...\\n```
      - No fences at all (returned as-is)

    Args:
        raw: Raw string from the LLM response.

    Returns:
        Clean code string with fences and language tags removed.
    """
    stripped = raw.strip()

    # Check if the response is wrapped in code fences
    if stripped.startswith("```"):
        lines = stripped.splitlines()
        # Remove the opening fence line (```python, ```sql, or just ```)
        start = 1
        # Remove the closing fence if present
        end = len(lines)
        if lines[-1].strip() == "```":
            end = len(lines) - 1
        return "\n".join(lines[start:end]).strip()

    return stripped