//! A single atomic checkpoint for the existing DS4 KV index, not a second bank.
//! Tensor payloads keep the reference codec's dtypes, shapes and raw bytes.
use super::request::LiveSession;
use super::session_codec::decode_payload;
use super::stream::{Stream, Streams};
use crate::engine::metal::checkpoint::{read_u64, write_u64};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;

// V1 may bank an emitted MTP primary without its trunk/history state (#106).
// Reject it before decoding tensors; the owner rebuilds the selected cache.
const MAGIC: &[u8; 8] = b"QWMTKV02";
const HEADER_BYTES: u64 = 16;
const MAX_MANIFEST_BYTES: u64 = 64 * 1024 * 1024;
const IO_CHUNK: usize = 8 * 1024 * 1024;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(super) struct Identity {
    pub(super) model: [u8; 32],
    pub(super) context: u32,
    pub(super) vocabulary: u32,
    pub(super) mtp: bool,
    #[serde(default)]
    pub(super) head_chunk: bool,
}

#[derive(Serialize, Deserialize)]
struct Manifest {
    identity: Identity,
    tag: [u8; 32],
    #[serde(default)]
    verify_gate: bool,
    #[serde(default)]
    sorted_gather_guard: bool,
    #[serde(default)]
    images: Vec<crate::engine::qwen_vision::Identity>,
    tokens: Vec<i32>,
    lazy_kv: bool,
    payload: serde_json::Value,
    // Offsets are absolute; blob names are codec-local, never filesystem paths.
    blobs: BTreeMap<String, (u64, u64)>,
}

fn check_abort(abort: &impl Fn() -> bool) -> Result<(), String> {
    if abort() {
        Err("Qwen checkpoint interrupted".into())
    } else {
        Ok(())
    }
}

fn validate_tokens(tokens: &[i32], identity: Identity) -> Result<(), String> {
    if identity.context == 0
        || identity.vocabulary == 0
        || tokens.is_empty()
        || tokens.len() > identity.context as usize
        || tokens
            .iter()
            .any(|&token| token < 0 || token as u32 >= identity.vocabulary)
    {
        return Err("Qwen checkpoint token prefix is invalid".into());
    }
    Ok(())
}

