import re
import time
from datetime import datetime
from enum import Enum
from typing import Dict, List, Optional, Any, Literal, Union
from uuid import uuid4, UUID

from .registry.registry import register_model
from guardianhub.config.settings import settings

from pydantic import BaseModel, Field, field_validator, ConfigDict , model_validator

# =============================================================================
# ZONE 1: ADMINISTRATIVE ENUMS & BASE
# =============================================================================

class AgentStatus(str, Enum):
    # Use the values your DB is actually storing (lowercase 'draft')
    DRAFT = "draft"
    ACTIVE = "active"
    INACTIVE = "inactive"
    TRAINING = "training"


class AgentBase(BaseModel):
    """The fundamental identity with the Trinity of Names."""
    name: str = Field(..., description="Technical identifier (e.g., cmdb-agent-service)")

    title: Optional[str] = Field(..., description="Display name (e.g., Infrastructure Guardian)")
    semantic_identity: Optional[str] = Field(..., description="The key for the Agent's core DNA")

    description: str = Field(default="")
    system_prompt: str = Field(..., description="The core persona instructions.")

# =============================================================================
# ZONE 2: REGISTRY & DB MODELS (Internal & Admin API)
# =============================================================================

class AgentMission(BaseModel):
    """The 'Subscription' data for linking an Agent to a Semantic Tool."""
    semantic_tool: str = Field(..., description="The name of the semantic tool.")
    mission: str = Field(..., description="The specific intent for this tool in this agent's context.")
    constraints: Optional[str] = Field(None, description="Operational guardrails.")


@register_model
class AgentCreate(AgentBase):
    """The strict schema for Registry DB operations and LLM Generation."""
    domain: str = Field(..., description="The operational domain.", json_schema_extra={"example": "infrastructure"})
    status: AgentStatus = Field(AgentStatus.DRAFT)
    tags: List[str] = Field(default_factory=list)
    # 🟢 THE SOVEREIGN MOVE: Tools are now Missions
    missions: List[AgentMission] = Field(
        default_factory=list,
        description="List of authorized missions for this agent."
    )
    reflection_config: Dict[str, Any] = Field(default_factory=dict)
    metadata: Dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode='after')
    def settle_identity_defaults(self) -> 'AgentCreate':
        """Ensures the Trinity of Names is settled at birth."""
        # 1. Settle Title: "cmdb-agent" -> "Cmdb Agent"
        if not self.title:
            self.title = self.name.replace("-", " ").replace("_", " ").title()

        # 2. Settle Semantic Identity: "cmdb-agent" -> "aura_cmdb_agent"
        if not self.semantic_identity:
            clean_name = re.sub(r'[^a-zA-Z0-9]+', '_', self.name).lower()
            self.semantic_identity = f"aura_{clean_name}"

        return self

class Agent(AgentCreate):
    """The complete Agent record including DB-generated fields."""
    id: str = Field(..., description="UUID from the database.")
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

    model_config = ConfigDict(from_attributes=True)

@register_model
class AgentResponse(BaseModel):
    """Response model for agent operations."""
    id: Union[str, UUID] # 🟢 FIX: Accept both native UUID and String
    name: str  # Technical
    title: str  # Human
    semantic_identity: str  # Aura

    status: AgentStatus = Field(..., description="Current status of the agent")
    domain: str = Field(..., description="Domain of the agent")
    description: str = Field(default="", description="Description of the agent")

    system_prompt: str = Field(..., description="System prompt for the agent")
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    message: Optional[str] = Field(None, description="Additional details about the operation")
    metadata: Dict[str, Any] = Field(default_factory=dict)
    tags: List[str] = []

    model_config = ConfigDict(
        from_attributes=True,
        # Use str_strip_whitespace for better string handling
        str_strip_whitespace=True
    )

    @field_validator('id', mode='before')
    @classmethod
    def validate_id(cls, v):
        # Handle UUID to string conversion
        if isinstance(v, UUID):
            return str(v)
        return v

    @field_validator('metadata', mode='before')
    @classmethod
    def validate_metadata(cls, v):
        # If it's the SQLAlchemy MetaData object or None, return empty dict
        if v is None or not isinstance(v, dict):
            return {}
        return v

# =============================================================================
# ZONE 3: LLM STRUCTURED OUTPUT MODELS (Registered for LLM awareness)
# =============================================================================

@register_model
class SearchPhrasesResponse(BaseModel):
    search_phrases: List[str] = Field(..., description="1 to 3 relevant search phrases.")

