from pydantic import BaseModel, Field, ConfigDict, model_validator
from typing import List, Literal, Optional, Union
from datetime import datetime

# --- User Schemas ---

class OIDCLogin(BaseModel):
    """Schema for OIDC user data received from a login callback."""
    oidc_id: str = Field(..., description="The unique ID from the OIDC provider.")
    email: str = Field(..., description="The user's email address.")
    username: str = Field(..., description="The user's username or display name.")

class LocalLoginRequest(BaseModel):
    email: str
    password: str

class PasswordUpdateRequest(BaseModel):
    current_password: str
    new_password: str

class UserStatus(BaseModel):
    """Schema for the response when checking a user's status."""
    id: str = Field(..., description="The internal user ID.")
    email: Optional[str] = Field(None, description="The user's email address.")
    is_logged_in: bool = Field(True, description="Indicates if the user is currently authenticated.")
    is_anonymous: bool = Field(False, description="Indicates if the user is an anonymous user.")
    oidc_configured: bool = Field(False, description="Whether OIDC SSO is enabled on the server.")
    allow_password_login: bool = Field(True, description="Whether local password login is supported by the server.")

class UserProfile(BaseModel):
    id: str
    email: str
    username: Optional[str] = None
    full_name: Optional[str] = None
    role: str = "user"
    group_id: Optional[str] = None
    group_name: Optional[str] = None
    avatar_url: Optional[str] = None
    created_at: datetime
    last_login_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

class UserProfileUpdate(BaseModel):
    username: Optional[str] = None
    full_name: Optional[str] = None
    avatar_url: Optional[str] = None

class UserRoleUpdate(BaseModel):
    role: str

class UserGroupUpdate(BaseModel):
    group_id: str

# --- Group Schemas ---
class GroupBase(BaseModel):
    name: str
    description: Optional[str] = None
    # Policy: {"llm": ["openai"], "tts": ["gcloud"], "stt": ["google"]}
    policy: Optional[dict] = Field(default_factory=dict)

    @model_validator(mode='after')
    def ensure_policy_dict(self) -> 'GroupBase':
        if self.policy is None:
            self.policy = {}
        return self

class GroupCreate(GroupBase):
    pass