impl LiveSession {
    #[allow(clippy::too_many_arguments)]
    pub(super) fn save_checkpoint(
        &mut self,
        path: &Path,
        identity: Identity,
        tag: [u8; 32],
        streams: &Streams,
        stream: Stream,
        mut progress: impl FnMut(u64),
        abort: impl Fn() -> bool,
    ) -> Result<(), String> {
        check_abort(&abort)?;
        let saved = self
            .current
            .as_ref()
            .ok_or("Qwen has no committed checkpoint")?;
        validate_tokens(&saved.token_ids, identity)?;
        if !crate::engine::qwen_vision::compatible(&self.current_images) {
            return Err("Qwen vision cache policy differs; rebuild the cache".into());
        }
        if saved.mtp.is_some() != identity.mtp {
            return Err("Qwen checkpoint MTP configuration differs".into());
        }
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(|e| e.to_string())?;
        }
        // Same serialized-writer, staging-file and sync/rename contract as DS4.
        let temporary = path.with_extension("tmp");
        let result = (|| -> Result<(), String> {
            let mut file = File::create(&temporary).map_err(|e| e.to_string())?;
            file.write_all(MAGIC).map_err(|e| e.to_string())?;
            write_u64(&mut file, 0)?;
            progress(HEADER_BYTES);
            let mut blobs = BTreeMap::new();
            let payload = saved.encode_payload(
                256,
                streams,
                stream,
                |name, bytes| {
                    let offset = file.stream_position().map_err(|e| e.to_string())?;
                    for chunk in bytes.chunks(IO_CHUNK) {
                        check_abort(&abort)?;
                        file.write_all(chunk).map_err(|e| e.to_string())?;
                        progress(chunk.len() as u64);
                    }
                    if blobs.insert(name, (offset, bytes.len() as u64)).is_some() {
                        return Err("duplicate Qwen checkpoint blob".into());
                    }
                    Ok(())
                },
                &abort,
            )?;
            check_abort(&abort)?;
            let offset = file.stream_position().map_err(|e| e.to_string())?;
            let manifest = serde_json::to_vec(&Manifest {
                verify_gate: true,
                sorted_gather_guard: true,
                identity,
                tag,
                images: self.current_images.clone(),
                tokens: saved.token_ids.clone(),
                lazy_kv: saved.lazy_kv,
                payload,
                blobs,
            })
            .map_err(|e| e.to_string())?;
            if manifest.len() as u64 > MAX_MANIFEST_BYTES {
                return Err("Qwen checkpoint metadata exceeds limit".into());
            }
            file.write_all(&manifest).map_err(|e| e.to_string())?;
            progress(manifest.len() as u64);
            file.seek(SeekFrom::Start(8)).map_err(|e| e.to_string())?;
            write_u64(&mut file, offset)?;
            file.sync_all().map_err(|e| e.to_string())?;
            check_abort(&abort)?;
            fs::rename(&temporary, path).map_err(|e| e.to_string())
        })();
        if result.is_err() {
            let _ = fs::remove_file(&temporary);
        } else {
            self.tag = tag;
        }
        result
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn load_checkpoint(
        &mut self,
        path: &Path,
        identity: Identity,
        streams: &Streams,
        stream: Stream,
        mut progress: impl FnMut(u64),
        abort: impl Fn() -> bool,
    ) -> Result<bool, String> {
        check_abort(&abort)?;
        let mut file = match File::open(path) {
            Ok(file) => file,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
            Err(e) => return Err(e.to_string()),
        };
        let mut magic = [0; 8];
        file.read_exact(&mut magic).map_err(|e| e.to_string())?;
        if &magic != MAGIC {
            return Err("Qwen checkpoint format differs; rebuild the cache".into());
        }
        let offset = read_u64(&mut file)?;
        let length = file.metadata().map_err(|e| e.to_string())?.len();
        if offset < HEADER_BYTES || offset >= length || length - offset > MAX_MANIFEST_BYTES {
            return Err("Qwen checkpoint metadata range is invalid".into());
        }
        file.seek(SeekFrom::Start(offset))
            .map_err(|e| e.to_string())?;
        let mut raw = vec![0; (length - offset) as usize];
        file.read_exact(&mut raw).map_err(|e| e.to_string())?;
        let mut manifest: Manifest = serde_json::from_slice(&raw).map_err(|e| e.to_string())?;
        progress(HEADER_BYTES + raw.len() as u64);
        if (Identity {
            head_chunk: identity.head_chunk,
            ..manifest.identity
        }) != identity
        {
            return Err(
                "Qwen checkpoint does not match the current model/context/MTP policy".into(),
            );
        }
        validate_tokens(&manifest.tokens, identity)?;
        if !crate::engine::qwen_vision::compatible(&manifest.images) {
            return Err("Qwen vision cache policy differs; rebuild the cache".into());
        }
        if identity.mtp && !manifest.verify_gate {
            return Err(
                "Qwen checkpoint uses obsolete MTP attention-gate math; rebuild the cache".into(),
            );
        }
        // The route also changes short dense prefills, not just MTP verify.
        // A <=3-token AR prefix cannot have executed a >=3-row body, whereas
        // MTP can retain a short frontier from a wider verification window.
        if manifest.identity.head_chunk != identity.head_chunk
            && (identity.mtp || manifest.tokens.len() >= 4)
        {
            return Err("Qwen checkpoint SDPA head-chunk policy differs; rebuild the cache".into());
        }
        // Old files do not record prefill chunk sizes or producer hardware.
        // At ten routes/token this length can contain an affected chunk;
        // shorter sessions cannot. Reject before allocating/restoring state.
        if !manifest.sorted_gather_guard && manifest.tokens.len() >= 3277 {
            return Err(
                "Qwen checkpoint predates the sorted-gather row guard; rebuild the cache".into(),
            );
        }
        if !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp {
            return Err("Qwen checkpoint MTP payload differs".into());
        }
        // Validate all extents before allocating any GPU state. Require exact
        // coverage, so overlapping blobs and unindexed/trailing bytes fail closed.
        let mut ranges = manifest.blobs.values().copied().collect::<Vec<_>>();
        ranges.sort_unstable();
        let mut end = HEADER_BYTES;
        for (start, size) in ranges {
            if start != end {
                return Err("Qwen checkpoint blob range is invalid".into());
            }
            end = start
                .checked_add(size)
                .filter(|&n| n <= offset)
                .ok_or("Qwen checkpoint blob exceeds payload")?;
        }
        if end != offset {
            return Err("Qwen checkpoint payload coverage differs".into());
        }
        let saved = decode_payload(
            &manifest.payload,
            manifest.tokens,
            manifest.lazy_kv,
            true,
            streams,
            stream,
            |name| {
                check_abort(&abort)?;
                let (start, size) = manifest
                    .blobs
                    .remove(name)
                    .ok_or("missing or repeated Qwen checkpoint blob")?;
                file.seek(SeekFrom::Start(start))
                    .map_err(|e| e.to_string())?;
                let mut bytes = Vec::new();
                let size = usize::try_from(size).map_err(|_| "Qwen checkpoint blob too large")?;
                bytes.try_reserve_exact(size).map_err(|e| e.to_string())?;
                bytes.resize(size, 0);
                for chunk in bytes.chunks_mut(IO_CHUNK) {
                    check_abort(&abort)?;
                    file.read_exact(chunk).map_err(|e| e.to_string())?;
                    progress(chunk.len() as u64);
                }
                Ok(bytes)
            },
        )?;
        if !manifest.blobs.is_empty() {
            return Err("unreferenced Qwen checkpoint blobs".into());
        }
        check_abort(&abort)?;
        // Publish only after complete validation and materialization. Failure
        // leaves the live conversation and its prefix untouched.
        self.current = Some(saved);
        self.prefix = None;
        self.tag = manifest.tag;
        self.current_images = manifest.images;
        self.prefix_images.clear();
        Ok(true)
    }
}

