use serde::{Deserialize, Serialize};
use time::OffsetDateTime;

use crate::ids::{AgentId, MessageId, SessionId, WorkspaceId};

/// Deterministic hash bucket of an agent id, carried in every SessionKey
/// (§12). Stable across processes and releases: future partitioning relies
/// on it never changing for an existing agent.
pub fn shard_of(agent_id: AgentId) -> u16 {
    let bytes = agent_id.as_uuid().into_bytes();
    let mut acc: u16 = 0;
    for b in bytes {
        acc = acc.wrapping_mul(31).wrapping_add(u16::from(b));
    }
    acc % 16
}

/// A chat session between humans and one agent (spec §14 Session).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub workspace_id: WorkspaceId,
    pub title: String,
    pub shard: u16,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339")]
    pub last_active_at: OffsetDateTime,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MessageRole {
    User,
    Agent,
    System,
}

impl MessageRole {
    pub fn as_str(&self) -> &'static str {
        match self {
            MessageRole::User => "user",
            MessageRole::Agent => "agent",
            MessageRole::System => "system",
        }
    }
}

impl std::str::FromStr for MessageRole {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "user" => Ok(MessageRole::User),
            "agent" => Ok(MessageRole::Agent),
            "system" => Ok(MessageRole::System),
            other => Err(format!("unknown message role: {other}")),
        }
    }
}

/// One transcript message (spec §14 Message). `content` is provider-neutral
/// JSON: `{"text": "..."}` for plain messages.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Message {
    pub id: MessageId,
    pub session_id: SessionId,
    pub seq: i64,
    pub role: MessageRole,
    pub content: serde_json::Value,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StepStatus {
    Running,
    Ok,
    Error,
}

impl StepStatus {
    pub fn as_str(&self) -> &'static str {
        match self {
            StepStatus::Running => "running",
            StepStatus::Ok => "ok",
            StepStatus::Error => "error",
        }
    }
}

impl std::str::FromStr for StepStatus {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "running" => Ok(StepStatus::Running),
            "ok" => Ok(StepStatus::Ok),
            "error" => Ok(StepStatus::Error),
            other => Err(format!("unknown step status: {other}")),
        }
    }
}

/// One tool/reasoning step attached to an agent message — the "N steps"
/// trace (§5b). `taint` records untrusted content sources (§15).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Step {
    pub id: uuid::Uuid,
    pub message_id: MessageId,
    pub seq: i32,
    pub kind: String,
    pub tool_name: Option<String>,
    pub input: Option<serde_json::Value>,
    pub output: Option<serde_json::Value>,
    pub taint: Vec<String>,
    pub status: StepStatus,
}

/// A transcript message with its step trace (`?tools=true` history shape).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MessageWithSteps {
    #[serde(flatten)]
    pub message: Message,
    pub steps: Vec<Step>,
}

/// Lifecycle of one agent run (spec §15 approval interception suspends a
/// run at `AwaitingApproval` until a human decides).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunState {
    Running,
    AwaitingApproval,
    Completed,
    Failed,
    Cancelled,
}

impl RunState {
    pub fn as_str(&self) -> &'static str {
        match self {
            RunState::Running => "running",
            RunState::AwaitingApproval => "awaiting_approval",
            RunState::Completed => "completed",
            RunState::Failed => "failed",
            RunState::Cancelled => "cancelled",
        }
    }
}

impl std::str::FromStr for RunState {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "running" => Ok(RunState::Running),
            "awaiting_approval" => Ok(RunState::AwaitingApproval),
            "completed" => Ok(RunState::Completed),
            "failed" => Ok(RunState::Failed),
            "cancelled" => Ok(RunState::Cancelled),
            other => Err(format!("unknown run state: {other}")),
        }
    }
}

/// One execution of the agent loop within a session.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentRun {
    pub id: uuid::Uuid,
    pub session_id: SessionId,
    pub state: RunState,
    pub last_event_id: i64,
    pub error: Option<String>,
}