use crate::State;
use anyhow::{ensure, Result};
use relay_timelock::{GeneratedPuzzle, RandomXMode, SignedManifest};
use sha2::{Digest, Sha256};
use std::sync::{Arc, atomic::AtomicU64};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use zeroize::Zeroizing;

pub struct Epoch {
    pub manifest: SignedManifest,
    pub key: Zeroizing<[u8;32]>,
    pub activated_at_ms: u64,
    pub expires_at_ms: u64,
    pub activated_monotonic: Instant,
    pub puzzle_artifact: String,
    pub evidence_artifact: String,
    pub sequence: AtomicU64,
}

impl Epoch {
    pub fn accepts(&self, now_ms: u64) -> bool {
        self.accepts_at(now_ms, Instant::now())
    }

    fn accepts_at(&self, now_ms: u64, monotonic: Instant) -> bool {
        now_ms >= self.activated_at_ms && now_ms < self.expires_at_ms
            && !self.monotonic_expired(monotonic)
    }

    fn monotonic_expired(&self, now: Instant) -> bool {
        now.saturating_duration_since(self.activated_monotonic)
            >= Duration::from_millis(self.expires_at_ms.saturating_sub(self.activated_at_ms))
    }
}

/// Independent of generation and NSM availability: the active state's strong
/// reference is released when its monotonic lifetime ends. Existing request
/// leases are bounded by their capture and publication deadlines.
pub async fn expiration_watchdog(state: Arc<State>) {
    loop {
        expire_active(&state.active, Instant::now()).await;
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
}

async fn expire_active(active: &tokio::sync::RwLock<Option<Arc<Epoch>>>, now: Instant) {
    let mut active = active.write().await;
    if active.as_ref().is_some_and(|epoch| epoch.monotonic_expired(now)) {
        *active = None;
    }
}

/// Publication is acknowledged before activation or delivery of the response.
/// An untrusted host can still lie about durability; independent replicas and
/// client-held receipts are required for the availability promise.
pub async fn publish(state: &State, suffix: &str, data: &[u8]) -> Result<String> {
    ensure!(data.len() <= 32*1024*1024, "artifact exceeds publication bound");
    let name = format!("{}.{}", hex::encode(Sha256::digest(data)), suffix);
    for attempt in 0..3 {
        let result: Result<()> = async {
            let mut stream = state.transport.connect_records().await?;
            tlproxy_common::framing::write_frame(&mut stream, &name, data).await?;
            let mut ack = [0;3];
            stream.read_exact(&mut ack).await?;
            ensure!(ack == *b"OK\n", "artifact persistence rejected");
            stream.shutdown().await?;
            Ok(())
        }.await;
        if result.is_ok() { return Ok(name) }
        if attempt < 2 { tokio::time::sleep(Duration::from_millis(100)).await; }
    }
    anyhow::bail!("artifact persistence unavailable")
}

async fn activate(state: &State, generated: GeneratedPuzzle) -> Result<u64> {
    // Nothing about this candidate is disclosed until its evidence is durable.
    // Retry temporary attestation/storage failures without throwing away work.
    let (evidence_artifact, activated_monotonic, activated_at_ms) = loop {
        let attempt = async {
            let document = state.attestation_for_epoch(&rand::random::<[u8;32]>()).await?;
            let evidence = publish(state, "attestation.json", &serde_json::to_vec(&document)?).await?;
            let anchor = Instant::now();
            Ok::<_, anyhow::Error>((evidence, anchor, state.attester.trusted_time_ms_uncached()?))
        };
        match tokio::time::timeout(Duration::from_secs(60), attempt).await {
            Ok(Ok(prepared)) => break prepared,
            _ => {
                crate::log("epoch evidence unavailable; retrying before puzzle disclosure");
                tokio::time::sleep(Duration::from_secs(2)).await;
            }
        }
    };
    // Anchor ONCE before first possible disclosure. Never refresh these clocks
    // when retrying: a lost ACK may mean the puzzle is already public.
    let lifetime_ms = state.config.epoch_seconds.checked_mul(1000)
        .ok_or_else(|| anyhow::anyhow!("epoch duration overflow"))?;
    let expires_at_ms = activated_at_ms.checked_add(lifetime_ms)
        .ok_or_else(|| anyhow::anyhow!("epoch time overflow"))?;
    let puzzle_data = serde_json::to_vec(&generated.manifest)?;
    let puzzle_artifact = format!("{}.puzzle.json", hex::encode(Sha256::digest(&puzzle_data)));
    let bundle_data = serde_json::to_vec(&serde_json::json!({
        "version": 2, "puzzle": &puzzle_artifact, "attestation": &evidence_artifact,
        "manifest_id": generated.manifest.id()?, "epoch": generated.manifest.puzzle.epoch,
    }))?;
    let active = Arc::new(Epoch { manifest: generated.manifest, key: generated.epoch_key,
        activated_at_ms, expires_at_ms, activated_monotonic, puzzle_artifact, evidence_artifact,
        sequence: AtomicU64::new(0) });
    loop {
        ensure!(!active.monotonic_expired(Instant::now()), "publication candidate expired");
        let remaining = Duration::from_millis(lifetime_ms)
            .saturating_sub(activated_monotonic.elapsed());
        let attempt = async {
            let now = state.attester.trusted_time_ms_uncached()?;
            ensure!(active.accepts(now), "publication candidate outside its lifetime");
            publish(state, "puzzle.json", &puzzle_data).await?;
            publish(state, "bundle.json", &bundle_data).await?;
            let now = state.attester.trusted_time_ms_uncached()?;
            ensure!(active.accepts(now), "published candidate expired before activation");
            Ok::<_, anyhow::Error>(())
        };
        match tokio::time::timeout(remaining.min(Duration::from_secs(60)), attempt).await {
            Ok(Ok(())) => {
                *state.active.write().await = Some(active);
                crate::log("published puzzle activated");
                return Ok(expires_at_ms);
            }
            _ => {
                crate::log("epoch puzzle or bundle publication unavailable; retrying same candidate");
                tokio::time::sleep(remaining.min(Duration::from_secs(2))).await;
            }
        }
    }
}

// Linux applies nice values per thread. RandomX's bounded scoped workers inherit
// this priority, while the async network/runtime threads keep their priority.
// Nitro CPUs remain dedicated to the enclave; this does not lend CPUs to the host.
fn lower_generation_priority() -> Result<()> {
    #[cfg(target_os = "linux")]
    {
        let result = unsafe { libc::setpriority(libc::PRIO_PROCESS, 0, 19) };
        ensure!(result == 0, "cannot lower puzzle generation priority");
        ensure!(unsafe { libc::getpriority(libc::PRIO_PROCESS, 0) } == 19,
            "puzzle generation priority did not take effect");
    }
    Ok(())
}

pub async fn run(state: Arc<State>) {
    let mut epoch = rand::random::<u64>() & (u64::MAX >> 1);
    let mut activated_once = false;
    // Keep the publication deadline even after the watchdog erases an expired
    // key. Monotonic early erasure must never permit early future publication
    // while trusted NSM time is unavailable or has not reached that deadline.
    let mut publication_not_before: Option<u64> = None;
    loop {
        if state.dev && activated_once {
            if let Ok(delay) = std::env::var("RELAY_DEV_NEXT_GENERATION_DELAY_MS") {
                let delay = match delay.parse::<u64>() { Ok(n) if n <= 60_000 => n, _ => return };
                tokio::time::sleep(Duration::from_millis(delay)).await;
            }
        }
        let iterations = state.config.iterations;
        let signer = state.signer.clone();
        let mode = if state.dev { RandomXMode::Light } else { RandomXMode::Full };
        let entropy_state = state.clone();
        let low_priority = !state.dev;
        let workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1).min(relay_timelock::SEGMENTS);
        state.production.begin(iterations * relay_timelock::SEGMENTS as u64, workers as u64);
        let generation = tokio::task::spawn_blocking(move || {
            if low_priority { lower_generation_priority()?; }
            #[cfg(target_os = "linux")]
            { relay_timelock::generate_with_process_progress(&std::env::current_exe()?, epoch, iterations, &signer, mode, workers,
                |bytes| entropy_state.attester.fill_random(bytes),
                |hashes, done| entropy_state.production.advance(hashes, done)) }
            #[cfg(not(target_os = "linux"))]
            { relay_timelock::generate_with_progress(epoch, iterations, &signer, mode, workers,
                |bytes| entropy_state.attester.fill_random(bytes),
                |hashes, done| entropy_state.production.advance(hashes, done)) }
        }).await;
        let generated = match generation {
            Ok(Ok(generated)) => generated,
            _ => { state.production.phase(tlproxy_common::production::Phase::Failed); crate::log("puzzle generation failed; requests will expire closed"); return }
        };
        state.production.phase(tlproxy_common::production::Phase::AwaitingPublication);
        // Completed future puzzle and its seeds never leave enclave memory
        // before the currently active epoch ends.
        loop {
            let Some(expiration) = publication_not_before else { break };
            let now = match state.attester.trusted_time_ms_uncached() {
                Ok(now) => now,
                Err(_) => { tokio::time::sleep(Duration::from_secs(1)).await; continue }
            };
            if now >= expiration { break }
            tokio::time::sleep(Duration::from_millis((expiration-now).min(1000))).await;
        }
        // Requests have a bounded lease on the old epoch so they can finish
        // sealing; no new requests can acquire it after its deadline.
        *state.active.write().await = None;
        state.production.phase(tlproxy_common::production::Phase::Publishing);
        match activate(&state, generated).await {
            Ok(expiration) => { activated_once = true; publication_not_before = Some(expiration); },
            Err(_) => {
                state.production.phase(tlproxy_common::production::Phase::Failed);
                crate::log("publication candidate expired or invalid; generating replacement");
                tokio::time::sleep(Duration::from_secs(2)).await;
            }
        }
        let Some(next_epoch) = epoch.checked_add(1) else { return };
        epoch = next_epoch;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(target_os = "linux")]
    #[test]
    fn generation_priority_is_inherited_by_worker_threads() {
        let (generator, worker) = std::thread::spawn(|| {
            lower_generation_priority().unwrap();
            let generator = unsafe { libc::getpriority(libc::PRIO_PROCESS, 0) };
            let worker = std::thread::spawn(|| unsafe {
                libc::getpriority(libc::PRIO_PROCESS, 0)
            }).join().unwrap();
            (generator, worker)
        }).join().unwrap();
        assert_eq!((generator, worker), (19, 19));
    }

    #[test]
    fn epoch_admission_rejects_expiry_and_backward_clock() {
        let manifest = SignedManifest {
            puzzle: relay_timelock::Manifest {
                version: 1, epoch: 1, randomx: relay_timelock::RANDOMX_VERSION.into(),
                randomx_commit: relay_timelock::RANDOMX_COMMIT.into(), randomx_algorithm: "v2".into(),
                dataset_key: "00".repeat(32), iterations: 8, segments: 7,
                seed_1: "00".repeat(32), wrapped_keys: vec!["00".repeat(60); 7],
                key_commitment: "00".repeat(32),
            }, service_public_key: "00".repeat(32), service_signature: "00".repeat(64),
        };
        let epoch = Epoch {manifest, key: Zeroizing::new([0;32]), activated_at_ms: 100,
            activated_monotonic: Instant::now(),
            expires_at_ms: 200, puzzle_artifact: String::new(), evidence_artifact: String::new(),
            sequence: AtomicU64::new(0)};
        assert!(!epoch.accepts(99)); assert!(epoch.accepts(100));
        assert!(epoch.accepts(199)); assert!(!epoch.accepts(200));
        assert!(!epoch.accepts(u64::MAX)); assert!(!epoch.accepts(0));
        assert!(!epoch.accepts_at(150, epoch.activated_monotonic + Duration::from_millis(100)));
    }

    #[tokio::test]
    async fn watchdog_releases_expired_key_without_generation_or_nsm() {
        let manifest: SignedManifest = serde_json::from_str(include_str!("../../../python/attested-relay/tests/fixtures/native-puzzle.json")).unwrap();
        let activated = Instant::now();
        let epoch = Arc::new(Epoch { manifest, key: Zeroizing::new([42;32]),
            activated_at_ms: 100, expires_at_ms: 200, activated_monotonic: activated,
            puzzle_artifact: String::new(), evidence_artifact: String::new(), sequence: AtomicU64::new(0) });
        let weak = Arc::downgrade(&epoch);
        let active = tokio::sync::RwLock::new(Some(epoch));
        expire_active(&active, activated + Duration::from_millis(99)).await;
        assert!(active.read().await.is_some());
        // No generator or NSM call is involved; a failed/hung generation cannot
        // retain the old secret through the global active-state reference.
        expire_active(&active, activated + Duration::from_millis(100)).await;
        assert!(active.read().await.is_none());
        assert!(weak.upgrade().is_none());
    }

    #[tokio::test]
    async fn expired_key_survives_only_until_existing_request_lease_is_released() {
        let manifest: SignedManifest = serde_json::from_str(include_str!("../../../python/attested-relay/tests/fixtures/native-puzzle.json")).unwrap();
        let activated = Instant::now();
        let epoch = Arc::new(Epoch { manifest, key: Zeroizing::new([42;32]),
            activated_at_ms: 100, expires_at_ms: 200, activated_monotonic: activated,
            puzzle_artifact: String::new(), evidence_artifact: String::new(), sequence: AtomicU64::new(0) });
        let request_lease = epoch.clone();
        let weak = Arc::downgrade(&epoch);
        let active = tokio::sync::RwLock::new(Some(epoch));
        expire_active(&active, activated + Duration::from_millis(100)).await;
        assert!(active.read().await.is_none());
        assert_eq!(Arc::strong_count(&request_lease), 1);
        drop(request_lease);
        assert!(weak.upgrade().is_none());
    }
}