use std::collections::BTreeSet;
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

use miette::{IntoDiagnostic, Result};
use sha2::{Digest, Sha256};
use tracing::{error, info, warn};

use super::indexer::stored_to_chunk_value;
use crate::db::keys;
use crate::db::stream_chunks::{Chunk, FANIN, did_hash};
use crate::state::AppState;
use crate::types::StoredEvent;

const INTERVAL: Duration = Duration::from_secs(30);
const ZSTD_LEVEL: i32 = 3;
const STAR_TYPE: &str = "systems.gaze.hydrant.streamChunk";
const STAR_MAGIC: [u8; 4] = [0x2a, 0x6c, 0x01, 0x00];
/// once this many firehose frames are read but not committed, ingest counts as busy
const BACKLOG: usize = 64;

pub(crate) fn chunk_worker(state: Arc<AppState>) {
    lower_priority();
    loop {
        if let Err(e) = build(&state) {
            error!(err = %e, "can't build stream chunks");
        }
        if state.stop.wait(INTERVAL) {
            break;
        }
    }
}

fn lower_priority() {
    // linux nice is per thread and this has to be the tid
    #[cfg(target_os = "linux")]
    if unsafe { libc::setpriority(libc::PRIO_PROCESS, libc::gettid() as libc::id_t, 19) } != 0 {
        warn!(err = %std::io::Error::last_os_error(), "can't lower the stream chunk worker's priority");
    }
}

/// under a firehose backlog a round seals only one span
pub(crate) fn build(state: &AppState) -> Result<()> {
    let Some(chunks) = state.db.stream.chunks() else {
        return Ok(());
    };
    chunks.settle()?;
    for dirty in chunks.dirty() {
        write(
            state,
            dirty.first,
            dirty.last,
            dirty.level,
            std::slice::from_ref(&dirty),
        )?;
    }
    merge_all(state)?;
    loop {
        let first = chunks.sealed_next();
        let last = first + chunks.span - 1;
        if state.db.stream.ids.head().is_none_or(|head| head < last) {
            break;
        }
        write(state, first, last, 0, &[])?;
        merge_all(state)?;
        if state.firehose_backlog() >= BACKLOG {
            break;
        }
    }
    chunks.collect_garbage()
}

/// panics when chunks are off
fn chunks(state: &AppState) -> &crate::db::stream_chunks::StreamChunks {
    state.db.stream.chunks().expect("stream chunks are on")
}

fn merge_all(state: &AppState) -> Result<()> {
    while let Some(run) = chunks(state).mergeable() {
        merge(state, run)?;
    }
    Ok(())
}

struct Writer {
    out: zstd::Encoder<'static, ChunkFile>,
    first: u64,
    last: u64,
    level: u32,
    events: u64,
    dids: BTreeSet<u64>,
}

struct ChunkFile {
    tmp: PathBuf,
    file: BufWriter<std::fs::File>,
    hash: Sha256,
    bytes: u64,
}

impl Write for ChunkFile {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let n = self.file.write(buf)?;
        self.hash.update(&buf[..n]);
        self.bytes += n as u64;
        Ok(n)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.file.flush()
    }
}

impl ChunkFile {
    fn create(state: &AppState, first: u64, last: u64) -> Result<Self> {
        let tmp = chunks(state).tmp_path();
        let mut this = Self {
            file: BufWriter::new(std::fs::File::create(&tmp).into_diagnostic()?),
            tmp,
            hash: Sha256::new(),
            bytes: 0,
        };
        this.write_all(&header_frame(state, first, last)?)
            .into_diagnostic()?;
        Ok(this)
    }

    fn finish(
        self,
        first: u64,
        last: u64,
        level: u32,
        events: u64,
        dids: BTreeSet<u64>,
    ) -> Result<(Chunk, PathBuf)> {
        self.file
            .into_inner()
            .map_err(|e| e.into_error())
            .into_diagnostic()?
            .sync_all()
            .into_diagnostic()?;
        let chunk = Chunk {
            first,
            last,
            level,
            events,
            bytes: self.bytes,
            sha256: crate::util::hex::encode(self.hash.finalize()),
            dids: dids.into_iter().collect(),
            dirty: false,
        };
        Ok((chunk, self.tmp))
    }
}

