"""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, Union
from uuid import uuid4
from ..registry.registry import register_model
from ...config.settings import settings


"""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, ConfigDict,field_validator
from typing import List, Dict, Any, Optional
from datetime import datetime


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

# =============================================================================
# Unified Workflow Spine
# =============================================================================
class SchemaKeyValuePair(BaseModel):
    """Gemini-safe way to handle dynamic arguments."""
    model_config = ConfigDict(extra='forbid')
    key: str
    value: str


class PlanStep(BaseModel):
    """The Architect's Intent: Defined by the LLM/Planner."""
    # 🎯 GEMINI FIX: Forbid dynamic/extra properties
    model_config = ConfigDict(extra='forbid')

    step_id: str = Field(description="Unique ID for the step.")
    step_name: str = Field(..., description="Human-readable name of the step.")
    tool_name: str = Field(..., description="Target tool identifier.")
    # 🎯 GEMINI TACTIC: Using Dict[str, str] is safer than Dict[str, Any]
    # if you want to avoid 'additionalProperties' issues in nested maps.
    tool_args: Dict[str, str] = Field(default_factory=dict)
    dependencies: List[str] = Field(default_factory=list)
    status: str = "pending"

    @field_validator('tool_args', mode='before')
    @classmethod
    def normalize_tool_args(cls, v: Any) -> Dict[str, str]:
        """
        🚀 THE BRIDGE: Converts Sutram's [{'key': 'k', 'value': 'v'}]
        into a standard dictionary {'k': 'v'}.
        """
        if isinstance(v, list):
            # If it's a list of dicts with 'key' and 'value'
            return {item['key']: str(item['value']) for item in v if 'key' in item and 'value' in item}
        if isinstance(v, dict):
            # Already normalized
            return v
        return {}

class PlanStepIntent(BaseModel):
    """Clean 'Intent' model for the LLM Architect."""
    model_config = ConfigDict(extra='forbid')

    step_name: str = Field(..., description="Name of the operation.")
    tool_name: str = Field(..., description="Target tool ID.")
    # 🎯 FIX: Replace Dict[str, Any] with a List of pairs
    tool_args: List[SchemaKeyValuePair] = Field(default_factory=list)
    dependencies: List[str] = Field(
        default_factory=list,
        description="Names of steps this depends on."
    )


    def to_plan_step(self) -> Dict[str, Any]:
        """Converts intent list to a standard PlanStep dictionary."""
        return {
            "step_id": f"step_{uuid4().hex[:4]}",
            "step_name": self.step_name,
            "tool_name": self.tool_name,
            "tool_args": {kv.key: kv.value for kv in self.tool_args},
            "dependencies": self.dependencies,
            "status": "pending"
        }

@register_model
class ActionStep(PlanStep):
    # 🎯 ALIAS maps the incoming 'tool_name' from PlanStep to 'action_name'
    model_config = ConfigDict(
        extra='ignore',  # 🚀 THE FIX: Silently drop extra fields like template_id
        populate_by_name=True,
        arbitrary_types_allowed=True
    )

    action_name: str = Field(..., alias="tool_name")
    action_input: Dict[str, Any] = Field(..., alias="tool_args")  # 🎯 BRIDGE THE GAP

    session_id: str
    trace_id: Optional[str] = None
    agent_name: str = "specialist"
    is_dry_run: bool = False


# =============================================================================
# Mission Wrappers
# =============================================================================



@register_model
class MacroPlan(BaseModel):
    # 🎯 GEMINI FIX: Forbid additional properties at the root level

    model_config = ConfigDict(extra='forbid')


    confidence_score: float = Field(default=0.0, description="Confidence score")
    # Fixed metadata fields instead of dynamic dict
    reflection: str = Field(default="", description="The Architect's reasoning for this plan")

    metadata: List[SchemaKeyValuePair] = Field(default_factory=list)
    steps: Union[List[PlanStepIntent], List[PlanStep]] = Field(default_factory=list)

    @field_validator('steps', mode='after')
    @classmethod
    def process_steps(cls, v: List[PlanStepIntent]) -> List[PlanStepIntent]:
        """
        No more attribute errors.
        We simply validate that names are unique here.
        The conversion to PlanStep (with IDs) happens in the Orchestrator.
        """
        step_names = [s.step_name for s in v]
        if len(step_names) != len(set(step_names)):
            raise ValueError("Duplicate step names found in MacroPlan")
        return v


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


class MissionRequest(BaseModel):
    # 🎯 Data validation logic
    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)

    session_id: str
    trace_id: Optional[str] = None
    agent_name: str = "specialist"
    is_dry_run: bool = False
    template_id: str = "NO_TEMPLATE_ID"
    sub_objective: str = "NO_SUB_OBJECTIVE"
    assigned_mission_intent: str= "NO_MISSION_ASSIGNED"
    active_template_id: str = "NO_TEMPLATE_ID"
    mission_dna: str = "Audit Only"
    callback_url: str = settings.endpoints.SUTRAM_CALLBACK_URL
    history_limit: int = 5
    min_success_score: float = 0.7
    risk_threshold: float = 0.7

    # 🚀 The Core Fix: MacroPlan is now a proper model, not just a list
    macro_plan: MacroPlan

    metadata: Dict[str, Any] = Field(default_factory=dict)
    debug_mode: bool = True

    @field_validator('macro_plan')
    @classmethod
    def must_have_steps(cls, v: MacroPlan) -> MacroPlan:
        if not v.steps or len(v.steps) == 0:
            raise ValueError("Mission rejected: Sutram provided an empty MacroPlan.")
        return v

        # 🚀 ADD THIS PROPERTY HERE
    @property
    def technical_context(self) -> Dict[str, Any]:
        """Provides a unified bundle for activity arguments."""
        return {
            "session_id": self.session_id,
            "trace_id": self.trace_id,
            "template_id": self.active_template_id,
            "agent_name": self.agent_name,
            "is_dry_run": self.is_dry_run
        }