class GroupUpdate(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    policy: Optional[dict] = None

class GroupInfo(GroupBase):
    id: str
    created_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

# --- General Schemas ---
class UserPreferences(BaseModel):
    """Schema for user-specific LLM, TTS, STT preferences."""
    llm: dict = Field(default_factory=dict)
    tts: dict = Field(default_factory=dict)
    stt: dict = Field(default_factory=dict)
    statuses: Optional[dict] = Field(default_factory=dict)
    
# --- Config & General Schemas ---
class ConfigResponse(BaseModel):
    """Schema for returning user preferences alongside effective settings."""
    preferences: UserPreferences
    effective: dict = Field(default_factory=dict)

class SystemStatus(BaseModel):
    """Schema for overall system status, including TLS and OIDC state."""
    status: str
    oidc_enabled: bool
    tls_enabled: bool
    external_endpoint: Optional[str] = None
    version: str

class OIDCConfigUpdate(BaseModel):
    enabled: Optional[bool] = None
    client_id: Optional[str] = None
    client_secret: Optional[str] = None
    server_url: Optional[str] = None
    redirect_uri: Optional[str] = None
    allow_oidc_login: Optional[bool] = None

class SwarmConfigUpdate(BaseModel):
    external_endpoint: Optional[str] = None
    grpc_endpoint: Optional[str] = None

class AppConfigUpdate(BaseModel):
    allow_password_login: Optional[bool] = None

# --- Skill Schemas ---
class SkillBase(BaseModel):
    name: str
    description: Optional[str] = None
    skill_type: str = "local" # local, remote_grpc, mcp
    is_enabled: bool = True
    features: List[str] = Field(default_factory=lambda: ["chat"])
    is_system: bool = False
    extra_metadata: dict = Field(default_factory=dict)

class SkillCreate(SkillBase):
    pass

class SkillUpdate(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    skill_type: Optional[str] = None
    is_enabled: Optional[bool] = None
    features: Optional[List[str]] = None
    is_system: Optional[bool] = None
    extra_metadata: Optional[dict] = None

class SkillResponse(SkillBase):
    id: Union[int, str]
    owner_id: str
    created_at: datetime
    model_config = ConfigDict(from_attributes=True)

class SkillFileBase(BaseModel):
    file_path: str
    content: Optional[str] = None

class SkillFileCreate(SkillFileBase):
    pass

class SkillFileUpdate(BaseModel):
    content: str

class SkillFileResponse(SkillFileBase):
    id: Union[int, str]
    skill_id: Union[int, str]
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

# --- Chat Schemas ---
class ChatRequest(BaseModel):
    """Defines the shape of a request to the /chat endpoint."""
    prompt: str = Field(..., min_length=1)
    # The 'model' can now be specified in the request body to switch models mid-conversation.
    provider_name: str = Field("gemini")
    # Add a new optional boolean field to control the retriever
    load_faiss_retriever: Optional[bool] = Field(False, description="Whether to use the FAISS DB retriever for the chat.")


class ChatResponse(BaseModel):
    """Defines the shape of a successful response from the /chat endpoint."""
    answer: str
    provider_used: str
    message_id: Optional[int] = None

# --- Document Schemas ---
class DocumentCreate(BaseModel):
    title: str
    text: str
    source_url: Optional[str] = None
    author: Optional[str] = None
    user_id: str = "default_user"

class DocumentResponse(BaseModel):
    message: str

class SessionBase(BaseModel):
    user_id: str
    title: Optional[str] = None
    provider_name: Optional[str] = None
    feature_name: Optional[str] = "default"
    status: str
    created_at: datetime
    model_config = ConfigDict(from_attributes=True)

class DocumentInfo(BaseModel):
    id: int
    title: str
    source_url: Optional[str] = None
    status: str
    created_at: datetime
    model_config = ConfigDict(from_attributes=True)

class DocumentListResponse(BaseModel):
    documents: List[DocumentInfo]

class DocumentDeleteResponse(BaseModel):
    message: str
    document_id: int

# --- Session Schemas ---
class SessionCreate(BaseModel):
    """Defines the shape for starting a new conversation session."""
    user_id: Optional[str] = None
    provider_name: str = "deepseek"
    model_name: Optional[str] = None
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None
    feature_name: Optional[str] = "default"

class SessionUpdate(BaseModel):
    title: Optional[str] = None
    provider_name: Optional[str] = None
    model_name: Optional[str] = None
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None
    restrict_skills: Optional[bool] = None
    allowed_skill_ids: Optional[List[int]] = None
    allowed_skill_names: Optional[List[str]] = None
    system_prompt_override: Optional[str] = None
    is_locked: Optional[bool] = None
    auto_clear_history: Optional[bool] = None

class Session(BaseModel):
    """Defines the shape of a session object returned by the API."""
    id: int
    user_id: str
    title: Optional[str] = None
    provider_name: Optional[str] = None
    model_name: Optional[str] = None
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None
    feature_name: str
    created_at: datetime
    # M3: Node attachment
    sync_workspace_id: Optional[str] = None
    attached_node_ids: Optional[List[str]] = Field(default_factory=list)
    node_sync_status: Optional[dict] = Field(default_factory=dict)
    sync_config: Optional[dict] = Field(default_factory=dict)
    
    restrict_skills: bool = False
    allowed_skill_names: Optional[List[str]] = Field(default_factory=list)
    system_prompt_override: Optional[str] = None
    is_locked: bool = False
    auto_clear_history: bool = False
    
    model_config = ConfigDict(from_attributes=True)

# --- M3: Session Node Attachment Schemas ---

class NodeWorkspaceConfig(BaseModel):
    """How a node should seed its workspace for a session."""
    source: Literal["empty", "server", "node_local"] = "empty"
    path: Optional[str] = None  # root path on node when source='node_local'
    source_node_id: Optional[str] = None # which node to pull from if source='node_local'
    read_only_node_ids: List[str] = Field(default_factory=list) # M3: Nodes that should not push changes back

class NodeAttachRequest(BaseModel):
    """Attach one or more nodes to a session."""
    node_ids: List[str]
    config: Optional[NodeWorkspaceConfig] = None

class NodeSyncStatusEntry(BaseModel):
    """Per-node sync status within a session."""
    node_id: str
    status: str           # 'pending' | 'syncing' | 'synced' | 'error'
    last_sync: Optional[str] = None
    error: Optional[str] = None

class SessionNodeStatusResponse(BaseModel):
    """Response showing all attached nodes and their sync state."""
    session_id: int
    sync_workspace_id: Optional[str] = None
    nodes: List[NodeSyncStatusEntry] = []
    sync_config: Optional[dict] = Field(default_factory=dict)

# --- M4: Node Config YAML ---

class NodeConfigYamlResponse(BaseModel):
    """The generated config YAML content an admin downloads to set up a node."""
    node_id: str
    config_yaml: str   # Full YAML string ready to save as agent_config.yaml

class Message(BaseModel):
    """Defines the shape of a single message within a session's history."""
    id: int
    # The sender can only be the user, assistant or system.
    sender: Literal["user", "assistant", "system"]
    # The text content of the message.
    content: str
    # The progressive reasoning or 'thinking' step for models that support it
    reasoning_content: Optional[str] = None
    # The timestamp for when the message was created.
    created_at: datetime

    # Metadata such as quality audits, tool usage stats, etc.
    message_metadata: Optional[dict] = None

    # URL to the saved audio file
    audio_url: Optional[str] = None
    # Whether audio exists for this message
    has_audio: bool = False

    # Enables creating this schema from a SQLAlchemy database object.
    model_config = ConfigDict(from_attributes=True)

class MessageHistoryResponse(BaseModel):
    """Defines the response for retrieving a session's chat history."""
    session_id: int
    messages: List[Message]

class SessionTokenUsageResponse(BaseModel):
    """Defines the response for retrieving a session's token usage."""
    token_count: int
    token_limit: int
    percentage: float
    error: Optional[str] = None  # Set when model limit could not be resolved

class SpeechRequest(BaseModel):
    text: str

# --- STT Schemas ---
class STTResponse(BaseModel):
    """Defines the shape of a successful response from the /stt endpoint."""
    transcript: str

class VerifyProviderRequest(BaseModel):
    provider_name: str
    provider_type: Optional[str] = None
    api_key: Optional[str] = None
    model: Optional[str] = None
    voice: Optional[str] = None

class VerifyProviderResponse(BaseModel):
    success: bool
    message: str

class ModelInfoResponse(BaseModel):
    model_name: str
    max_tokens: Optional[int] = None
    max_input_tokens: Optional[int] = None


# ---------------------------------------------------------------------------
# Agent Node Schemas
# ---------------------------------------------------------------------------

# --- Skill Toggles (admin-configured) ---

class SandboxConfig(BaseModel):
    """Granular command execution policy for the shell skill."""
    mode: Literal["STRICT", "PERMISSIVE"] = "STRICT"
    allowed_commands: List[str] = Field(default_factory=lambda: ["ls", "cat", "echo", "pwd", "uname", "curl", "python3", "git"])
    denied_commands: List[str] = Field(default_factory=list)
    sensitive_commands: List[str] = Field(default_factory=list)

class SkillConfig(BaseModel):
    """Per-skill enable/disable with optional config."""
    enabled: bool = True
    cwd_jail: Optional[str] = None          # shell only: restrict working directory
    max_file_size_mb: Optional[int] = None  # sync only: file size cap
    sandbox: Optional[SandboxConfig] = None # NEW: shell sandbox config

class NodeSkillConfig(BaseModel):
    """Admin-controlled skill configuration for a node."""
    status: Optional[str] = "configured"
    shell: SkillConfig = SkillConfig(enabled=True, sandbox=SandboxConfig())
    sync: SkillConfig = SkillConfig(enabled=True)

# --- Admin Create / Update ---

class AgentNodeCreate(BaseModel):
    """Payload for admin creating a new node registration."""
    node_id: str = Field(..., description="Stable identifier used in the node's config YAML, e.g. 'dev-macbook-m3'")
    display_name: str = Field(..., description="Human-readable name shown in the UI")
    description: Optional[str] = None
    skill_config: NodeSkillConfig = NodeSkillConfig()

class AgentNodeUpdate(BaseModel):
    """Payload for admin updating node configuration."""
    display_name: Optional[str] = None
    description: Optional[str] = None
    skill_config: Optional[NodeSkillConfig] = None
    is_active: Optional[bool] = None

# --- Group Access ---

class NodeAccessGrant(BaseModel):
    """Admin grants a group access to a node."""
    group_id: str
    access_level: str = Field("use", description="'view', 'use', or 'admin'")

class NodeAccessResponse(BaseModel):
    id: int
    node_id: str
    group_id: str
    access_level: str
    granted_at: datetime
    model_config = ConfigDict(from_attributes=True)

# --- Live Stats ---

class AgentNodeStats(BaseModel):
    """Live performance stats reported via heartbeat."""
    active_worker_count: int = 0
    cpu_usage_percent: float = 0.0
    memory_usage_percent: float = 0.0
    cpu_count: int = 0
    memory_used_gb: float = 0.0
    memory_total_gb: float = 0.0
    running: List[str] = []
    # M6 Rich Metrics
    cpu_usage_per_core: List[float] = Field(default_factory=list)
    cpu_freq_mhz: float = 0.0
    memory_available_gb: float = 0.0
    load_avg: List[float] = Field(default_factory=list)

# --- Node Responses ---

class AgentNodeAdminDetail(BaseModel):
    """Full node detail for admin view — includes invite_token and skill config."""
    node_id: str
    display_name: str
    description: Optional[str] = None
    skill_config: dict = {}
    capabilities: dict = {}
    invite_token: Optional[str] = None
    is_active: bool = True
    last_status: str
    last_seen_at: Optional[datetime] = None
    created_at: datetime
    registered_by: str
    group_access: List[NodeAccessResponse] = []
    stats: AgentNodeStats = AgentNodeStats()
    model_config = ConfigDict(from_attributes=True)

class AgentNodeUserView(BaseModel):
    """Node as seen by a user — no invite_token, no admin config details."""
    node_id: str
    display_name: str
    description: Optional[str] = None
    capabilities: dict = {}
    # Which skills are available to this user (derived from skill_config.enabled)
    available_skills: List[str] = []
    last_status: str    # 'online' | 'offline' | 'stale'
    last_seen_at: Optional[datetime] = None
    stats: AgentNodeStats = AgentNodeStats()
    model_config = ConfigDict(from_attributes=True)

class AgentNodeStatusResponse(BaseModel):
    """Full live status of an agent node (used internally)."""
    node_id: str
    display_name: Optional[str] = None
    stats: AgentNodeStats = AgentNodeStats()
    status: str
    connected_at: Optional[str] = None
    last_heartbeat_at: Optional[str] = None

# --- User Node Preferences ---

class NodeDataSourceConfig(BaseModel):
    """How a node should seed its workspace for a session."""
    source: str = Field("empty", description="'empty' | 'server' | 'node_local'")
    path: Optional[str] = None  # root path on node when source='node_local'
    source_node_id: Optional[str] = None  # root node to pull from when source='node_local'

class UserNodePreferences(BaseModel):
    """Stored in User.preferences['nodes']."""
    default_node_ids: List[str] = Field(
        default_factory=list,
        description="Node IDs auto-attached when starting a new session"
    )
    data_source: NodeDataSourceConfig = NodeDataSourceConfig()

# --- Task Dispatch ---

class NodeDispatchRequest(BaseModel):
    """Dispatch a shell action to a specific node."""
    task_id: Optional[str] = None # NEW: Support client-side generated task IDs
    command: str = ""
    session_id: Optional[str] = None
    timeout_ms: int = 30000

class NodeDispatchResponse(BaseModel):
    task_id: str
    status: str   # 'accepted' | 'rejected'
    reason: Optional[str] = None

# --- FS Explorer Schemas ---

class FileNodeInfo(BaseModel):
    path: str
    name: str
    is_dir: bool
    size: int = 0
    hash: Optional[str] = None
    is_synced: bool = False

class DirectoryListing(BaseModel):
    node_id: str
    path: str
    files: List[FileNodeInfo]

class FileReadRequest(BaseModel):
    path: str

class FileWriteRequest(BaseModel):
    path: str
    content: str = "" # Plain text content
    is_dir: bool = False
    session_id: Optional[str] = "__fs_explorer__"

class FileDeleteRequest(BaseModel):
    path: str
    session_id: Optional[str] = "__fs_explorer__"

class FileMoveRequest(BaseModel):
    old_path: str
    new_path: str
    session_id: Optional[str] = "__fs_explorer__"

class FileCopyRequest(BaseModel):
    old_path: str
    new_path: str
    session_id: Optional[str] = "__fs_explorer__"

class FileStatResponse(BaseModel):
    path: str
    name: str
    size: int
    is_dir: bool
    is_file: bool
    is_link: bool
    mtime: float
    exists: bool = True

# Keep backward-compat alias
AgentNodeSummary = AgentNodeUserView


# ---------------------------------------------------------------------------
# Agent Schemas
# ---------------------------------------------------------------------------

class AgentTemplateBase(BaseModel):
    name: str
    description: Optional[str] = None
    system_prompt_path: Optional[str] = None
    system_prompt_content: Optional[str] = None
    max_loop_iterations: int = 20
    co_worker_quality_gate: bool = False
    rework_threshold: int = 80
    max_rework_attempts: int = 3
    user_id: Optional[str] = None

class AgentTemplateCreate(AgentTemplateBase):
    pass

class AgentTemplateResponse(AgentTemplateBase):
    id: str
    model_config = ConfigDict(from_attributes=True)


class AgentInstanceBase(BaseModel):
    template_id: str
    user_id: Optional[str] = None
    session_id: Optional[int] = None
    mesh_node_id: Optional[str] = None
    status: str = "idle"
    current_workspace_jail: Optional[str] = None

class AgentInstanceCreate(AgentInstanceBase):
    pass

class AgentInstanceStatusUpdate(BaseModel):
    status: str

class AgentInstanceResponse(AgentInstanceBase):
    id: str
    last_heartbeat: Optional[datetime] = None
    template: Optional[AgentTemplateResponse] = None
    session: Optional[Session] = None
    
    # Metrics
    total_runs: Optional[int] = 0
    successful_runs: Optional[int] = 0
    total_tokens_accumulated: Optional[int] = 0
    total_input_tokens: Optional[int] = 0
    total_output_tokens: Optional[int] = 0
    total_running_time_seconds: Optional[int] = 0
    tool_call_counts: Optional[dict] = {}
    last_reasoning: Optional[str] = None
    last_error: Optional[str] = None
    evaluation_status: Optional[str] = None
    current_rework_attempt: Optional[int] = 0
    latest_quality_score: Optional[int] = None

    
    model_config = ConfigDict(from_attributes=True)

class AgentTriggerBase(BaseModel):
    instance_id: str
    trigger_type: str
    cron_expression: Optional[str] = None
    interval_seconds: Optional[int] = None
    default_prompt: Optional[str] = None
    webhook_secret: Optional[str] = None
    webhook_mapping_schema: Optional[dict] = None

class AgentTriggerCreate(BaseModel):
    trigger_type: str
    cron_expression: Optional[str] = None
    interval_seconds: Optional[int] = None
    default_prompt: Optional[str] = None

class AgentTriggerResponse(AgentTriggerBase):
    id: str
    model_config = ConfigDict(from_attributes=True)


class DeployAgentRequest(BaseModel):
    """One-click deployment: creates template + session + instance atomically."""
    name: str
    description: Optional[str] = None
    system_prompt: Optional[str] = None
    mesh_node_id: str
    max_loop_iterations: int = 20
    initial_prompt: Optional[str] = None  # First message to kick off the loop
    provider_name: Optional[str] = None
    trigger_type: Optional[str] = "manual"
    cron_expression: Optional[str] = None
    interval_seconds: Optional[int] = None
    default_prompt: Optional[str] = None
    co_worker_quality_gate: Optional[bool] = None
    rework_threshold: Optional[int] = None
    max_rework_attempts: Optional[int] = None

class AgentConfigUpdate(BaseModel):
    """Day 2 Agent Configuration edits"""
    name: Optional[str] = None
    system_prompt: Optional[str] = None
    max_loop_iterations: Optional[int] = None
    mesh_node_id: str
    provider_name: Optional[str] = None
    model_name: Optional[str] = None
    allowed_skill_ids: Optional[List[int]] = None
    restrict_skills: Optional[bool] = None
    is_locked: Optional[bool] = None
    auto_clear_history: Optional[bool] = None
    co_worker_quality_gate: Optional[bool] = None
    rework_threshold: Optional[int] = None
    max_rework_attempts: Optional[int] = None