impl Writer {
    fn new(state: &AppState, first: u64, last: u64, level: u32) -> Result<Self> {
        let out = ChunkFile::create(state, first, last)?;
        Ok(Self {
            out: zstd::Encoder::new(out, ZSTD_LEVEL).into_diagnostic()?,
            first,
            last,
            level,
            events: 0,
            dids: BTreeSet::new(),
        })
    }

    fn push(&mut self, did: u64, id: u64, value: &[u8]) -> Result<()> {
        self.dids.insert(did);
        let mut lens = Vec::with_capacity(20);
        varint(&mut lens, 8);
        lens.extend_from_slice(&id.to_be_bytes());
        varint(&mut lens, value.len() as u64);
        self.out.write_all(&lens).into_diagnostic()?;
        self.out.write_all(value).into_diagnostic()?;
        self.events += 1;
        Ok(())
    }

    fn finish(self) -> Result<(Chunk, PathBuf)> {
        let out = self.out.finish().into_diagnostic()?;
        out.finish(self.first, self.last, self.level, self.events, self.dids)
    }
}

#[derive(serde::Serialize)]
struct Header<'a> {
    #[serde(rename = "$type")]
    kind: &'a str,
    stream: &'a str,
    first: u64,
    last: u64,
}

/// its own zstd frame because a merge cuts it off and writes a new one
fn header_frame(state: &AppState, first: u64, last: u64) -> Result<Vec<u8>> {
    let meta = serde_ipld_dagcbor::to_vec(&Header {
        kind: STAR_TYPE,
        stream: state.db.stream.id(),
        first,
        last,
    })
    .into_diagnostic()?;
    let mut header = STAR_MAGIC.to_vec();
    varint(&mut header, meta.len() as u64);
    header.extend_from_slice(&meta);
    zstd::bulk::compress(&header, ZSTD_LEVEL).into_diagnostic()
}

fn varint(out: &mut Vec<u8>, mut n: u64) {
    while n >= 0x80 {
        out.push(n as u8 | 0x80);
        n >>= 7;
    }
    out.push(n as u8);
}

/// built from the event rows, not copied from any chunk it replaces
fn write(state: &AppState, first: u64, last: u64, level: u32, replaces: &[Chunk]) -> Result<()> {
    let started = Instant::now();
    let chunks = chunks(state);
    let mut writer = Writer::new(state, first, last, level)?;
    let rows = state
        .db
        .stream
        .event_range(keys::event_key(first)..=keys::event_key(last));
    for guard in rows {
        let (key, value) = guard.into_inner().into_diagnostic()?;
        let id = u64::from_be_bytes(key.as_ref().try_into().into_diagnostic()?);
        // replay skips undecodable rows and chunks have to match it
        let Ok(stored) = rmp_serde::from_slice::<StoredEvent>(&value) else {
            continue;
        };
        let did = did_hash(stored.did.to_did().as_str());
        writer.push(did, id, &stored_to_chunk_value(state, id, stored))?;
    }
    let (chunk, tmp) = writer.finish()?;
    let (events, bytes) = (chunk.events, chunk.bytes);
    chunks.publish(replaces, chunk, &tmp)?;
    let rebuilt = !replaces.is_empty();
    info!(first, last, events, bytes, rebuilt, elapsed = ?started.elapsed(), "wrote a stream chunk");
    Ok(())
}

fn concat(state: &AppState, run: &[Chunk]) -> Result<(Chunk, PathBuf)> {
    let chunks = chunks(state);
    let (first, last) = (run[0].first, run[run.len() - 1].last);
    let mut out = ChunkFile::create(state, first, last)?;
    for chunk in run {
        let bytes = std::fs::read(chunks.path(&chunk.sha256)).into_diagnostic()?;
        let header = zstd::zstd_safe::find_frame_compressed_size(&bytes).map_err(|code| {
            miette::miette!(
                "a chunk without a header frame: {}",
                zstd::zstd_safe::get_error_name(code)
            )
        })?;
        out.write_all(&bytes[header..]).into_diagnostic()?;
    }
    let events = run.iter().map(|c| c.events).sum();
    let dids = run.iter().flat_map(|c| c.dids.iter().copied()).collect();
    out.finish(first, last, run[0].level + 1, events, dids)
}

