use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

pub(crate) enum PreparedCryptoWork {
    Open { work: CryptoWork, cipher: AeadKey },
    Seal { work: OutboundCryptoWork, cipher: AeadKey },
}

const DATAPLANE_AEAD_WORKER_FAIRNESS_PACKETS: usize = 8;
const DATAPLANE_AEAD_JOB_PACKETS: usize = 128;

impl PreparedCryptoWork {
    pub(crate) fn open(work: CryptoWork, cipher: AeadKey) -> Self {
        Self::Open { work, cipher }
    }

    pub(crate) fn seal(work: OutboundCryptoWork, cipher: AeadKey) -> Self {
        Self::Seal { work, cipher }
    }

    fn lane(&self) -> Lane {
        match self {
            Self::Open { work, .. } => work.reservation.lane,
            Self::Seal { work, .. } => work.reservation.lane,
        }
    }

    fn reservation(&self) -> &OwnerReservation {
        match self {
            Self::Open { work, .. } => &work.reservation,
            Self::Seal { work, .. } => &work.reservation,
        }
    }

    fn cipher(&self) -> &AeadKey {
        match self {
            Self::Open { cipher, .. } | Self::Seal { cipher, .. } => cipher,
        }
    }

    fn is_open(&self) -> bool {
        matches!(self, Self::Open { .. })
    }

    fn is_open_fsp_session_payload(&self) -> bool {
        match self {
            Self::Open { work, .. } => work.is_open_fsp_session_payload(),
            Self::Seal { .. } => false,
        }
    }

    fn into_owner_item(self) -> (CryptoOwnerRunItem, AeadKey) {
        match self {
            Self::Open { work, cipher } => (CryptoOwnerRunItem::open(work), cipher),
            Self::Seal { work, cipher } => (CryptoOwnerRunItem::seal(work), cipher),
        }
    }
}

struct CryptoOwnerRunBuilder {
    cipher: Option<AeadKey>,
    run: Option<CryptoOwnerRun>,
}

impl CryptoOwnerRunBuilder {
    fn new() -> Self {
        Self {
            cipher: None,
            run: None,
        }
    }

    fn push(&mut self, pool: &mut DataplaneAeadWorkerPool, work: PreparedCryptoWork) {
        if !self.matches_run(&work)
            || self
                .run
                .as_ref()
                .is_some_and(|run| run.len() >= DATAPLANE_AEAD_JOB_PACKETS)
        {
            self.flush(pool);
        }
        let (work, cipher) = work.into_owner_item();
        match &mut self.run {
            Some(run) => run.push(work),
            None => {
                self.run = Some(CryptoOwnerRun::new(work, DATAPLANE_AEAD_JOB_PACKETS));
                self.cipher = Some(cipher);
            }
        }
    }

    fn flush(&mut self, pool: &mut DataplaneAeadWorkerPool) {
        let Some(run) = self.run.take() else {
            return;
        };
        let cipher = self
            .cipher
            .take()
            .expect("crypto run cipher exists when work is non-empty");
        pool.submit_run(run, cipher);
    }

    fn matches_run(&self, work: &PreparedCryptoWork) -> bool {
        let Some(run) = self.run.as_ref() else {
            return true;
        };
        let Some(current_cipher) = self.cipher.as_ref() else {
            return true;
        };
        Arc::ptr_eq(current_cipher, work.cipher())
            && run.matches(
                work.reservation(),
                work.is_open(),
                work.is_open_fsp_session_payload(),
            )
    }
}

#[derive(Clone, Debug)]
struct DataplaneAeadWorkerCounters {
    in_flight: Arc<AtomicUsize>,
    bulk_in_flight: Arc<AtomicUsize>,
}

impl DataplaneAeadWorkerCounters {
    fn new() -> Self {
        Self {
            in_flight: Arc::new(AtomicUsize::new(0)),
            bulk_in_flight: Arc::new(AtomicUsize::new(0)),
        }
    }

    fn add(&self, count: usize, bulk_count: usize) {
        self.in_flight.fetch_add(count, Relaxed);
        if bulk_count > 0 {
            self.bulk_in_flight.fetch_add(bulk_count, Relaxed);
        }
    }

    fn finish(&self, count: usize, bulk_count: usize) {
        self.in_flight.fetch_sub(count, Relaxed);
        if bulk_count > 0 {
            self.bulk_in_flight.fetch_sub(bulk_count, Relaxed);
        }
    }
}