@register_model
class LLMReflectionConfig(BaseModel):
    enabled: Optional[bool] = Field(None)
    optimization_types: List[str] = Field(default_factory=list)
    synthesis_directives: Optional[Dict[str, Any]] = Field(None)

@register_model
class LLMKnowledgeSuggestion(BaseModel):
    type: str = Field(..., description="e.g., 'document', 'query'")
    value: str = Field(..., description="ID or URL")

@register_model
class AgentLLMConfigSuggestion(BaseModel):
    """Structured output for LLM-generated agent proposals."""
    system_prompt: str = Field(...)
    reflection_config: LLMReflectionConfig = Field(...)
    initial_warmup_query: str = Field(...)
    knowledge_suggestions: List[LLMKnowledgeSuggestion] = Field(default_factory=list)

# =============================================================================
# ZONE 4: STRATEGIC DISPATCH & A2A (Wire Formats - Not Registered)
# =============================================================================
@register_model
class AgentSubMission(BaseModel):
    """Tactical 'Marching Orders' dispatched from Sutram to a Specialist."""
    mission_id: str = Field(default_factory=lambda: str(uuid4()))
    session_id: str = Field(..., description="The unique session ID for the OODA loop.")

    # 🚀 Strategic Context: Handed to Config Defaults
    default_dna: str = "Standard Sovereign Protocol"
    template_id: str = Field(
        description="The OODA template ID being executed."
    )

    # Core Objective
    sub_objective: str = Field(..., description="The specialist's specific goal.")
    assigned_mission_intent: str = Field(..., description="The mission intent from the registry.")

    # 🚀 Shared Situational Awareness
    context_snapshot: Dict[str, Any] = Field(
        default_factory=dict,
        description="Facts discovered by previous agents in this session."
    )

    # Constraints & Auth
    constraints: List[str] = Field(default_factory=list)
    auth_context: List[str] = Field(default_factory=list)

    # 🎯 RESILIENCE FIX: Ensure these are simple dicts to avoid
    # Pydantic recursive validation errors in Temporal
    relevant_beliefs: List[Dict[str, Any]] = Field(default_factory=list)

    # Metadata & Callback
    metadata: Dict[str, Any] = Field(default_factory=dict)

    # 🎯 THE GLOBAL CALLBACK: Handed to Config Defaults
    callback_url: str = Field(
        default=settings.endpoints.SUTRAM_CALLBACK_URL,
        description="Standardized callback endpoint."
    )

    model_config = ConfigDict(arbitrary_types_allowed=True)
# =============================================================================
# ZONE 3: LLM STRUCTURED OUTPUT MODELS (Registered for LLM awareness)
# =============================================================================

@register_model
class A2AExchangeSchema(BaseModel):
    """
    The 'Handshake' used for Peer-to-Peer capability negotiation.
    The LLM generates this when it needs to consult another agent.
    """
    negotiation_id: str = Field(
        default_factory=lambda: str(uuid4()),
        description="Unique identifier for this peer-to-peer negotiation."
    )

    capability_requested: str = Field(
        ...,
        description="The specific capability being requested (e.g., 'get_k8s_logs')."
    )

    required_format: Literal["json", "graph_nodes", "markdown_report"] = Field(
        ...,
        description="The expected format of the response from the peer."
    )

    priority: Literal["low", "medium", "high", "critical"] = Field(
        "medium",
        description="Priority level of the request."
    )

    # The 'Belief' Exchange
    shared_context: Dict[str, Any] = Field(
        default_factory=dict,
        description="A slice of local context/facts to help the receiver process the request."
    )

    token_budget: int = Field(2000, description="Max tokens for the response.")
    hop_limit: int = Field(2, description="Prevents infinite recursive delegation.")


@register_model
class ExplorationPlan(BaseModel):
    """
    The LLM generates this when it identifies a 'Fog of War'
    and needs to propose an inquiry path to the Orchestrator.
    """
    status: Literal["EXPLORING"] = "EXPLORING"
    reason: str = Field(..., description="The gap in knowledge or toolset.")
    target_agents: List[str] = Field(..., description="Peer agents to consult.")
    inquiry: str = Field(..., description="The specific question for the peers.")


class A2AMessage(BaseModel):
    """Carrier for Peer-to-Peer communication."""
    message_id: str = Field(default_factory=lambda: str(uuid4()))
    sender: str
    receiver: str
    trace_parent: str
    session_id: str
    message_type: Literal["DELEGATE_EXPLORATION", "NEGOTIATE_CAPABILITY", "DATA_PROVISION"]
    payload: Dict[str, Any]
    hop_count: int = 0
    timestamp: float = Field(default_factory=time.time)