fn merge(state: &AppState, run: Vec<Chunk>) -> Result<()> {
    let started = Instant::now();
    let chunks = chunks(state);
    let (chunk, tmp) = concat(state, &run)?;
    let (first, last, bytes) = (chunk.first, chunk.last, chunk.bytes);
    chunks.publish(&run, chunk, &tmp)?;
    info!(first, last, bytes, elapsed = ?started.elapsed(), "merged stream chunks");
    Ok(())
}

#[cfg(test)]
mod tests {
    use jacquard_common::types::did::Did;
    use jacquard_common::types::string::Tid;
    use tokio::sync::mpsc;

    use super::*;
    use crate::config::Config;
    use crate::db::Txn;
    use crate::db::types::{DbAction, DbRkey, DbTid};

    const ALICE: &str = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa";
    const BOB: &str = "did:plc:bbbbbbbbbbbbbbbbbbbbbbbb";
    const COLLECTION: &str = "app.bsky.feed.post";

    fn state(span: u64) -> (tempfile::TempDir, Arc<AppState>) {
        let tmp = tempfile::tempdir().unwrap();
        let state = open(&tmp, span);
        (tmp, state)
    }

    fn open(dir: &tempfile::TempDir, span: u64) -> Arc<AppState> {
        let config = Config {
            database_path: dir.path().to_path_buf(),
            stream_chunk_span: span,
            ..Default::default()
        };
        Arc::new(AppState::new(&config).unwrap())
    }

    fn frame(line: &[u8]) -> serde_json::Value {
        serde_json::from_slice(line).unwrap()
    }

    #[derive(serde::Deserialize)]
    struct ChunkFrame {
        cid: Option<jacquard_common::types::cid::IpldCid>,
        did: String,
        rev: String,
        live: bool,
        rkey: String,
        action: String,
        #[serde(default, with = "serde_bytes")]
        record: Option<Vec<u8>>,
        collection: String,
    }

