agents/planning_agent.py

"""
agents/planning_agent.py
------------------------
Planning Agent for AutoDevAgent.

Responsible for two things:
  1. Breaking the user's task into a structured, step-by-step plan
     before any code is written. This is the ReAct "Reason" step —
     the agent thinks before it acts.
  2. For SQL tasks: inferring the database schema (tables, columns,
     dummy data) needed to execute and test the query.

Design:
  - Uses Groq's primary model (Llama 3.1 70B) for quality reasoning.
  - Returns fully typed Pydantic models — no raw strings leave this module.
  - Prompt instructs the LLM to respond in JSON only, which we parse
    into PlanStep / SQLSchema models via Pydantic validation.

Usage:
    from agents.planning_agent import PlanningAgent
    from pipeline.state import PipelineState, Language

    agent = PlanningAgent()
    state = PipelineState(task="reverse a string", language=Language.PYTHON)
    updated = agent.run(state)
    print(updated["plan"])        # List[PlanStep]
    print(updated["sql_schema"])  # SQLSchema | None
"""

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

logger = logging.getLogger(__name__)


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

PYTHON_PLAN_SYSTEM = """
You are a senior software engineer planning how to solve a coding task.
Your job is to break the task into clear, logical steps BEFORE writing any code.

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "steps": [
    {
      "step_number": 1,
      "description": "What this step does in plain English",
      "reasoning": "Why this step is necessary"
    }
  ]
}

Rules:
- Between 3 and 7 steps. No more, no less.
- Each step should be a single, clear action.
- reasoning must explain WHY, not just repeat the description.
- Do not write any code in the plan.
""".strip()

SQL_PLAN_SYSTEM = """
You are a senior data engineer planning how to solve a SQL task.
Your job is to:
  1. Infer what database tables and columns the query will need.
  2. Break the task into clear steps.
  3. Generate dummy data that is GUARANTEED to satisfy the query conditions.

Return ONLY a JSON object in this exact format — no markdown, no explanation:
{
  "steps": [
    {
      "step_number": 1,
      "description": "What this step does in plain English",
      "reasoning": "Why this step is necessary"
    }
  ],
  "schema": {
    "create_statements": [
      "CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, revenue REAL);"
    ],
    "insert_statements": [
      "INSERT INTO customers VALUES (1, 'Alice', 5000.0);",
      "INSERT INTO customers VALUES (2, 'Bob', 3200.0);"
    ],
    "table_descriptions": [
      "customers — stores customer name and total revenue"
    ]
  }
}

Rules:
- Between 3 and 7 steps.
- Schema must include at least 6 realistic dummy rows per table.
- Use standard SQLite-compatible SQL only.
- Do not include any explanation outside the JSON.

CRITICAL RULE — DATE-BASED QUERIES:
If the query filters by a date range (e.g. "last 30 days", "last 7 days", "this month",
"this year", "yesterday"), you MUST use SQLite relative date expressions in INSERT statements
so the dummy data is always within the filter window.

Use: date('now', '-N days')  — NOT hardcoded dates like '2024-01-15'

Examples for "last 30 days":
  CORRECT: INSERT INTO sales VALUES (1, 'ProductA', 500.0, date('now', '-5 days'));
  CORRECT: INSERT INTO sales VALUES (2, 'ProductB', 300.0, date('now', '-12 days'));
  CORRECT: INSERT INTO sales VALUES (3, 'ProductC', 200.0, date('now', '-20 days'));
  WRONG:   INSERT INTO sales VALUES (1, 'ProductA', 500.0, '2024-01-15');  ← NEVER do this

For "last 30 days" queries: spread dates from date('now', '-1 days') to date('now', '-28 days').
For "last 7 days": use date('now', '-1 days') through date('now', '-6 days').
For "this month": use date('now', 'start of month') plus a few days offset.

Generate enough rows (at least 6) so LIMIT/TOP-N queries return full results.
""".strip()


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

class PlanningAgent:
    """
    Breaks a user task into a typed step-by-step plan.

    For SQL tasks, also infers the schema and generates dummy data
    that will be loaded into SQLite before query execution.

    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=1500,
            request_timeout=settings.groq_request_timeout,
        )

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

        Sends the task to the LLM with a language-specific system prompt,
        parses the JSON response into typed Pydantic models, and returns
        a partial state dict that LangGraph merges back into the full state.

        Args:
            state: The current PipelineState. Reads: task, language.

        Returns:
            Partial state dict with keys:
              - "plan": List[PlanStep]
              - "sql_schema": SQLSchema | None  (SQL tasks only)
              - "status": PipelineStatus.PLANNING
        """
        ma = state.model_assignments or {}
        model = ma.get("planner", settings.groq_model_primary)
        self.llm = self._build_llm(model)
        logger.info("PlanningAgent using model: %s", model)

        logger.info("PlanningAgent running for task: %s", state.task[:60])

        # Select the right system prompt based on language
        system_prompt = (
            SQL_PLAN_SYSTEM
            if state.language == Language.SQL
            else PYTHON_PLAN_SYSTEM
        )

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

        # ── LLM call ──────────────────────────────────────────────── #
        try:
            response = self.llm.invoke(messages)
            raw = response.content.strip()
            logger.debug("PlanningAgent raw response: %s", raw[:300])
        except Exception as e:
            logger.error("PlanningAgent LLM call failed: %s", e)
            raise RuntimeError(f"Planning agent 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)

        # ── Parse JSON ────────────────────────────────────────────── #
        parsed = _parse_json(raw, agent="PlanningAgent")

        # ── Build typed plan steps ────────────────────────────────── #
        plan: list[PlanStep] = []
        for raw_step in parsed.get("steps", []):
            try:
                plan.append(PlanStep(**raw_step))
            except (ValidationError, TypeError) as e:
                logger.warning("Skipping malformed plan step: %s — %s", raw_step, e)

        if not plan:
            raise ValueError("PlanningAgent returned no valid plan steps.")

        logger.info("PlanningAgent produced %d steps", len(plan))

        # ── Build SQL schema (SQL tasks only) ─────────────────────── #
        sql_schema: SQLSchema | None = None
        if state.language == Language.SQL and "schema" in parsed:
            try:
                sql_schema = SQLSchema(**parsed["schema"])
                logger.info(
                    "PlanningAgent inferred schema with %d table(s)",
                    len(sql_schema.create_statements),
                )
            except (ValidationError, TypeError) as e:
                logger.warning("SQL schema parsing failed: %s", e)

        return {
            "plan":        plan,
            "sql_schema":  sql_schema,
            "status":      PipelineStatus.PLANNING,
            "token_usage": updated_token_usage,
        }


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

def _parse_json(raw: str, agent: str) -> dict:
    """
    Safely parse a JSON string returned by the LLM.

    LLMs occasionally wrap JSON in markdown code fences (```json ... ```)
    even when instructed not to. This function strips fences before
    attempting to parse.

    Args:
        raw:   Raw string content from the LLM response.
        agent: Name of the calling agent, used in error messages.

    Returns:
        Parsed dict from the JSON response.

    Raises:
        ValueError: If the string cannot be parsed as valid JSON.
    """
    # Strip common markdown fences the LLM might add despite instructions
    cleaned = raw
    if cleaned.startswith("```"):
        # Remove opening fence (```json or ```)
        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\nRaw output: %s", agent, e, raw[:500])
        raise ValueError(
            f"{agent} returned invalid JSON. "
            f"Parse error: {e}. Raw output: {raw[:200]}"
        ) from e