# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Data models for the Agent service."""

from typing import List, Dict, Any, Optional
from enum import Enum
from pydantic import BaseModel, Field
from modules.shared.timeUtils import getUtcTimestamp
import uuid


class AgentStatusEnum(str, Enum):
    RUNNING = "running"
    COMPLETED = "completed"
    MAX_ROUNDS_REACHED = "maxRoundsReached"
    BUDGET_EXCEEDED = "budgetExceeded"
    ERROR = "error"
    STOPPED = "stopped"


class AgentEventTypeEnum(str, Enum):
    MESSAGE = "message"
    CHUNK = "chunk"
    TOOL_CALL = "toolCall"
    TOOL_RESULT = "toolResult"
    AGENT_PROGRESS = "agentProgress"
    AGENT_SUMMARY = "agentSummary"
    FILE_CREATED = "fileCreated"
    FILE_UPDATED = "fileUpdated"
    FILE_EDIT_PROPOSAL = "fileEditProposal"
    FILE_VERSION = "fileVersion"
    FILE_EDIT_REJECTED = "fileEditRejected"
    DATA_SOURCE_ACCESS = "dataSourceAccess"
    VOICE_RESPONSE = "voiceResponse"
    FINAL = "final"
    ERROR = "error"


class ToolDefinition(BaseModel):
    """Schema for a tool available to the agent."""
    name: str = Field(description="Unique tool name")
    description: str = Field(description="What this tool does")
    parameters: Dict[str, Any] = Field(
        default_factory=dict,
        description="JSON Schema for tool parameters"
    )
    readOnly: bool = Field(
        default=False,
        description="If True, tool can run in parallel with other readOnly tools"
    )
    featureType: Optional[str] = Field(
        default=None,
        description="Feature scope for this tool (None = available to all)"
    )
    toolSet: Optional[str] = Field(
        default=None,
        description="Tool-set scope (None = available to all sets, e.g. 'core', 'workspace')"
    )


class ToolCallRequest(BaseModel):
    """A tool call requested by the AI model."""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str
    args: Dict[str, Any] = Field(default_factory=dict)


class ToolResult(BaseModel):
    """Result from executing a tool."""
    toolCallId: str
    toolName: str
    success: bool = True
    data: str = ""
    error: Optional[str] = None
    durationMs: int = 0
    sideEvents: Optional[List[Dict[str, Any]]] = None


class AgentEvent(BaseModel):
    """Event emitted during agent execution for SSE streaming."""
    type: AgentEventTypeEnum
    content: Optional[str] = None
    data: Optional[Dict[str, Any]] = None


class AgentConfig(BaseModel):
    """Configuration for an agent run."""
    maxRounds: int = Field(default=25, ge=1, le=100)
    maxCostCHF: Optional[float] = Field(default=None, ge=0.0)
    toolSet: str = Field(default="core")
    temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)


class AgentState(BaseModel):
    """Tracks state across an agent loop execution."""
    workflowId: str
    currentRound: int = 0
    maxRounds: int = 25
    totalAiCalls: int = 0
    totalToolCalls: int = 0
    totalCostCHF: float = 0.0
    totalProcessingTime: float = 0.0
    status: AgentStatusEnum = AgentStatusEnum.RUNNING
    abortReason: Optional[str] = None


class ToolCallLog(BaseModel):
    """Log of a single tool call for observability."""
    toolName: str
    args: Dict[str, Any] = Field(default_factory=dict)
    success: bool = True
    durationMs: int = 0
    error: Optional[str] = None
    resultData: str = Field(default="", description="Short result summary for artifact tracking")


class AgentRoundLog(BaseModel):
    """Log of a single agent round for observability."""
    roundNumber: int
    aiModel: str = ""
    inputTokens: int = 0
    outputTokens: int = 0
    costCHF: float = 0.0
    toolCalls: List[ToolCallLog] = Field(default_factory=list)
    durationMs: int = 0


class AgentTrace(BaseModel):
    """Full trace of an agent workflow for observability."""
    workflowId: str
    userId: str = ""
    featureInstanceId: str = ""
    startedAt: float = Field(default_factory=getUtcTimestamp)
    completedAt: Optional[float] = None
    status: AgentStatusEnum = AgentStatusEnum.RUNNING
    totalRounds: int = 0
    totalToolCalls: int = 0
    totalCostCHF: float = 0.0
    abortReason: Optional[str] = None
    rounds: List[AgentRoundLog] = Field(default_factory=list)


class PendingFileEdit(BaseModel):
    """A proposed file edit awaiting user approval."""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    fileId: str
    fileName: str
    mimeType: str = ""
    oldContent: str = ""
    newContent: str = ""
    status: str = Field(default="pending", description="pending | accepted | rejected")
    toolCallId: str = ""
    workflowId: str = ""