use anyhow::Context;
use jiff::Timestamp;
use tokio_util::sync::CancellationToken;

use crate::{
    Operator, OperatorState,
    control_plane::backing_store::BackingStore,
    manifests::{Action, ManifestName, ManifestState, ManifestStatusState, WorkerId},
    reconcile_queue::ReconcileQueue,
};

#[derive(Clone)]
pub struct Reconciler<T: Operator> {
    worker_id: WorkerId,
    store: BackingStore<T::Specifications>,
    operator: OperatorState<T>,
    reconcile_queue: ReconcileQueue<T::Specifications>,
}

impl<T: Operator> Reconciler<T> {
    pub fn new(
        worker_id: WorkerId,
        store: &BackingStore<T::Specifications>,
        operator: OperatorState<T>,
    ) -> Self {
        Self {
            worker_id,
            store: store.clone(),
            operator,
            reconcile_queue: ReconcileQueue::new(),
        }
    }

    /// Run the reconciler. This starts three concurrent tasks:
    /// 1. A sync task that periodically checks for new/changed manifests and enqueues them
    /// 2. A worker task that processes the queue and runs reconciliations
    /// 3. A periodic resync task that enqueues all manifests at a configurable interval
    pub async fn run(&self, cancellation_token: &CancellationToken) -> anyhow::Result<()> {
        let now = Timestamp::now();
        tracing::debug!(%self.worker_id, %now, "starting reconciler");

        tokio::select! {
            result = self.sync_manifests(cancellation_token) => {
                result.context("sync manifests task failed")?;
            }
            result = self.process_queue(cancellation_token) => {
                result.context("process queue task failed")?;
            }
            result = self.periodic_resync(cancellation_token) => {
                result.context("periodic resync task failed")?;
            }
            _ = cancellation_token.cancelled() => {
                tracing::debug!("reconciler received cancellation");
            }
        }

        tracing::debug!("reconciler shutting down");
        Ok(())
    }

    /// Periodically sync manifests from the backing store.
    /// Acquires leases and enqueues manifests that need reconciliation.
    async fn sync_manifests(&self, cancellation_token: &CancellationToken) -> anyhow::Result<()> {
        loop {
            if cancellation_token.is_cancelled() {
                break;
            }

            let now = Timestamp::now();
            tracing::trace!(%self.worker_id, %now, "syncing manifests");

            if let Err(e) = self.sync_once().await {
                tracing::warn!(error = %e, "failed to sync manifests");
            }

            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
        }

        Ok(())
    }

    /// Single sync iteration - check for manifests, acquire leases, enqueue work.
    async fn sync_once(&self) -> anyhow::Result<()> {
        for manifest_state in self.store.get_owned_and_potential_leases().await? {
            let manifest_name = manifest_state.manifest.name.clone();

            match &manifest_state.lease {
                Some(lease) if lease.owner == self.worker_id => {
                    // We own the lease, update it
                    self.store
                        .update_lease(&manifest_state)
                        .await
                        .context("update lease")?;

                    // Check if there are unhandled changes
                    if self.needs_reconciliation(&manifest_state) {
                        self.reconcile_queue.enqueue(manifest_name).await;
                    }
                }
                None => {
                    // No lease, try to acquire
                    self.store
                        .acquire_lease(&manifest_state, &self.worker_id)
                        .await
                        .context("acquire lease")?;

                    // Enqueue for reconciliation
                    self.reconcile_queue.enqueue(manifest_name).await;
                }
                _ => {
                    // Someone else owns the lease, skip
                    continue;
                }
            }
        }

        Ok(())
    }

    /// Periodically enqueue all manifests for reconciliation.
    /// This ensures manifests are reconciled at least every resync_interval.
    async fn periodic_resync(&self, cancellation_token: &CancellationToken) -> anyhow::Result<()> {
        let resync_interval = self.operator.config().resync_interval;
        tracing::debug!(%self.worker_id, ?resync_interval, "starting periodic resync");

        loop {
            tokio::select! {
                _ = tokio::time::sleep(resync_interval) => {}
                _ = cancellation_token.cancelled() => break,
            }

            if cancellation_token.is_cancelled() {
                break;
            }

            tracing::debug!(%self.worker_id, "periodic resync: enqueuing all manifests");

            if let Err(e) = self.enqueue_all_manifests().await {
                tracing::warn!(error = %e, "failed to enqueue all manifests for resync");
            }
        }

        Ok(())
    }