#[derive(Debug)]
pub(crate) struct DataplaneAeadWorkerPool {
    completion_tx: tokio::sync::mpsc::Sender<CryptoCompletionBatch>,
    completion_rx: tokio::sync::mpsc::Receiver<CryptoCompletionBatch>,
    completion_notify: Arc<tokio::sync::Notify>,
    pending_completion_batch: Option<CryptoCompletionBatch>,
    counters: DataplaneAeadWorkerCounters,
    max_in_flight: usize,
    runtime: Option<tokio::runtime::Handle>,
    tasks: tokio::task::JoinSet<()>,
}

impl DataplaneAeadWorkerPool {
    pub(crate) fn new(max_in_flight: usize) -> Self {
        let max_in_flight = max_in_flight.max(1);
        let (completion_tx, completion_rx) = tokio::sync::mpsc::channel(max_in_flight);

        Self {
            completion_tx,
            completion_rx,
            completion_notify: Arc::new(tokio::sync::Notify::new()),
            pending_completion_batch: None,
            counters: DataplaneAeadWorkerCounters::new(),
            max_in_flight,
            runtime: tokio::runtime::Handle::try_current().ok(),
            tasks: tokio::task::JoinSet::new(),
        }
    }

    pub(crate) fn completion_notify(&self) -> Arc<tokio::sync::Notify> {
        Arc::clone(&self.completion_notify)
    }

    pub(crate) fn has_ready_completions(&self) -> bool {
        self.pending_completion_batch.is_some() || !self.completion_rx.is_empty()
    }

    pub(crate) fn drain_completion_batches_into_sink<S>(
        &mut self,
        limit: usize,
        sink: &mut S,
    ) -> usize
    where
        S: DataplaneCompletionSink,
    {
        self.drain_completion_batches_with(limit, |batch| {
            sink.push_completion_batch(batch);
        })
    }

    pub(crate) fn record_perf_depths(&self) {
        if !crate::perf_profile::enabled() {
            return;
        }
        crate::perf_profile::record_event_count(
            crate::perf_profile::Event::DataplaneAeadInFlight,
            self.counters.in_flight.load(Relaxed) as u64,
        );
        let pending_completion_depth = self
            .pending_completion_batch
            .as_ref()
            .map_or(0, CryptoCompletionBatch::len);
        let pending_completion_batches = self.pending_completion_batch.is_some() as usize;
        let rx_queued_messages = self.completion_rx.len();
        let completion_depth = pending_completion_depth.saturating_add(rx_queued_messages);
        crate::perf_profile::record_event_count(
            crate::perf_profile::Event::DataplaneAeadCompletionQueueDepth,
            completion_depth as u64,
        );
        crate::perf_profile::record_dataplane_aead_completion_backlog(
            rx_queued_messages,
            pending_completion_batches,
            pending_completion_depth,
        );
    }

    fn finish_drained_completions(&self, count: usize, bulk_count: usize) {
        self.counters.finish(count, bulk_count);
    }

    fn drain_completion_batch(
        &mut self,
        mut batch: CryptoCompletionBatch,
        limit: usize,
        push_batch: &mut impl FnMut(CryptoCompletionBatch),
    ) -> (usize, Option<CryptoCompletionBatch>) {
        crate::perf_profile::record_dataplane_aead_completion_batch(batch.len());
        let drained = batch.len().min(limit);
        if drained == 0 {
            return (0, Some(batch));
        }
        let pending = if drained < batch.len() {
            let pending = batch.split_off(drained);
            crate::perf_profile::record_dataplane_aead_completion_split(pending.len());
            Some(pending)
        } else {
            None
        };
        let bulk_count = if batch.lane() == Lane::Bulk { drained } else { 0 };
        self.finish_drained_completions(drained, bulk_count);
        push_batch(batch);
        (drained, pending)
    }

    fn drain_completion_batches_with(
        &mut self,
        limit: usize,
        mut push_batch: impl FnMut(CryptoCompletionBatch),
    ) -> usize {
        self.reap_finished_tasks();
        let mut drained = 0usize;
        while drained < limit {
            let Some(batch) = self
                .pending_completion_batch
                .take()
                .or_else(|| self.completion_rx.try_recv().ok())
            else {
                break;
            };
            let (got, pending) = self.drain_completion_batch(
                batch,
                limit.saturating_sub(drained),
                &mut push_batch,
            );
            drained = drained.saturating_add(got);
            if let Some(pending) = pending {
                self.pending_completion_batch = Some(pending);
                break;
            }
        }
        self.reap_finished_tasks();
        drained
    }