@register_model
class TacticalAuditReport(BaseModel):
    """The formal response from the Safety Officer."""
    success: bool = Field(True)
    decision: Literal["PROCEED", "HALT", "ADVISE"] = Field(...)
    risk_score: float = Field(..., ge=0.0, le=1.0)
    # 🎯 ALIGNMENT: Match the workflow's expectation
    justification: str = Field(..., description="The reasoning behind the safety decision")

@register_model
class RiskAnalysisReport(BaseModel):
    """Structured output for the Tactical Safety Audit."""
    decision: Literal["PROCEED", "HALT", "ADVISE"] = Field(
        ...,
        description="The final safety verdict"
    )
    risk_score: float = Field(
        ...,
        ge=0.0, le=1.0,
        description="Quantitative risk (0=Safe, 1=Catastrophic)"
    )
    reason: str = Field(
        ...,
        description="Detailed explanation for the decision"
    )
    mitigation_suggested: Optional[str] = Field(
        None,
        description="Steps to reduce risk if decision is ADVISE or HALT"
    )
    impact_level: Literal["low", "medium", "high", "critical"] = Field("medium")


@register_model
class MissionDebriefReport(BaseModel):
    """The final analytical report synthesized at the end of a mission."""
    summary: str = Field(..., description="High-level overview of mission outcomes")
    key_takeaways: List[str] = Field(
        default_factory=list,
        description="Bullet points of lessons learned or technical findings"
    )
    status_attained: Literal["SUCCESS", "PARTIAL_SUCCESS", "FAILURE"] = Field(...)
    technical_debt_identified: Optional[str] = Field(
        None,
        description="Any CMDB gaps or misconfigurations found during the run"
    )
    next_steps_suggested: List[str] = Field(default_factory=list)



from pydantic import BaseModel, Field
from typing import List, Dict, Any

class DiagnosticQuestion(BaseModel):
    """A strategic inquiry designed to isolate failure modes."""
    question: str = Field(..., description="The specific question to ask the system or logs.")
    rationale: str = Field(..., description="Why this question is critical for this mission.")
    expected_evidence: str = Field(..., description="What 'good' looks like in the response.")

class FailureMode(BaseModel):
    """A known structural gap or common misconfiguration."""
    mode: str = Field(..., description="e.g., 'Orphaned Application Service'")
    symptoms: List[str] = Field(..., description="Observable traits of this failure.")
    remediation_hint: str = Field(..., description="Initial guidance for the Specialist.")

@register_model
class MissionBriefingPacket(BaseModel):
    """
    The 'Cognitive Aura' for a specific mission.
    Generated by the Sync Agent, consumed by the Specialist Agent.
    """
    mission_intent: str = Field(..., description="The unique key for the intent.")
    diagnostic_pipeline: List[DiagnosticQuestion] = Field(
        ...,
        description="A sequence of questions to resolve the mission."
    )
    known_failure_modes: List[FailureMode] = Field(default_factory=list)
    success_criteria: List[str] = Field(
        ...,
        description="Measurable signals that the mission is complete."
    )
    suggested_search_queries: List[str] = Field(
        default_factory=list,
        description="Queries to run if local knowledge is insufficient."
    )

class IncarnationRecord(BaseModel):
    """A single life of the agent (from Vector/Leela)."""
    id: str
    content: str
    similarity: float
    produced_at: Optional[datetime] = None

class WisdomInsight(BaseModel):
    """A learned lesson (from Graph/Dharma)."""
    lesson_topic: str
    advice: str
    past_outcome: str

class AgentAncestry(BaseModel):
    """The collection of past lives."""
    incarnations: List[IncarnationRecord] = []
    aha_moments: List[WisdomInsight] = []


class AgentDossier(BaseModel):
    """The Complete Strategic Resume + Institutional Memory."""
    identity: AgentResponse = Field(..., description="The Trinity: Name, Title, Aura Identity")
    system_prompt: str = Field(..., description="The core persona and instructions")
    authorized_missions: List[AgentSubMission] = Field(default_factory=list)
    ancestry: AgentAncestry = Field(..., description="The 10 incarnations and learned wisdom")

    # 🟢 THE CRITICAL PLUMBING
    config: Dict[str, Any] = Field(
        default_factory=dict,
        description="Faculties, expertise, and domain settings"
    )
    memory_config: Dict[str, Any] = Field(
        default_factory=dict,
        description="Persistence strategy and memory tiering"
    )

    model_config = ConfigDict(from_attributes=True)