    /// Enqueue all manifests that we own for reconciliation.
    async fn enqueue_all_manifests(&self) -> anyhow::Result<()> {
        for manifest_state in self.store.get_owned_and_potential_leases().await? {
            // Only enqueue manifests we own
            if let Some(lease) = &manifest_state.lease
                && lease.owner == self.worker_id
            {
                self.reconcile_queue
                    .enqueue(manifest_state.manifest.name.clone())
                    .await;
            }
        }

        Ok(())
    }

    /// Check if a manifest needs reconciliation.
    fn needs_reconciliation(&self, manifest: &ManifestState<T::Specifications>) -> bool {
        // Has unhandled changes
        if let Some(change) = manifest.status.changes.first()
            && !change.handled
        {
            return true;
        }

        false
    }

    /// Process the reconciliation queue.
    /// Takes items from the queue and reconciles them, re-enqueuing with delay if needed.
    async fn process_queue(&self, cancellation_token: &CancellationToken) -> anyhow::Result<()> {
        loop {
            if cancellation_token.is_cancelled() {
                break;
            }

            // Wait for the next item from the queue
            let manifest_name = tokio::select! {
                name = self.reconcile_queue.take() => name,
                _ = cancellation_token.cancelled() => break,
            };

            tracing::trace!(%self.worker_id, %manifest_name, "processing manifest from queue");

            if let Err(e) = self.reconcile_manifest(&manifest_name).await {
                tracing::warn!(error = %e, %manifest_name, "failed to reconcile manifest");
                // On error, requeue with a short delay
                self.requeue_with_delay(&manifest_name, std::time::Duration::from_secs(5))
                    .await;
            }
        }

        Ok(())
    }

    /// Reconcile a single manifest by name.
    async fn reconcile_manifest(&self, manifest_name: &ManifestName) -> anyhow::Result<()> {
        // Fetch the current manifest state
        let Some(mut manifest) = self.store.get(manifest_name).await? else {
            tracing::debug!(%manifest_name, "manifest not found, skipping");
            return Ok(());
        };

        // Verify we still own the lease
        match &manifest.lease {
            Some(lease) if lease.owner == self.worker_id => {}
            _ => {
                tracing::debug!(%manifest_name, "we don't own the lease, skipping");
                return Ok(());
            }
        }

        let now = Timestamp::now();
        manifest.status.status = ManifestStatusState::Started;
        manifest.updated = now;

        // Run the reconciliation
        let action = match self.operator.reconcile(&mut manifest).await {
            Ok(action) => action,
            Err(e) => {
                manifest.status.status = ManifestStatusState::Errored;
                manifest.updated = Timestamp::now();

                // Let the operator handle the error
                self.operator.on_error(&mut manifest, &e).await
            }
        };

        // Update status
        let now = Timestamp::now();
        manifest.status.status = ManifestStatusState::Running;
        manifest.updated = now;

        // Mark change as handled
        if let Some(change) = manifest.status.changes.first_mut() {
            change.handled = true;
        }

        // Persist the updated state
        self.store.update_state(&manifest).await?;

        // Handle requeue action
        match action {
            Action::None => {}
            Action::Requeue(delay) => {
                self.requeue_with_delay(manifest_name, delay).await;
            }
        }

        Ok(())
    }

    /// Requeue a manifest with a delay.
    async fn requeue_with_delay(&self, manifest_name: &ManifestName, delay: std::time::Duration) {
        let scheduled_at = Timestamp::now()
            .checked_add(jiff::Span::try_from(delay).unwrap_or(jiff::Span::new().seconds(5)))
            .unwrap_or_else(|_| Timestamp::now());

        tracing::trace!(%manifest_name, ?delay, %scheduled_at, "requeuing manifest");
        self.reconcile_queue
            .enqueue_at(manifest_name.clone(), scheduled_at)
            .await;
    }

    // Keep the old method name for backwards compatibility
    pub async fn reconcile(&self, cancellation_token: &CancellationToken) -> anyhow::Result<()> {
        self.run(cancellation_token).await
    }
}