    fn available_capacity(&self) -> usize {
        self.max_in_flight.saturating_sub(
            self.counters.in_flight.load(Relaxed),
        )
    }

    fn available_capacity_for_lane(&self, lane: Lane) -> usize {
        let total_available = self.available_capacity();
        if lane == Lane::Priority {
            return total_available;
        }
        let bulk_limit =
            self.max_in_flight
                .saturating_sub(dataplane_aead_worker_priority_reserve(
                    self.max_in_flight,
                ));
        let bulk_in_flight = self
            .counters
            .bulk_in_flight
            .load(Relaxed);
        bulk_limit.saturating_sub(bulk_in_flight).min(total_available)
    }

    fn submit_run(&mut self, run: CryptoOwnerRun, cipher: AeadKey) {
        self.reap_finished_tasks();
        let chunk_len = run.len();
        let bulk_count = run.bulk_count();
        let queued_at = crate::perf_profile::stamp();
        self.counters.add(chunk_len, bulk_count);
        let completion_tx = self.completion_tx.clone();
        let completion_notify = Arc::clone(&self.completion_notify);
        let counters = self.counters.clone();
        let runtime = self
            .runtime
            .get_or_insert_with(tokio::runtime::Handle::current)
            .clone();
        self.tasks.spawn_on(
            async move {
                crate::perf_profile::record_since(
                    crate::perf_profile::Stage::DataplaneAeadWorkerQueueWait,
                    queued_at,
                );
                let completion = execute_crypto_owner_run(run, cipher);
                let completed_count = completion.len();
                let completed_bulk_count = if completion.lane() == Lane::Bulk {
                    completed_count
                } else {
                    0
                };
                crate::perf_profile::record_dataplane_aead_completion_send(
                    1,
                    1,
                    completed_count,
                );
                if completion_tx.send(completion).await.is_err() {
                    counters.finish(completed_count, completed_bulk_count);
                    return;
                }
                completion_notify.notify_one();
            },
            &runtime,
        );
        crate::perf_profile::record_dataplane_aead_prepared_job(chunk_len);
    }

    fn reap_finished_tasks(&mut self) {
        while let Some(result) = self.tasks.try_join_next() {
            result.expect("dataplane AEAD task failed");
        }
    }

    fn submit_prepared_chunk(
        &mut self,
        prepared: &mut Vec<PreparedCryptoWork>,
    ) {
        if prepared.is_empty() {
            return;
        }

        let mut runs = CryptoOwnerRunBuilder::new();
        for work in prepared.drain(..) {
            runs.push(self, work);
        }
        runs.flush(self);
    }
}

fn execute_crypto_owner_run(
    mut run: CryptoOwnerRun,
    cipher: AeadKey,
) -> CryptoCompletionBatch {
    let _open_timer = run.is_open().then(|| {
        crate::perf_profile::Timer::start(crate::perf_profile::Stage::DataplaneAeadOpen)
    });
    for item in &mut run.items {
        let state = std::mem::replace(
            &mut item.state,
            CryptoOwnerRunItemState::Completed(CryptoResult::Failed(CryptoFailureKind::Open)),
        );
        let result = match state {
            CryptoOwnerRunItemState::Open(packet) => {
                execute_open_crypto_work(packet, &item.reservation, &cipher)
            }
            CryptoOwnerRunItemState::Seal(packet) => {
                execute_seal_crypto_work(packet, &item.reservation, &cipher)
            }
            CryptoOwnerRunItemState::Completed(_) => {
                panic!("crypto owner run executed twice")
            }
        };
        item.state = CryptoOwnerRunItemState::Completed(result);
    }
    CryptoCompletionBatch::from_owner_run(run)
}

fn dataplane_aead_worker_priority_reserve(max_in_flight: usize) -> usize {
    max_in_flight
        .saturating_sub(DATAPLANE_AEAD_WORKER_FAIRNESS_PACKETS)
        .min(DATAPLANE_AEAD_WORKER_FAIRNESS_PACKETS)
}

impl std::fmt::Debug for PreparedCryptoWork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Open { work, .. } => f
                .debug_struct("PreparedCryptoWork::Open")
                .field("reservation", &work.reservation)
                .finish_non_exhaustive(),
            Self::Seal { work, .. } => f
                .debug_struct("PreparedCryptoWork::Seal")
                .field("reservation", &work.reservation)
                .finish_non_exhaustive(),
        }
    }
}