    fn frame_json(id: u64, f: ChunkFrame) -> Vec<u8> {
        let json = |s: &str| serde_json::to_string(s).unwrap();
        let mut out = format!(
            r#"{{"id":{},"type":"record","record":{{"live":{},"did":{},"rev":{},"collection":{},"rkey":{},"action":{}"#,
            id,
            f.live,
            json(&f.did),
            json(&f.rev),
            json(&f.collection),
            json(&f.rkey),
            json(&f.action),
        )
        .into_bytes();
        let cid = f
            .cid
            .or_else(|| f.record.as_deref().map(crate::car::compute_cid));
        if let Some(body) = f.record {
            let mark = out.len();
            out.extend_from_slice(br#","record":"#);
            if crate::types::write_record_json(&body, &mut out).is_err() {
                out.truncate(mark);
            }
        }
        if let Some(cid) = cid {
            out.extend_from_slice(format!(r#","cid":"{cid}""#).as_bytes());
        }
        out.extend_from_slice(b"}}");
        out
    }

    #[derive(serde::Deserialize)]
    struct Header {
        #[serde(rename = "$type")]
        kind: String,
        stream: String,
        first: u64,
        last: u64,
    }

    fn varint(bytes: &[u8], at: &mut usize) -> usize {
        let (mut n, mut shift) = (0, 0);
        loop {
            let b = bytes[*at];
            *at += 1;
            n |= usize::from(b & 0x7f) << shift;
            if b & 0x80 == 0 {
                return n;
            }
            shift += 7;
        }
    }

    fn read_chunk(path: &std::path::Path) -> (Header, Vec<Vec<u8>>) {
        let bytes = zstd::decode_all(std::fs::File::open(path).unwrap()).unwrap();
        assert_eq!(bytes[..4], STAR_MAGIC);
        let mut at = 4;
        let len = varint(&bytes, &mut at);
        let header: Header = serde_ipld_dagcbor::from_slice(&bytes[at..at + len]).unwrap();
        at += len;
        let mut frames = Vec::new();
        while at < bytes.len() {
            let len = varint(&bytes, &mut at);
            let id = u64::from_be_bytes(bytes[at..at + len].try_into().unwrap());
            at += len;
            let len = varint(&bytes, &mut at);
            let value = serde_ipld_dagcbor::from_slice(&bytes[at..at + len]).unwrap();
            at += len;
            frames.push(frame_json(id, value));
        }
        (header, frames)
    }

    fn rkey(n: usize) -> DbRkey {
        DbRkey::Str(smol_str::format_smolstr!("r{n}"))
    }

    fn commit(state: &AppState, did: &'static str, ops: &[(usize, Option<usize>)]) {
        let did = Did::new_static(did).unwrap();
        let mut txn = Txn::new(&state.db);
        let mut records = txn
            .records(state, &DbTid::from(&Tid::now_0()), &did)
            .unwrap();
        for (n, version) in ops {
            match version {
                Some(v) => {
                    let body = serde_json::json!({ "$type": COLLECTION, "n": n, "v": v });
                    let block =
                        crate::car::CarBlock::from_body(serde_ipld_dagcbor::to_vec(&body).unwrap());
                    let action = if *v == 0 {
                        DbAction::Create
                    } else {
                        DbAction::Update
                    };
                    records
                        .put_record(COLLECTION, &rkey(*n), &block, action)
                        .unwrap();
                }
                None => records.delete_record(COLLECTION, &rkey(*n)).unwrap(),
            }
        }
        txn.commit().unwrap();
    }

    fn rounds(state: &AppState, rounds: std::ops::Range<usize>) {
        for round in rounds {
            for n in 0..5 {
                let op = match (round + n) % 4 {
                    3 => None,
                    v => Some(round * 4 + v),
                };
                commit(state, ALICE, &[(n, op)]);
                commit(state, BOB, &[(n, Some(round)), (n + 10, op)]);
            }
        }
    }

    /// what the real stream thread sends, which chunks have to match
    fn replay(state: &Arc<AppState>, cursor: u64) -> Vec<Vec<u8>> {
        let head = state.db.stream.event_head().unwrap();
        if head.is_none_or(|head| head < cursor) {
            return Vec::new();
        }
        let (tx, mut rx) = mpsc::channel(1 << 16);
        let thread_state = state.clone();
        let opts = crate::control::stream::StreamOptions::from_config(&Config::default());
        std::thread::spawn(move || {
            super::super::event_stream_thread::<super::super::Frame>(
                thread_state,
                tx,
                Some(cursor),
                opts,
            )
        });
        let mut frames = Vec::new();
        loop {
            let line = rx.blocking_recv().unwrap().unwrap().to_vec();
            let id = frame(&line)["id"].as_u64();
            frames.push(line);
            if id == head {
                return frames;
            }
        }
    }

    fn chunk_frames(state: &AppState) -> (Vec<Vec<u8>>, u64) {
        let chunks = chunks(state);
        let manifest = chunks.manifest();
        let mut frames = Vec::new();
        for c in &manifest.chunks {
            let path = chunks
                .servable(&c.sha256)
                .expect("a listed chunk is served");
            let file = std::fs::read(&path).unwrap();
            assert_eq!(crate::util::hex::encode(Sha256::digest(&file)), c.sha256);
            let (header, chunk) = read_chunk(&path);
            assert_eq!(
                (
                    header.kind.as_str(),
                    header.stream.as_str(),
                    header.first,
                    header.last
                ),
                (STAR_TYPE, state.db.stream.id(), c.first, c.last)
            );
            frames.extend(chunk);
        }
        (frames, manifest.through.map_or(0, |through| through + 1))
    }

    fn with_tail(state: &Arc<AppState>) -> Vec<Vec<u8>> {
        let (mut frames, next) = chunk_frames(state);
        frames.extend(replay(state, next));
        frames
    }

    #[test]
    fn chunks_and_the_tail_are_the_replay() {
        let (_tmp, state) = state(8);
        rounds(&state, 0..3);
        build(&state).unwrap();
        // later writes hit records already in sealed chunks
        rounds(&state, 3..6);
        build(&state).unwrap();
        let (frames, next) = chunk_frames(&state);
        assert!(next > 64 && !frames.is_empty(), "nothing sealed");
        assert_eq!(with_tail(&state), replay(&state, 0));
        assert!(chunks(&state).manifest().chunks.len() < 8);
    }

    #[test]
    fn the_worker_builds_until_hydrant_stops() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        let worker = {
            let state = state.clone();
            std::thread::spawn(move || chunk_worker(state))
        };
        while chunks(&state).manifest().chunks.is_empty() {
            std::thread::sleep(Duration::from_millis(10));
        }
        state.stop.trigger();
        worker.join().unwrap();
        assert_eq!(with_tail(&state), replay(&state, 0));
    }

    #[test]
    fn merges_stop_at_sixteen_spans() {
        let (_tmp, state) = state(1);
        rounds(&state, 0..6);
        build(&state).unwrap();
        let manifest = chunks(&state).manifest();
        let widest = manifest.chunks.iter().map(|c| c.last + 1 - c.first).max();
        assert_eq!(widest, Some(16));
        assert_eq!(with_tail(&state), replay(&state, 0));
    }

    #[test]
    fn a_restart_keeps_the_chunks_it_built() {
        let tmp = tempfile::tempdir().unwrap();
        let state = open(&tmp, 4);
        rounds(&state, 0..6);
        build(&state).unwrap();
        let before = serde_json::to_value(chunks(&state).manifest()).unwrap();
        assert!(!before["chunks"].as_array().unwrap().is_empty());
        drop(state);

        let state = open(&tmp, 4);
        let chunks = chunks(&state);
        assert_eq!(serde_json::to_value(chunks.manifest()).unwrap(), before);
        let through = before["through"].as_u64().unwrap();
        assert_eq!(chunks.sealed_next(), through + 1);
        rounds(&state, 6..8);
        build(&state).unwrap();
        assert_eq!(with_tail(&state), replay(&state, 0));
    }

    #[test]
    fn a_busy_ingest_gets_one_span_a_round() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        let firehose = url::Url::parse("wss://relay.invalid").unwrap();
        let tracker = futures::executor::block_on(state.firehose_tracker(&firehose));
        let tickets: Vec<_> = (0..BACKLOG as i64)
            .map(|seq| tracker.ticket(seq, 0))
            .collect();
        let chunks = chunks(&state);
        build(&state).unwrap();
        assert_eq!(chunks.sealed_next(), 4);
        build(&state).unwrap();
        assert_eq!(chunks.sealed_next(), 8);

        for ticket in tickets {
            ticket.done(|_| {});
        }
        build(&state).unwrap();
        let head = state.db.stream.ids.head().unwrap();
        assert!(
            chunks.sealed_next() + 4 > head,
            "an idle ingest left spans unsealed"
        );
    }

    fn bodies(frames: &[Vec<u8>]) -> usize {
        frames
            .iter()
            .filter(|line| frame(line)["record"].get("record").is_some())
            .count()
    }

    fn holds(state: &AppState, sha256: &str, did: &str) -> bool {
        read_chunk(&chunks(state).path(sha256))
            .1
            .iter()
            .any(|line| frame(line)["record"]["did"] == did)
    }

    #[test]
    fn a_redaction_reaches_every_chunk_holding_the_record() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        build(&state).unwrap();
        let before = chunks(&state).manifest();

        crate::db::redact_record_bodies(
            &state.db,
            &Did::new_static(BOB).unwrap(),
            COLLECTION,
            &rkey(1),
            crate::types::DeleteBodyTarget::All,
        )
        .unwrap();
        let chunks = chunks(&state);
        for chunk in &before.chunks {
            let served = chunks.servable(&chunk.sha256).is_some();
            assert_eq!(served, !holds(&state, &chunk.sha256, BOB));
        }
        let dirty = chunks.manifest();
        assert!(dirty.through < before.through, "no chunk went dirty");
        assert!(dirty.generation > before.generation);

        build(&state).unwrap();
        assert_eq!(with_tail(&state), replay(&state, 0));
        assert_eq!(chunks.manifest().through, before.through);
    }

