from typing import List, Dict, Any, Optional, Literal
from pydantic import BaseModel, Field, conlist
from ..registry.registry import register_model


# --- STEP 1: Define the individual tool-call step model ---
class PlanStep(BaseModel):
    """A single step in the agent's overall macro-plan."""

    # The exact name of the tool to be called (must match a registered tool)
    tool_name: str = Field(..., description="The exact name of the tool to execute.")

    # A unique name for this step, used as the key to store the result in macro_context.
    step_name: str = Field(...,
                           description="A unique identifier for this plan step. Its result will be stored under this name in the macro_context.")

    # The arguments for the tool call, passed as a dictionary.
    tool_args: Dict[str, Any] = Field(..., description="The dictionary of arguments required for the tool call.")

    # The list of results (by step_name) this step requires before it can be executed.
    dependencies: List[str] = Field(
        default_factory=list,
        description="List of step_name strings whose results must be available in macro_context before this step can run."
    )

    # Internal status, handled by the graph executor.
    status: Literal["pending", "running", "complete", "error"] = Field(
        default="pending",
        description="The current status of the step (pending, running, complete, error). Must be 'pending' initially."
    )


# --- STEP 2: Define the Macro Plan model ---
@register_model
class MacroPlan(BaseModel):
    """The complete structured plan generated by the LLM."""

    # Must be an array of PlanStep objects.
    plan: conlist(PlanStep, min_length=0) = Field(
        default_factory=list,  # Use default_factory for mutable types
        description="The sequence of PlanStep objects required to fulfill the user request. Can be empty if no tools are needed."
    )

    # A final reflection or thought process before synthesis/execution.
    # 🟢 FIX: Make 'reflection' Optional and provide an empty string default.
    # This allows the model to validate even if the key is missing from the LLM's JSON.
    reflection: Optional[str] = Field(
        default="",
        description="A brief thought process explaining why this plan was chosen and how it addresses the user's input."
    )
    metadata: Dict[str, Any] = Field(
        ...,
        description=(
            "Additional metadata about the plan in a Pydantic-compatible JSON Schema format. "
            "This can include configuration parameters, execution context, or any other "
            "plan-specific metadata needed for execution or analysis."
        )
    )



class MacroPlanResponse(BaseModel):
    """Response model for the LLM's planning output."""
    plan: List[PlanStep] = Field(default_factory=list, description="List of plan steps")
    reflection: Optional[str] = Field(default="", description="Reasoning behind the plan")