fn failed_crypto_completion(
    reservation: OwnerReservation,
    kind: CryptoFailureKind,
) -> CryptoCompletion {
    CryptoCompletion {
        reservation,
        result: CryptoResult::Failed(kind),
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum AeadHeader {
    Fmp([u8; FMP_ESTABLISHED_HEADER_SIZE]),
    Fsp([u8; FSP_HEADER_SIZE]),
}

impl AeadHeader {
    fn as_aad(&self) -> &[u8] {
        match self {
            Self::Fmp(header) => header,
            Self::Fsp(header) => header,
        }
    }
}

fn execute_open_crypto_work(
    mut packet: SocketPacket,
    reservation: &OwnerReservation,
    cipher: &LessSafeKey,
) -> CryptoResult {
    let parsed = match packet.owner.protocol {
        PacketProtocol::Fmp => FmpWireHeader::parse(packet.payload.as_slice()).map(|header| {
            (
                AeadHeader::Fmp(header.header_bytes()),
                header.ciphertext_offset(),
                header.counter(),
            )
        }),
        PacketProtocol::Fsp => FspWireHeader::parse(packet.payload.as_slice()).map(|header| {
            (
                AeadHeader::Fsp(header.header_bytes()),
                header.ciphertext_offset(),
                header.counter(),
            )
        }),
    };
    let Ok((header, ciphertext_offset, counter)) = parsed else {
        return CryptoResult::Failed(CryptoFailureKind::Open);
    };
    if counter != packet.counter {
        return CryptoResult::Failed(CryptoFailureKind::Open);
    }

    let target = packet.output;
    let source_wire_len = packet.payload.len();
    let plaintext_len = {
        let Some(ciphertext) = packet.payload.as_mut_slice().get_mut(ciphertext_offset..) else {
            return CryptoResult::Failed(CryptoFailureKind::Open);
        };
        let nonce = aead_nonce(reservation.counter);
        let Ok(plaintext) = cipher.open_in_place(nonce, Aad::from(header.as_aad()), ciphertext)
        else {
            return CryptoResult::Failed(CryptoFailureKind::Open);
        };
        plaintext.len()
    };
    packet.payload.truncate(ciphertext_offset + plaintext_len);
    CryptoResult::Opened(PacketOutput {
        owner: reservation.owner,
        counter: reservation.counter,
        ingress_seq: reservation.ingress_seq,
        lane: reservation.lane,
        target,
        source_path: reservation.source_path.clone(),
        previous_hop: reservation.previous_hop,
        ce_flag: reservation.ce_flag,
        path_mtu: reservation.path_mtu,
        source_peer: reservation.source_peer,
        path: reservation.output_path.clone(),
        activity_tick: reservation.activity_tick,
        fmp_timestamp_ms: reservation.fmp_timestamp_ms,
        source_wire_len: Some(source_wire_len),
        fsp_send_receipt: None,
        send_token: reservation.send_token,
        payload: packet.payload,
    })
}

fn execute_seal_crypto_work(
    mut packet: OutboundPacket,
    reservation: &OwnerReservation,
    cipher: &LessSafeKey,
) -> CryptoResult {
    let _timer = crate::perf_profile::Timer::start(crate::perf_profile::Stage::DataplaneAeadSeal);
    let inner_prefix = match packet.crypto_plaintext_prefix(
        reservation.fmp_timestamp_ms,
        reservation.fsp_timestamp_ms,
    ) {
        Ok(prefix) => prefix,
        Err(_) => return CryptoResult::Failed(CryptoFailureKind::Seal),
    };
    let Ok(payload_len) = u16::try_from(inner_prefix.len().saturating_add(packet.payload.len()))
    else {
        return CryptoResult::Failed(CryptoFailureKind::Seal);
    };
    let (header, coord_prefix, ciphertext_offset) = match (packet.owner.protocol, packet.wire) {
            (
                PacketProtocol::Fmp,
                OutboundWire::Fmp {
                    receiver_idx,
                    flags,
                },
            ) => (
                AeadHeader::Fmp(build_fmp_established_header(
                    receiver_idx,
                    reservation.counter,
                    flags,
                    payload_len,
                )),
                Vec::new(),
                FMP_ESTABLISHED_HEADER_SIZE,
            ),
            (PacketProtocol::Fsp, OutboundWire::Fsp { flags }) => {
                let coord_prefix = std::mem::take(&mut packet.fsp_cleartext_prefix);
                if validate_fsp_cleartext_prefix(flags, &coord_prefix).is_err() {
                    return CryptoResult::Failed(CryptoFailureKind::Seal);
                }
                let ciphertext_offset = FSP_HEADER_SIZE + coord_prefix.len();
                let Ok(header) = build_fsp_established_header(
                    reservation.counter,
                    flags,
                    payload_len,
                ) else {
                    return CryptoResult::Failed(CryptoFailureKind::Seal);
                };
                (
                    AeadHeader::Fsp(header),
                    coord_prefix,
                    ciphertext_offset,
                )
            }
            _ => return CryptoResult::Failed(CryptoFailureKind::Seal),
        };

    let aad = header.as_aad();
    let aad_len = aad.len();
    let prefix_len = aad
        .len()
        .saturating_add(coord_prefix.len())
        .saturating_add(inner_prefix.len());
    if packet.payload.try_prepend_slices(
        &[aad, coord_prefix.as_slice(), inner_prefix.as_slice()],
        AEAD_TAG_SIZE,
    ) {
        crate::perf_profile::record_event(crate::perf_profile::Event::DataplaneSealInPlace);
    } else {
        crate::perf_profile::record_event(crate::perf_profile::Event::DataplaneSealAllocated);
        let plaintext = std::mem::take(&mut packet.payload);
        let mut payload = Vec::with_capacity(
            prefix_len
                .saturating_add(plaintext.len())
                .saturating_add(AEAD_TAG_SIZE),
        );
        payload.extend_from_slice(aad);
        payload.extend_from_slice(&coord_prefix);
        payload.extend_from_slice(&inner_prefix);
        payload.extend_from_slice(plaintext.as_slice());
        packet.payload = PacketBuffer::new(payload);
    }

    if aad_len > ciphertext_offset || ciphertext_offset > packet.payload.len() {
        return CryptoResult::Failed(CryptoFailureKind::Seal);
    }
    let nonce = aead_nonce(reservation.counter);
    let (prefix, plaintext) = packet
        .payload
        .as_mut_slice()
        .split_at_mut(ciphertext_offset);
    let Some(aad) = prefix.get(..aad_len) else {
        return CryptoResult::Failed(CryptoFailureKind::Seal);
    };
    let Ok(tag) = cipher.seal_in_place_separate_tag(nonce, Aad::from(aad), plaintext) else {
        return CryptoResult::Failed(CryptoFailureKind::Seal);
    };
    packet.payload.extend_from_slice(tag.as_ref());

    match packet.post_seal {
        OutboundPostSeal::Transport => CryptoResult::Sealed(PacketOutput {
            owner: reservation.owner,
            counter: reservation.counter,
            ingress_seq: reservation.ingress_seq,
            lane: reservation.lane,
            target: OutputTarget::Transport,
            source_path: reservation.source_path.clone(),
            previous_hop: reservation.previous_hop,
            ce_flag: reservation.ce_flag,
            path_mtu: reservation.path_mtu,
            source_peer: reservation.source_peer,
            path: reservation.output_path.clone(),
            activity_tick: reservation.activity_tick,
            fmp_timestamp_ms: reservation.fmp_timestamp_ms,
            source_wire_len: None,
            fsp_send_receipt: packet.fsp_send_receipt,
            send_token: reservation.send_token,
            payload: packet.payload,
        }),
        OutboundPostSeal::FmpWrap(route) => {
            let mut output = route
                .into_fmp_outbound(packet.class, packet.payload)
                .with_fsp_send_receipt(DataplaneFspSendReceipt {
                    owner: reservation.owner,
                    counter: reservation.counter,
                });
            if let Some(send_token) = packet.send_token {
                output = output.with_send_token(send_token);
            }
            if let Some(tick) = reservation.activity_tick {
                output = output.with_activity_tick(tick);
            }
            CryptoResult::Outbound(output)
        }
    }
}

fn validate_fsp_cleartext_prefix(flags: u8, prefix: &[u8]) -> Result<(), WireBuildError> {
    if flags & crate::node::session_wire::FSP_FLAG_CP == 0 {
        return if prefix.is_empty() {
            Ok(())
        } else {
            Err(WireBuildError::BadFspCoords)
        };
    }

    crate::node::session_wire::parse_encrypted_coords(prefix)
        .map(|_| ())
        .map_err(|_| WireBuildError::BadFspCoords)
}

fn aead_nonce(counter: u64) -> Nonce {
    let mut nonce_bytes = [0u8; 12];
    nonce_bytes[4..12].copy_from_slice(&counter.to_le_bytes());
    Nonce::assume_unique_for_key(nonce_bytes)
}
