use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{PgPool, Pool, Postgres};
use uuid::Uuid;

use crate::{
    datastore::Datastore,
    error::JobResult,
    job::{Job, JobStatus},
};

/// PostgreSQL implementation of the datastore
#[derive(Clone)]
pub struct PostgresDatastore {
    pool: PgPool,
}

impl PostgresDatastore {
    /// Create a new PostgreSQL datastore
    pub async fn new(pool: Pool<Postgres>) -> JobResult<Self> {
            Ok(Self { pool })
    }

    /// Run database migrations
    pub async fn migrate(&self) -> JobResult<()> {
        let migrations = [
            include_str!("migrations/20251201000000_create_jobs_table.sql"),
            // future migrations
        ];

        for sql in migrations {
            sqlx::raw_sql(sql).execute(&self.pool).await?;
        }

        Ok(())
    }

    /// Get a reference to the connection pool
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }
}

#[async_trait]
impl Datastore for PostgresDatastore {
    async fn insert_job(&self, record: Job) -> JobResult<Uuid> {
        let job_schedule = serde_json::to_value(&record.job_schedule)?;
        let retry = serde_json::to_value(&record.retry)?;

        sqlx::query(
            r#"
            INSERT INTO _runat_jobs (
                id, job_type, payload, job_schedule, status, retry,
                attempts, max_attempts, scheduled_at, locked_until, error,
                created_at, updated_at, started_at, completed_at
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
            "#
        )
        .bind(record.id)
        .bind(&record.job_type)
        .bind(&record.payload)
        .bind(job_schedule)
        .bind(record.status.to_i32())
        .bind(retry)
        .bind(record.attempts)
        .bind(record.max_attempts)
        .bind(record.scheduled_at)
        .bind(record.locked_until)
        .bind(&record.error)
        .bind(record.created_at)
        .bind(record.updated_at)
        .bind(record.started_at)
        .bind(record.completed_at)
        .execute(&self.pool)
        .await?;

        Ok(record.id)
    }

    async fn get_job(&self, job_id: Uuid) -> JobResult<Option<Job>> {
        let row = sqlx::query_as::<_, Job>(
            r#"
            SELECT id, job_type, payload, job_schedule, status, retry,
                   attempts, max_attempts, scheduled_at, locked_until, error,
                   created_at, updated_at, started_at, completed_at
            FROM _runat_jobs
            WHERE id = $1
            "#
        )
        .bind(job_id)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row)
    }

    async fn get_all_jobs(&self) -> JobResult<Vec<Job>> {
        let rows = sqlx::query_as::<_, Job>(
            r#"
            SELECT id, job_type, payload, job_schedule, status, retry,
                   attempts, max_attempts, scheduled_at, locked_until, error,
                   created_at, updated_at, started_at, completed_at
            FROM _runat_jobs
            ORDER BY scheduled_at ASC
            "#
        )
        .fetch_all(&self.pool)
        .await?;

        Ok(rows)
    }

    async fn update_job(&self, job: Job) -> JobResult<()> {
        let job_schedule = serde_json::to_value(&job.job_schedule)?;
        let retry = serde_json::to_value(&job.retry)?;

        sqlx::query(
            r#"
            UPDATE _runat_jobs
            SET payload = $2,
                job_schedule = $3,
                status = $4,
                retry = $5,
                attempts = $6,
                scheduled_at = $7,
                locked_until = $8,
                error = $9,
                started_at = $10,
                completed_at = $11,
                updated_at = NOW()
            WHERE id = $1
            "#
        )
        .bind(job.id)
        .bind(job.payload.clone())
        .bind(job_schedule)
        .bind(job.status.to_i32())
        .bind(retry)
        .bind(job.attempts)
        .bind(job.scheduled_at)
        .bind(job.locked_until)
        .bind(job.error.clone())
        .bind(job.started_at)
        .bind(job.completed_at)
        .execute(&self.pool)
        .await?;

    Ok(())
    }

    async fn delete_job(&self, job_id: Uuid) -> JobResult<()> {
        sqlx::query("DELETE FROM _runat_jobs WHERE id = $1")
            .bind(job_id)
            .execute(&self.pool)
            .await?;

        Ok(())
    }

    async fn claim_next_job(
        &self,
        locked_until: DateTime<Utc>,
    ) -> JobResult<Option<Job>> {
        let row = sqlx::query_as::<_, Job>(
            r#"
            UPDATE _runat_jobs
            SET status = $1,
                locked_until = $2,
                started_at = NOW(),
                updated_at = NOW()
            WHERE id = (
                SELECT id
                FROM _runat_jobs
                WHERE status = $3
                  AND scheduled_at <= NOW()
                  AND (locked_until IS NULL OR locked_until < NOW())
                ORDER BY scheduled_at ASC
                LIMIT 1
                FOR UPDATE SKIP LOCKED
            )
            RETURNING id, job_type, payload, job_schedule, status, retry,
                      attempts, max_attempts, scheduled_at, locked_until, error,
                      created_at, updated_at, started_at, completed_at
            "#
        )
        .bind(JobStatus::Running.to_i32())
        .bind(locked_until)
        .bind(JobStatus::Queued.to_i32())
        .fetch_optional(&self.pool)
        .await?;

        Ok(row)
    }
}