#[test]
#[ignore = "model-free Metal checkpoint; external supervisor and DS4_QWEN38_CHECKPOINT_DIR required"]
fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
    use super::super::gpu::Context;
    use super::session_cache::{Boundary, SessionSnapshot, State};
    use super::stream::Device;
    use super::{
        allocator::Allocator,
        array::{Array, Dtype},
        configure_sources, scalar_buffer,
    };
    use std::cell::Cell;
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(Allocator::new().unwrap()));
    let stream = streams.default_stream(Device::Gpu).unwrap();
    let leaf = Array::new(
        &[1, 1, 4],
        Dtype::BF16,
        scalar_buffer(&[0, 0, 128, 63, 0, 64, 64, 64]).unwrap(),
    )
    .unwrap();
    let state = Some(State::Recurrent(vec![Some(leaf.clone()), None]));
    let mut live = LiveSession {
        current_images: Vec::new(),
        prefix_images: Vec::new(),
        current: Some(SessionSnapshot {
            token_ids: vec![1, 2, 3, 4],
            logits: leaf.clone(),
            hidden: Some(leaf.clone()),
            trunk: vec![state.clone()],
            mtp: Some(vec![None]),
            lazy_kv: true,
            boundaries: vec![Boundary {
                tokens: 2,
                state: vec![state],
                hidden: Some(leaf),
            }],
        }),
        ..Default::default()
    };
    let identity = Identity {
        model: [7; 32],
        context: 64,
        vocabulary: 256,
        mtp: true,
        head_chunk: super::policy::sdpa_head_chunk_enabled(|key| std::env::var(key).ok()),
    };
    let directory =
        std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_CHECKPOINT_DIR").unwrap());
    fs::create_dir_all(&directory).unwrap();
    let path = directory.join("roundtrip.bin");
    let fingerprint = |live: &LiveSession| {
        let mut blobs = BTreeMap::new();
        let payload = live
            .current
            .as_ref()
            .unwrap()
            .encode_payload(
                256,
                &streams,
                stream,
                |name, bytes| {
                    blobs.insert(name, bytes);
                    Ok(())
                },
                || false,
            )
            .unwrap();
        (payload, blobs)
    };
    let original = fingerprint(&live);
    let mut written = 0;
    live.save_checkpoint(
        &path,
        identity,
        [9; 32],
        &streams,
        stream,
        |n| written += n,
        || false,
    )
    .unwrap();
    let valid = fs::read(&path).unwrap();
    assert_eq!(written, valid.len() as u64);
    assert_eq!(live.tag, [9; 32]);
    live = LiveSession::default();
    let mut read = 0;
    assert!(
        live.load_checkpoint(&path, identity, &streams, stream, |n| read += n, || false)
            .unwrap()
    );
    assert_eq!(read, written);
    assert_eq!(fingerprint(&live), original);
    assert_eq!(live.tag, [9; 32]);
    assert_eq!(live.current.as_ref().unwrap().token_ids, [1, 2, 3, 4]);
    assert!(live.current.as_ref().unwrap().lazy_kv);
    assert!(
        !live
            .load_checkpoint(
                &directory.join("missing.bin"),
                identity,
                &streams,
                stream,
                |_| {},
                || false
            )
            .unwrap()
    );

    // Cancellation after IO has started must not replace a valid file/tag.
    let interrupted = Cell::new(false);
    assert!(
        live.save_checkpoint(
            &path,
            identity,
            [8; 32],
            &streams,
            stream,
            |_| interrupted.set(true),
            || interrupted.get()
        )
        .is_err()
    );
    assert_eq!(fs::read(&path).unwrap(), valid);
    assert!(!path.with_extension("tmp").exists());
    interrupted.set(false);
    assert!(
        live.load_checkpoint(
            &path,
            identity,
            &streams,
            stream,
            |_| interrupted.set(true),
            || interrupted.get()
        )
        .is_err()
    );
    for changed in [
        Identity {
            context: 65,
            ..identity
        },
        Identity {
            mtp: false,
            ..identity
        },
        Identity {
            model: [8; 32],
            ..identity
        },
    ] {
        assert!(
            live.load_checkpoint(&path, changed, &streams, stream, |_| {}, || false)
                .is_err()
        );
    }

    let offset = u64::from_le_bytes(valid[8..16].try_into().unwrap()) as usize;
    let broken = directory.join("broken.bin");
    let mutations: [fn(&mut Manifest); 5] = [
        |m| m.tokens[0] = -1,
        |m| m.tokens.resize(65, 1),
        |m| m.blobs.values_mut().next().unwrap().1 = u64::MAX,
        |m| m.blobs.values_mut().next().unwrap().0 = 0,
        |m| m.payload["logits"]["nbytes"] = serde_json::json!(99),
    ];
    for mutate in mutations {
        let mut manifest: Manifest = serde_json::from_slice(&valid[offset..]).unwrap();
        mutate(&mut manifest);
        let mut bytes = valid[..offset].to_vec();
        bytes.extend(serde_json::to_vec(&manifest).unwrap());
        fs::write(&broken, bytes).unwrap();
        assert!(
            live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
                .is_err()
        );
    }
    for bytes in [&b"DS4QWN01"[..], &valid[..offset], &valid[..7]] {
        fs::write(&broken, bytes).unwrap();
        assert!(
            live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
                .is_err()
        );
    }
    assert_eq!(live.tag, [9; 32]);
    assert_eq!(fingerprint(&live), original);
    // Files predating the verify-phase law lack this field. Only MTP state
    // changed numerically; old AR files must remain exactly restorable.
    let without_gate_policy = |bytes: &[u8]| {
        let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
        let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
        manifest.as_object_mut().unwrap().remove("verify_gate");
        let mut result = bytes[..offset].to_vec();
        result.extend(serde_json::to_vec(&manifest).unwrap());
        result
    };
    fs::write(&broken, without_gate_policy(&valid)).unwrap();
    assert!(
        live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
            .unwrap_err()
            .contains("obsolete MTP attention-gate")
    );
    assert_eq!(fingerprint(&live), original);
    live.current.as_mut().unwrap().mtp = None;
    let ar_identity = Identity {
        mtp: false,
        ..identity
    };
    let ar_original = fingerprint(&live);
    live.save_checkpoint(
        &path,
        ar_identity,
        [10; 32],
        &streams,
        stream,
        |_| {},
        || false,
    )
    .unwrap();
    fs::write(&broken, without_gate_policy(&fs::read(&path).unwrap())).unwrap();
    assert!(
        live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
            .unwrap()
    );
    assert_eq!(fingerprint(&live), ar_original);
    let bytes = fs::read(&path).unwrap();
    let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
    let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
    manifest
        .as_object_mut()
        .unwrap()
        .remove("sorted_gather_guard");
    let write_manifest = |manifest: &serde_json::Value| {
        let mut old = bytes[..offset].to_vec();
        old.extend(serde_json::to_vec(manifest).unwrap());
        fs::write(&broken, old).unwrap();
    };
    write_manifest(&manifest);
    assert!(
        live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
            .unwrap()
    );
    assert_eq!(fingerprint(&live), ar_original);
    manifest["tokens"] = serde_json::json!(vec![1; 3277]);
    manifest["identity"]["context"] = 4096.into();
    write_manifest(&manifest);
    let large_identity = Identity {
        context: 4096,
        ..ar_identity
    };
    assert!(
        live.load_checkpoint(&broken, large_identity, &streams, stream, |_| {}, || false)
            .unwrap_err()
            .contains("sorted-gather row guard")
    );
    assert_eq!(fingerprint(&live), ar_original);
    // Policy changes in either direction invalidate affected checkpoints;
    // missing metadata is the actual pre-port (unchunked) policy.
    for mtp in [false, true] {
        for count in [1, 3, 4] {
            for produced in [false, true] {
                let saved = live.current.as_mut().unwrap();
                saved.token_ids = vec![1; count];
                saved.mtp = mtp.then(|| vec![None]);
                saved.boundaries.clear();
                let producer = Identity {
                    mtp,
                    head_chunk: produced,
                    ..identity
                };
                live.save_checkpoint(
                    &path,
                    producer,
                    [113; 32],
                    &streams,
                    stream,
                    |_| {},
                    || false,
                )
                .unwrap();
                let original = fingerprint(&live);
                let bytes = fs::read(&path).unwrap();
                let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
                for legacy in [false, true] {
                    if legacy && produced {
                        continue;
                    }
                    let mut manifest: serde_json::Value =
                        serde_json::from_slice(&bytes[offset..]).unwrap();
                    if legacy {
                        manifest["identity"]
                            .as_object_mut()
                            .unwrap()
                            .remove("head_chunk");
                    }
                    let raw = serde_json::to_vec(&manifest).unwrap();
                    let mut file = bytes[..offset].to_vec();
                    file.extend(&raw);
                    fs::write(&broken, file).unwrap();
                    for consumed in [false, true] {
                        let consumer = Identity {
                            head_chunk: consumed,
                            ..producer
                        };
                        let mut read = 0;
                        let result = live.load_checkpoint(
                            &broken,
                            consumer,
                            &streams,
                            stream,
                            |n| read += n,
                            || false,
                        );
                        if produced != consumed && (mtp || count >= 4) {
                            assert!(
                                result
                                    .unwrap_err()
                                    .contains("SDPA head-chunk policy differs")
                            );
                            assert_eq!(
                                read,
                                HEADER_BYTES + raw.len() as u64,
                                "rejected before tensor reads"
                            );
                        } else {
                            assert!(result.unwrap());
                        }
                        assert_eq!(fingerprint(&live), original);
                        assert_eq!(live.current.as_ref().unwrap().token_ids, vec![1; count]);
                    }
                }
            }
        }
    }
    println!(
        "#113 SDPA policy: legacy/current, both directions, AR/MTP, 1/3/4-token frontiers; exact compatible restore and incompatible rejection before tensor reads"
    );
    println!(
        "Qwen sorted-gather checkpoint policy: old short state exact, possibly affected long state rejected before restore"
    );
    println!(
        "Qwen checkpoint: exact roundtrip, cancellation/corruption controls, old MTP rejected and old AR restored"
    );
}