use super::*;
use time::OffsetDateTime;

impl PostgresRegistry {
    pub async fn create_stream_session(
        &self,
        request: CreateStreamSessionRequest<'_>,
    ) -> Result<StreamSession, RegistryError> {
        sqlx::query(
            "insert into stream_sessions (
                id,
                workspace_id,
                agent_id,
                operation_id,
                protocol,
                mode,
                status,
                cursor_json,
                state_json,
                expires_at,
                last_poll_at,
                created_at,
                closed_at
             ) values (
                $1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10::timestamptz,
                $11::timestamptz, $12::timestamptz, $13::timestamptz
             )",
        )
        .bind(request.session.id.as_str())
        .bind(request.session.workspace_id.as_str())
        .bind(
            request
                .session
                .agent_id
                .as_ref()
                .map(|value| value.as_str()),
        )
        .bind(request.session.operation_id.as_str())
        .bind(serialize_enum_text(&request.session.protocol, "protocol")?)
        .bind(serialize_enum_text(&request.session.mode, "mode")?)
        .bind(serialize_enum_text(&request.session.status, "status")?)
        .bind(request.session.cursor.clone().map(Json))
        .bind(Json(request.session.state.clone()))
        .bind(request.session.expires_at)
        .bind(request.session.last_poll_at)
        .bind(request.session.created_at)
        .bind(request.session.closed_at)
        .execute(&self.pool)
        .await?;

