#!/usr/bin/env python3
"""Planner Agent — Generated by ABI-Core"""

import json

from planner import AbiPlannerAgent
from abi_core.agent import AbiCore
from abi_core.common.utils import abi_logging, clean_llm_json
from abi_core.common.plan_models import PlannerOutput
from abi_core.common.semantic_tools import (
    tool_find_agent,
)

agent = AbiCore()


@agent.task(
    name="analyze_query",
    input_map={
        "query": "$input.query",
        "context": "$input.context",
    },
)
async def analyze_query(query, context):
    """Prepare the planning prompt from query + context."""
    planning_query = f"User request: {query}\nContext: {json.dumps(context, indent=2)}"
    return {"planning_query": planning_query}


@agent.task(
    name="parse_plan",
    depends_on=["analyze_query"],
    input_map={"raw_response": "$input.llm_response"},
)
def parse_plan(raw_response):
    """Clean, parse, and validate the LLM response into a structured plan.

    Uses PlannerOutput (Pydantic) to enforce schema — descriptions are
    automatically cleaned of LLM reasoning artifacts by field validators.
    Falls back to clean_llm_json if Pydantic validation fails.
    """
    if not raw_response:
        return {"status": "error", "message": "Empty LLM response"}

    parsed = clean_llm_json(raw_response)

    # Validate with Pydantic — cleans descriptions automatically
    try:
        validated = PlannerOutput.model_validate(parsed)
        plan_dict = validated.to_dict()
        abi_logging(f"[✅] Plan validated by PlannerOutput schema")
        # Log the full plan for observability
        plan = plan_dict.get("plan", {})
        abi_logging(f"[📋] PLAN: objective='{plan.get('objective', '')}' strategy={plan.get('execution_strategy', '')}")
        for t in plan.get("tasks", []):
            abi_logging(f"[📋]   {t.get('task_id')}: {t.get('description', '')[:120]} deps={t.get('dependencies', [])}")
        return plan_dict
    except Exception as e:
        abi_logging(f"[⚠️] Pydantic validation failed, using raw parsed: {e}")
        # Fallback: still clean descriptions manually
        from abi_core.common.utils import _clean_description
        plan = parsed.get("plan", {})
        for task in plan.get("tasks", []):
            if "description" in task:
                task["description"] = _clean_description(task["description"])
        return parsed


@agent.task(
    name="assign_agents",
    depends_on=["parse_plan"],
    input_map={"plan_data": "$parse_plan"},
)
async def assign_agents(plan_data):
    """Assign agents to each task in the plan.

    For each task:
    1. Search for an existing agent → type: "execute"
    2. No agent found → type: "build_and_execute" (builder resolves tools)
    
    The builder is responsible for searching tools, deciding what to create,
    and building the ephemeral agent. The planner only finds existing agents.
    """
    if plan_data.get("status") != "ready":
        return plan_data  # pass through clarification / error

    plan = plan_data.get("plan", {})
    tasks = plan.get("tasks", [])

    # Infrastructure agents that should never be assigned as task executors
    INFRA_AGENTS = {"builder", "planner", "orchestrator", "guardian", "semantic-layer"}

    abi_logging(f"[🔍] Assigning agents to {len(tasks)} tasks...")

    for task in tasks:
        task_desc = task.get("description", "")
        task_id = task.get("task_id", "unknown")

        # Try to find an existing agent
        found_agent = await tool_find_agent.ainvoke(task_desc)

        if found_agent:
            agent_name = (
                found_agent.name if hasattr(found_agent, "name") else
                found_agent.get("name", "") if isinstance(found_agent, dict) else ""
            ).lower()
            # Skip infrastructure agents — they are not task executors
            if any(infra in agent_name for infra in INFRA_AGENTS):
                abi_logging(
                    f"[⚠️] Task '{task_id}': found '{agent_name}' but it's infrastructure, skipping"
                )
                found_agent = None

        if found_agent:
            agent_data = found_agent.model_dump() if hasattr(found_agent, "model_dump") else found_agent
            task["type"] = "execute"
            task["agents"] = [agent_data]
            abi_logging(f"[✅] Task '{task_id}': agent found → execute")
            continue

        # No agent → builder handles everything (tool search + creation + ephemeral)
        task["type"] = "build_and_execute"
        task["agents"] = []
        task["builder_spec"] = {
            "system_prompt": task_desc,
            "ephemeral": True,
        }
        abi_logging(f"[🏗️] Task '{task_id}': no agent → build_and_execute (builder resolves tools)")

    # Log final plan with assignments
    abi_logging(f"[📋] FINAL PLAN ({len(tasks)} tasks):")
    for t in tasks:
        tid = t.get("task_id", "?")
        ttype = t.get("type", "?")
        desc = t.get("description", "")[:100]
        agents = [a.get("name", "?") if isinstance(a, dict) else str(a) for a in t.get("agents", [])]
        abi_logging(f"[📋]   {tid} [{ttype}] {desc} agents={agents}")

    return plan_data


agent.run(AbiPlannerAgent())
