use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::{
    ids::{AgentId, AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
    protocol::Protocol,
    streaming::ExecutionMode,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamStatus {
    Created,
    Running,
    Stopped,
    Failed,
    Expired,
}

impl StreamStatus {
    pub fn is_terminal(self) -> bool {
        matches!(self, Self::Stopped | Self::Failed | Self::Expired)
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamSession {
    pub id: StreamSessionId,
    pub workspace_id: WorkspaceId,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<AgentId>,
    pub operation_id: OperationId,
    pub protocol: Protocol,
    pub mode: ExecutionMode,
    pub status: StreamStatus,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cursor: Option<Value>,
    pub state: Value,
    pub expires_at: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub last_poll_at: Option<String>,
    pub created_at: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub closed_at: Option<String>,
}

impl StreamSession {
    pub fn is_expired(&self, now: &str) -> bool {
        self.expires_at.as_str() <= now
    }

    pub fn can_poll(&self, now: &str) -> bool {
        !self.status.is_terminal() && !self.is_expired(now)
    }

    pub fn mark_polled(&mut self, now: impl Into<String>) {
        self.last_poll_at = Some(now.into());
    }

    pub fn mark_closed(&mut self, now: impl Into<String>) {
        let now = now.into();
        self.status = StreamStatus::Stopped;
        self.last_poll_at = Some(now.clone());
        self.closed_at = Some(now);
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JobStatus {
    Created,
    Running,
    Completed,
    Failed,
    Cancelled,
    Expired,
}

impl JobStatus {
    pub fn is_terminal(self) -> bool {
        matches!(
            self,
            Self::Completed | Self::Failed | Self::Cancelled | Self::Expired
        )
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AsyncJobHandle {
    pub id: AsyncJobId,
    pub workspace_id: WorkspaceId,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<AgentId>,
    pub operation_id: OperationId,
    pub status: JobStatus,
    pub progress: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<String>,
    pub created_at: String,
    pub updated_at: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finished_at: Option<String>,
}

impl AsyncJobHandle {
    pub fn is_finished(&self) -> bool {
        self.status.is_terminal()
    }

    pub fn can_cancel(&self) -> bool {
        matches!(self.status, JobStatus::Created | JobStatus::Running)
    }

    pub fn mark_finished(&mut self, now: impl Into<String>, result: Value) {
        let now = now.into();
        self.status = JobStatus::Completed;
        self.result = Some(result);
        self.error = None;
        self.updated_at = now.clone();
        self.finished_at = Some(now);
    }

    pub fn mark_failed(&mut self, now: impl Into<String>, error: Value) {
        let now = now.into();
        self.status = JobStatus::Failed;
        self.error = Some(error);
        self.updated_at = now.clone();
        self.finished_at = Some(now);
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::{AsyncJobHandle, JobStatus, StreamSession, StreamStatus};
    use crate::{
        ids::{AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
        protocol::Protocol,
        streaming::ExecutionMode,
    };

    #[test]
    fn stream_session_tracks_poll_and_close_transitions() {
        let mut session = StreamSession {
            id: StreamSessionId::new("stream_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            protocol: Protocol::Rest,
            mode: ExecutionMode::Session,
            status: StreamStatus::Running,
            cursor: None,
            state: json!({"cursor":"abc"}),
            expires_at: "2026-04-06T12:05:00Z".to_owned(),
            last_poll_at: None,
            created_at: "2026-04-06T12:00:00Z".to_owned(),
            closed_at: None,
        };

        assert!(session.can_poll("2026-04-06T12:01:00Z"));

        session.mark_polled("2026-04-06T12:01:00Z");
        session.mark_closed("2026-04-06T12:02:00Z");

        assert_eq!(session.status, StreamStatus::Stopped);
        assert_eq!(session.closed_at.as_deref(), Some("2026-04-06T12:02:00Z"));
        assert!(!session.can_poll("2026-04-06T12:03:00Z"));
    }

    #[test]
    fn async_job_tracks_finish_and_failure() {
        let mut job = AsyncJobHandle {
            id: AsyncJobId::new("job_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            status: JobStatus::Running,
            progress: json!({"percent": 60}),
            result: None,
            error: None,
            expires_at: Some("2026-04-06T12:05:00Z".to_owned()),
            created_at: "2026-04-06T12:00:00Z".to_owned(),
            updated_at: "2026-04-06T12:00:00Z".to_owned(),
            finished_at: None,
        };

        assert!(job.can_cancel());

        job.mark_finished("2026-04-06T12:01:00Z", json!({"ok": true}));
        assert!(job.is_finished());
        assert_eq!(job.status, JobStatus::Completed);

        let mut failed_job = job.clone();
        failed_job.status = JobStatus::Running;
        failed_job.result = None;
        failed_job.finished_at = None;

        failed_job.mark_failed("2026-04-06T12:02:00Z", json!({"message": "boom"}));
        assert_eq!(failed_job.status, JobStatus::Failed);
        assert_eq!(
            failed_job.finished_at.as_deref(),
            Some("2026-04-06T12:02:00Z")
        );
    }
}