        Ok(request.session.clone())
    }

    pub async fn get_stream_session(
        &self,
        id: &StreamSessionId,
    ) -> Result<Option<StreamSession>, RegistryError> {
        let row = sqlx::query!(
            "select
                id,
                workspace_id,
                agent_id,
                operation_id,
                protocol,
                mode,
                status,
                cursor_json,
                state_json,
                expires_at as \"expires_at!: OffsetDateTime\",
                last_poll_at as \"last_poll_at: OffsetDateTime\",
                created_at as \"created_at!: OffsetDateTime\",
                closed_at as \"closed_at: OffsetDateTime\"
             from stream_sessions
             where id = $1",
            id.as_str(),
        )
        .fetch_optional(&self.pool)
        .await?;

        row.map(|row| {
            Ok(StreamSession {
                id: StreamSessionId::new(row.id),
                workspace_id: WorkspaceId::new(row.workspace_id),
                agent_id: row.agent_id.map(AgentId::new),
                operation_id: OperationId::new(row.operation_id),
                protocol: deserialize_enum_text(&row.protocol, "protocol")?,
                mode: deserialize_enum_text(&row.mode, "mode")?,
                status: deserialize_enum_text(&row.status, "status")?,
                cursor: row.cursor_json,
                state: row.state_json,
                expires_at: row.expires_at,
                last_poll_at: row.last_poll_at,
                created_at: row.created_at,
                closed_at: row.closed_at,
            })
        })
        .transpose()
    }

    pub async fn update_stream_session_state(
        &self,
        request: UpdateStreamSessionStateRequest<'_>,
    ) -> Result<StreamSession, RegistryError> {
        validate_stream_session_transition(
            request.session_id,
            request.current_status,
            request.next_status,
        )?;

        let row = sqlx::query(
            "update stream_sessions
                set status = $3,
                    cursor_json = $4::jsonb,
                    state_json = $5::jsonb,
                    expires_at = coalesce($6::timestamptz, expires_at),
                    last_poll_at = coalesce($7::timestamptz, last_poll_at),
                    closed_at = case
                        when $8::timestamptz is null then closed_at
                        else $8::timestamptz
                    end
              where id = $1
                and status = $2
            returning
                id,
                workspace_id,
                agent_id,
                operation_id,
                protocol,
                mode,
                status,
                cursor_json,
                state_json,
                expires_at,
                last_poll_at,
                created_at,
                closed_at",
        )
        .bind(request.session_id.as_str())
        .bind(serialize_enum_text(&request.current_status, "status")?)
        .bind(serialize_enum_text(&request.next_status, "status")?)
        .bind(request.cursor.cloned().map(Json))
        .bind(Json(request.state.clone()))
        .bind(request.expires_at)
        .bind(request.last_poll_at)
        .bind(request.closed_at)
        .fetch_optional(&self.pool)
        .await?;

        match row.as_ref() {
            Some(row) => map_stream_session(row),
            None => {
                if self.get_stream_session(request.session_id).await?.is_some() {
                    Err(RegistryError::InvalidStreamSessionTransition {
                        session_id: request.session_id.as_str().to_owned(),
                        from: serialize_enum_text(&request.current_status, "status")?,
                        to: serialize_enum_text(&request.next_status, "status")?,
                    })
                } else {
                    Err(RegistryError::StreamSessionNotFound {
                        session_id: request.session_id.as_str().to_owned(),
                    })
                }
            }
        }
    }

    pub async fn close_stream_session(
        &self,
        id: &StreamSessionId,
        now: &OffsetDateTime,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update stream_sessions
                set status = 'stopped',
                    last_poll_at = $2::timestamptz,
                    closed_at = $2::timestamptz
              where id = $1
                and status in ('created', 'running')",
        )
        .bind(id.as_str())
        .bind(now)
        .execute(&self.pool)
        .await?;

        if result.rows_affected() > 0 {
            return Ok(());
        }

        match self.get_stream_session(id).await? {
            Some(existing) => Err(RegistryError::InvalidStreamSessionTransition {
                session_id: id.as_str().to_owned(),
                from: serialize_enum_text(&existing.status, "status")?,
                to: "stopped".to_owned(),
            }),
            None => Err(RegistryError::StreamSessionNotFound {
                session_id: id.as_str().to_owned(),
            }),
        }
    }

    pub async fn list_stream_sessions(
        &self,
        filter: StreamSessionFilter<'_>,
    ) -> Result<Page<StreamSession>, RegistryError> {
        let status = filter
            .status
            .as_ref()
            .map(|value| serialize_enum_text(value, "status"))
            .transpose()?;
        let mode = filter
            .mode
            .as_ref()
            .map(|value| serialize_enum_text(value, "mode"))
            .transpose()?;
        let limit = i64::from(filter.limit);

        let rows = sqlx::query!(
            "select
                id,
                workspace_id,
                agent_id,
                operation_id,
                protocol,
                mode,
                status,
                cursor_json,
                state_json,
                expires_at as \"expires_at!: OffsetDateTime\",
                last_poll_at as \"last_poll_at: OffsetDateTime\",
                created_at as \"created_at!: OffsetDateTime\",
                closed_at as \"closed_at: OffsetDateTime\"
             from stream_sessions
             where workspace_id = $1
               and ($2::text is null or agent_id = $2)
               and ($3::text is null or operation_id = $3)
               and ($4::text is null or status = $4)
               and ($5::text is null or mode = $5)
             order by created_at desc
             limit $6",
            filter.workspace_id.as_str(),
            filter.agent_id.map(|value| value.as_str()),
            filter.operation_id.map(|value| value.as_str()),
            status,
            mode,
            limit,
        )
        .fetch_all(&self.pool)
        .await?;

        let items = rows
            .into_iter()
            .map(|row| {
                Ok(StreamSession {
                    id: StreamSessionId::new(row.id),
                    workspace_id: WorkspaceId::new(row.workspace_id),
                    agent_id: row.agent_id.map(AgentId::new),
                    operation_id: OperationId::new(row.operation_id),
                    protocol: deserialize_enum_text(&row.protocol, "protocol")?,
                    mode: deserialize_enum_text(&row.mode, "mode")?,
                    status: deserialize_enum_text(&row.status, "status")?,
                    cursor: row.cursor_json,
                    state: row.state_json,
                    expires_at: row.expires_at,
                    last_poll_at: row.last_poll_at,
                    created_at: row.created_at,
                    closed_at: row.closed_at,
                })
            })
            .collect::<Result<Vec<_>, RegistryError>>()?;
        let total = items.len() as u64;

        Ok(Page { items, total })
    }

    pub async fn delete_expired_stream_sessions(
        &self,
        now: &OffsetDateTime,
    ) -> Result<u64, RegistryError> {
        let result = sqlx::query(
            "delete from stream_sessions
              where expires_at <= $1::timestamptz",
        )
        .bind(now)
        .execute(&self.pool)
        .await?;

        Ok(result.rows_affected())
    }

    pub async fn create_async_job(
        &self,
        request: CreateAsyncJobRequest<'_>,
    ) -> Result<AsyncJobHandle, RegistryError> {
        sqlx::query(
            "insert into async_jobs (
                id,
                workspace_id,
                agent_id,
                operation_id,
                status,
                progress_json,
                result_json,
                error_json,
                expires_at,
                last_poll_at,
                created_at,
                updated_at,
                finished_at
             ) values (
                $1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
                $10::timestamptz, $11::timestamptz, $12::timestamptz, $13::timestamptz
             )",
        )
        .bind(request.job.id.as_str())
        .bind(request.job.workspace_id.as_str())
        .bind(request.job.agent_id.as_ref().map(|value| value.as_str()))
        .bind(request.job.operation_id.as_str())
        .bind(serialize_enum_text(&request.job.status, "status")?)
        .bind(Json(request.job.progress.clone()))
        .bind(request.job.result.clone().map(Json))
        .bind(request.job.error.clone().map(Json))
        .bind(request.job.expires_at)
        .bind(request.job.last_poll_at)
        .bind(request.job.created_at)
        .bind(request.job.updated_at)
        .bind(request.job.finished_at)
        .execute(&self.pool)
        .await?;

        Ok(request.job.clone())
    }

    pub async fn get_async_job(
        &self,
        id: &AsyncJobId,
    ) -> Result<Option<AsyncJobHandle>, RegistryError> {
        let row = sqlx::query!(
            "select
                id,
                workspace_id,
                agent_id,
                operation_id,
                status,
                progress_json,
                result_json,
                error_json,
                expires_at as \"expires_at: OffsetDateTime\",
                last_poll_at as \"last_poll_at: OffsetDateTime\",
                created_at as \"created_at!: OffsetDateTime\",
                updated_at as \"updated_at!: OffsetDateTime\",
                finished_at as \"finished_at: OffsetDateTime\"
             from async_jobs
             where id = $1",
            id.as_str(),
        )
        .fetch_optional(&self.pool)
        .await?;

        row.map(|row| {
            Ok(AsyncJobHandle {
                id: AsyncJobId::new(row.id),
                workspace_id: WorkspaceId::new(row.workspace_id),
                agent_id: row.agent_id.map(AgentId::new),
                operation_id: OperationId::new(row.operation_id),
                status: deserialize_enum_text(&row.status, "status")?,
                progress: row.progress_json,
                result: row.result_json,
                error: row.error_json,
                expires_at: row.expires_at,
                last_poll_at: row.last_poll_at,
                created_at: row.created_at,
                updated_at: row.updated_at,
                finished_at: row.finished_at,
            })
        })
        .transpose()
    }

    pub async fn update_async_job_status(
        &self,
        request: UpdateAsyncJobStatusRequest<'_>,
    ) -> Result<AsyncJobHandle, RegistryError> {
        validate_async_job_transition(request.job_id, request.current_status, request.next_status)?;

        let row = sqlx::query(
            "update async_jobs
                set status = $3,
                    progress_json = $4::jsonb,
                    result_json = $5::jsonb,
                error_json = $6::jsonb,
                expires_at = coalesce($7::timestamptz, expires_at),
                updated_at = $8::timestamptz,
                    finished_at = case
                        when $9::timestamptz is null then finished_at
                        else $9::timestamptz
                    end
              where id = $1
                and status = $2
            returning
                id,
                workspace_id,
                agent_id,
                operation_id,
                status,
                progress_json,
                result_json,
                error_json,
                expires_at,
                last_poll_at,
                created_at,
                updated_at,
                finished_at",
        )
        .bind(request.job_id.as_str())
        .bind(serialize_enum_text(&request.current_status, "status")?)
        .bind(serialize_enum_text(&request.next_status, "status")?)
        .bind(Json(request.progress.clone()))
        .bind(request.result.cloned().map(Json))
        .bind(request.error.cloned().map(Json))
        .bind(request.expires_at)
        .bind(request.updated_at)
        .bind(request.finished_at)
        .fetch_optional(&self.pool)
        .await?;

        match row.as_ref() {
            Some(row) => map_async_job(row),
            None => {
                if self.get_async_job(request.job_id).await?.is_some() {
                    Err(RegistryError::InvalidAsyncJobTransition {
                        job_id: request.job_id.as_str().to_owned(),
                        from: serialize_enum_text(&request.current_status, "status")?,
                        to: serialize_enum_text(&request.next_status, "status")?,
                    })
                } else {
                    Err(RegistryError::AsyncJobNotFound {
                        job_id: request.job_id.as_str().to_owned(),
                    })
                }
            }
        }
    }

    pub async fn touch_async_job_poll(
        &self,
        id: &AsyncJobId,
        now: &OffsetDateTime,
    ) -> Result<AsyncJobHandle, RegistryError> {
        let row = sqlx::query(
            "update async_jobs
                set last_poll_at = $2::timestamptz
              where id = $1
            returning
                id,
                workspace_id,
                agent_id,
                operation_id,
                status,
                progress_json,
                result_json,
                error_json,
                expires_at,
                last_poll_at,
                created_at,
                updated_at,
                finished_at",
        )
        .bind(id.as_str())
        .bind(now)
        .fetch_optional(&self.pool)
        .await?;

        match row.as_ref() {
            Some(row) => map_async_job(row),
            None => Err(RegistryError::AsyncJobNotFound {
                job_id: id.as_str().to_owned(),
            }),
        }
    }

    pub async fn cancel_async_job(
        &self,
        id: &AsyncJobId,
        now: &OffsetDateTime,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update async_jobs
                set status = 'cancelled',
                    updated_at = $2::timestamptz,
                    finished_at = $2::timestamptz
              where id = $1
                and status in ('created', 'running')",
        )
        .bind(id.as_str())
        .bind(now)
        .execute(&self.pool)
        .await?;

        if result.rows_affected() > 0 {
            return Ok(());
        }

        match self.get_async_job(id).await? {
            Some(existing) => Err(RegistryError::InvalidAsyncJobTransition {
                job_id: id.as_str().to_owned(),
                from: serialize_enum_text(&existing.status, "status")?,
                to: "cancelled".to_owned(),
            }),
            None => Err(RegistryError::AsyncJobNotFound {
                job_id: id.as_str().to_owned(),
            }),
        }
    }

    pub async fn list_async_jobs(
        &self,
        filter: AsyncJobFilter<'_>,
    ) -> Result<Page<AsyncJobHandle>, RegistryError> {
        let status = filter
            .status
            .as_ref()
            .map(|value| serialize_enum_text(value, "status"))
            .transpose()?;
        let limit = i64::from(filter.limit);

        let rows = sqlx::query!(
            "select
                id,
                workspace_id,
                agent_id,
                operation_id,
                status,
                progress_json,
                result_json,
                error_json,
                expires_at as \"expires_at: OffsetDateTime\",
                last_poll_at as \"last_poll_at: OffsetDateTime\",
                created_at as \"created_at!: OffsetDateTime\",
                updated_at as \"updated_at!: OffsetDateTime\",
                finished_at as \"finished_at: OffsetDateTime\"
             from async_jobs
             where workspace_id = $1
               and ($2::text is null or agent_id = $2)
               and ($3::text is null or operation_id = $3)
               and ($4::text is null or status = $4)
             order by updated_at desc
             limit $5",
            filter.workspace_id.as_str(),
            filter.agent_id.map(|value| value.as_str()),
            filter.operation_id.map(|value| value.as_str()),
            status,
            limit,
        )
        .fetch_all(&self.pool)
        .await?;

        let items = rows
            .into_iter()
            .map(|row| {
                Ok(AsyncJobHandle {
                    id: AsyncJobId::new(row.id),
                    workspace_id: WorkspaceId::new(row.workspace_id),
                    agent_id: row.agent_id.map(AgentId::new),
                    operation_id: OperationId::new(row.operation_id),
                    status: deserialize_enum_text(&row.status, "status")?,
                    progress: row.progress_json,
                    result: row.result_json,
                    error: row.error_json,
                    expires_at: row.expires_at,
                    last_poll_at: row.last_poll_at,
                    created_at: row.created_at,
                    updated_at: row.updated_at,
                    finished_at: row.finished_at,
                })
            })
            .collect::<Result<Vec<_>, RegistryError>>()?;
        let total = items.len() as u64;

        Ok(Page { items, total })
    }

    pub async fn delete_expired_async_jobs(
        &self,
        now: &OffsetDateTime,
    ) -> Result<u64, RegistryError> {
        let result = sqlx::query(
            "delete from async_jobs
              where expires_at is not null
                and expires_at <= $1::timestamptz",
        )
        .bind(now)
        .execute(&self.pool)
        .await?;

        Ok(result.rows_affected())
    }
}