    #[test]
    fn an_erase_reaches_every_chunk_holding_the_repo() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        build(&state).unwrap();
        let with_bodies = bodies(&chunk_frames(&state).0);

        crate::db::erase_repo(&state.db, &Did::new_static(BOB).unwrap()).unwrap();
        build(&state).unwrap();
        let replayed = replay(&state, 0);
        assert!(bodies(&replayed) < with_bodies, "the erase took no body");
        assert_eq!(with_tail(&state), replayed);
    }

    fn change(state: &AppState, did: &str) {
        let mut batch = state.db.inner.batch();
        chunks(state).mark(&mut batch, did);
        batch.commit().unwrap();
        chunks(state).touch(did);
    }

    /// crashes right after change, before the chunk list gets written
    fn crash_after(change: impl FnOnce(&AppState)) {
        let tmp = tempfile::tempdir().unwrap();
        let state = open(&tmp, 4);
        rounds(&state, 0..6);
        build(&state).unwrap();
        let list = tmp.path().join("stream-chunks/state.json");
        let unaware = std::fs::read(&list).unwrap();
        change(&state);
        state.db.persist().unwrap();
        drop(state);
        std::fs::write(&list, unaware).unwrap();

        let state = open(&tmp, 4);
        let (served, _) = chunk_frames(&state);
        assert_eq!(served[..], replay(&state, 0)[..served.len()]);
        build(&state).unwrap();
        assert_eq!(with_tail(&state), replay(&state, 0));
    }

    #[test]
    fn turning_chunks_off_forgets_them() {
        let tmp = tempfile::tempdir().unwrap();
        let state = open(&tmp, 4);
        rounds(&state, 0..2);
        build(&state).unwrap();
        drop(state);
        drop(open(&tmp, 0));
        assert!(!tmp.path().join("stream-chunks").exists());
    }

    #[test]
    fn a_crash_before_a_redaction_reaches_the_chunks_leaves_none_holding_the_body() {
        use crate::types::DeleteBodyTarget::{Head, History};
        for target in [Head, History] {
            crash_after(|state| {
                let bob = Did::new_static(BOB).unwrap();
                crate::db::redact_record_bodies(&state.db, &bob, COLLECTION, &rkey(1), target)
                    .unwrap();
            });
        }
    }

    #[test]
    fn a_crash_before_an_erase_reaches_the_chunks_leaves_none_holding_the_bodies() {
        crash_after(|state| {
            crate::db::erase_repo(&state.db, &Did::new_static(BOB).unwrap()).unwrap();
        });
    }

    #[test]
    fn a_touch_while_chunks_merge_serves_none_of_them() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        build(&state).unwrap();
        rounds(&state, 6..12);
        while chunks(&state).sealed_next() + 4 <= state.db.stream.ids.head().unwrap() {
            let first = chunks(&state).sealed_next();
            write(&state, first, first + 3, 0, &[]).unwrap();
        }
        let chunks = chunks(&state);
        let run = chunks.mergeable().expect("a run to merge");
        change(&state, BOB);
        let (chunk, tmp) = concat(&state, &run).unwrap();
        chunks.publish(&run, chunk, &tmp).unwrap();
        for replaced in &run {
            assert!(chunks.servable(&replaced.sha256).is_none());
        }
    }

    #[test]
    fn a_touch_while_a_chunk_builds_leaves_it_dirty() {
        let (_tmp, state) = state(4);
        rounds(&state, 0..6);
        let chunks = chunks(&state);
        let mut writer = Writer::new(&state, 0, 3, 0).unwrap();
        writer.push(did_hash(BOB), 0, &[0xa0]).unwrap();
        change(&state, BOB);
        let (chunk, tmp) = writer.finish().unwrap();
        chunks.publish(&[], chunk, &tmp).unwrap();
        assert!(chunks.manifest().chunks.is_empty());
    }
}