"""Agent planning and execution models for the GuardianHub system.

This module contains Pydantic models for agent planning, mission execution,
and inter-agent communication within the GuardianHub ecosystem.
"""
import time
from typing import List, Dict, Any, Optional, Literal
from uuid import uuid4
from pydantic import BaseModel, Field, conlist
from ..registry.registry import register_model

# =============================================================================
# Core Planning Models
# =============================================================================

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."
        )
    )


@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,
        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.
    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"
    )


# =============================================================================
# Agent Mission and Communication Models
# =============================================================================

class AgentSubMission(BaseModel):
    """The 'Marching Orders' sent from Sutram to a Specialist Microservice."""
    mission_id: str = Field(
        default_factory=lambda: str(uuid4()),
        description="Unique identifier for this mission"
    )
    
    session_id: str = Field(
        ...,
        description="Original conversation/session ID for trace continuity."
    )

    # Strategic Context
    sub_objective: str = Field(
        ...,
        description="The specific goal for the specialist (e.g., 'Audit the Penrose DC')."
    )
    
    assigned_mission_intent: str = Field(
        ...,
        description="The mission intent string from the registry."
    )

    # Operational Guardrails
    constraints: List[str] = Field(
        default_factory=list,
        description="Rules inherited from the Strategic Architect."
    )
    
    auth_context: List[str] = Field(
        default_factory=list,
        description="Vault keys Sutram has authorized for this call."
    )

    # Situational Awareness
    relevant_beliefs: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Relevant facts/lessons from the Global Blackboard (Neo4j/ACE)."
    )

    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Additional metadata for the mission."
    )

    # Traceability
    parent_span_id: Optional[str] = Field(
        None,
        description="Parent span ID for distributed tracing"
    )
    
    callback_url: str = Field(
        "http://sutram-orchestrator:7002/v1/agent/callback",
        description="URL to send the mission completion callback to"
    )


class ExplorationPlan(BaseModel):
    """Returned when a Specialist identifies a Fog of War."""
    status: Literal["EXPLORING"] = "EXPLORING"
    
    reason: str = Field(
        ...,
        description="The information gap or missing tool identified."
    )
    
    target_agents: List[str] = Field(
        ...,
        description=(
            "The peer agents to be consulted "
            "(e.g., ['research-agent', 'k8s-agent'])."
        )
    )
    
    inquiry: str = Field(
        ...,
        description="The specific question or data needed from peers."
    )
    
    hindsight_required: bool = Field(
        True,
        description="Should Sutram provide historical logs for this gap?"
    )


class A2AExchangeSchema(BaseModel):
    """The 'Handshake' used for Peer-to-Peer capability negotiation."""
    negotiation_id: str = Field(
        default_factory=lambda: str(uuid4()),
        description="Unique identifier for this negotiation"
    )
    
    capability_requested: str = Field(
        ...,
        description="The capability being requested (e.g., 'get_prometheus_metrics')"
    )
    
    required_format: Literal["json", "graph_nodes", "markdown_report"] = Field(
        ...,
        description="The expected format of the response"
    )
    
    priority: Literal["low", "medium", "high", "critical"] = Field(
        ...,
        description="Priority level of the request"
    )

    # The 'Belief' Exchange
    shared_context: Dict[str, Any] = Field(
        default_factory=dict,
        description="A slice of the sender's local beliefs to help the receiver."
    )

    # Resource Constraints
    token_budget: int = Field(
        2000,
        description="Maximum token budget for the response"
    )
    
    hop_limit: int = Field(
        2,
        description="Maximum number of hops allowed (prevents infinite recursive delegation)"
    )


class A2AMessage(BaseModel):
    """The actual carrier for peer-to-peer agent communication.
    
    This message format is used for all inter-agent communication, including delegation,
    capability negotiation, and data exchange between agents in the GuardianHub ecosystem.
    """
    # Message Identification
    message_id: str = Field(
        default_factory=lambda: str(uuid4()),
        description="Unique identifier for this message"
    )
    
    # Agent Information
    sender: str = Field(
        ...,
        description="Agent ID of the message sender"
    )
    
    receiver: str = Field(
        ...,
        description="Agent ID of the intended recipient"
    )

    # Traceability
    trace_parent: str = Field(
        ...,
        description=(
            "OpenTelemetry trace parent ID for distributed tracing "
            "across agent boundaries"
        )
    )
    
    session_id: str = Field(
        ...,
        description="The original conversation/session ID for request correlation"
    )

    # Message Content
    message_type: Literal[
        "DELEGATE_EXPLORATION",
        "NEGOTIATE_CAPABILITY",
        "DATA_PROVISION"
    ] = Field(
        ...,
        description=(
            "Type of the message. Determines how the payload should be interpreted. "
            "Possible values: DELEGATE_EXPLORATION, NEGOTIATE_CAPABILITY, DATA_PROVISION"
        )
    )
    
    payload: Dict[str, Any] = Field(
        ...,
        description=(
            "The actual content of the message. Structure depends on the message_type. "
            "Should be a JSON-serializable dictionary."
        )
    )

    # System Metadata
    hop_count: int = Field(
        default=0,
        description=(
            "Number of hops this message has taken. "
            "Incremented by each forwarding agent to prevent infinite loops."
        )
    )
    
    timestamp: float = Field(
        default_factory=time.time,
        description="Unix timestamp (in seconds) when the message was created"
    )