use std::io::{Read, Seek, SeekFrom};

use hurray_io::file::{FileReaderCore, FileStep};

/// Answers every range the core asks for, and returns the step that ended the job.
fn drive<R: Read + Seek>(
    core: &mut FileReaderCore,
    source: &mut R,
) -> hurray_io::Result<FileStep> {
    loop {
        match core.step()? {
            FileStep::Range { offset, len } => {
                source.seek(SeekFrom::Start(offset))?;
                // `read`, not `read_exact`: a bulk helper that fails cannot say how many
                // bytes it moved, and that count is the read position.
                let n = source.read(&mut core.input_dst()[..len])?;
                core.report_filled(n)?;
            }
            other => return Ok(other),
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
use hurray_core::{
    BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
    TensorDescriptor, MIN_BUFFER_ALIGNMENT,
};
use hurray_io::file::{FileWriteOp, FileWriterCore};
use hurray_io::stream::WriteStep;
fn drain(mut op: FileWriteOp<'_>, sink: &mut Vec<u8>) -> hurray_io::Result<()> {
    while let WriteStep::Output { .. } = op.step()? {
        let src = op.output_src();
        sink.extend_from_slice(src);
        let n = src.len();
        op.report_written(n)?;
    }
    Ok(())
}
let handle = BufferHandle::new(
    256, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced,
)?;
let descriptor = TensorDescriptor::new(
    1, 0, ElementType::Float32, Shape::new(vec![8u64, 8])?, 0,
    LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
)?;
let data: Vec<u8> = (0..=255u8).collect();
let buffers: [&[u8]; 1] = [&data];
let mut bytes = Vec::new();
let mut writer = FileWriterCore::new();
drain(writer.header()?, &mut bytes)?;
drain(writer.tensor("embeddings", &descriptor, &buffers)?, &mut bytes)?;
drain(writer.finish(vec![])?, &mut bytes)?;
let mut source = std::io::Cursor::new(bytes);
let size = source.seek(SeekFrom::End(0))?;
let mut core = FileReaderCore::new(size);

// Opening is four ranges at most: header, trailer, index, KV section.
assert_eq!(drive(&mut core, &mut source)?, FileStep::Ready);
println!("{} tensor(s)", core.index().len());

// A descriptor on its own — one range, and not a byte of payload.
core.begin_descriptor("embeddings")?;
assert_eq!(drive(&mut core, &mut source)?, FileStep::Descriptor);
let descriptor = core.take_descriptor().expect("a descriptor is ready");
println!("{:?} {:?}", descriptor.element_type, descriptor.shape.dims());

// The whole tensor — the descriptor, then one range per buffer.
core.begin_tensor("embeddings")?;
assert_eq!(drive(&mut core, &mut source)?, FileStep::Tensor);
let tensor = core.take_tensor().expect("a tensor is ready");

// The buffer carries the alignment its handle declared, so an aligned SIMD load on it is
// sound. This is the guarantee, not a property of the allocator that happened to hold.
let declared = tensor.descriptor.buffers[0].alignment() as usize;
assert_eq!(tensor.buffers[0].as_ptr() as usize % declared, 0);
assert_eq!(tensor.buffers[0].as_ref(), &data[..]);
Ok(())
}

#include <fcntl.h>
#include <stdio.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include "hurray.h"

/* One absolute range. A positional read keeps no cursor, which is why re-issuing a
   failed request reads the same bytes rather than resuming somewhere else. */
static HurrayStatus file_read_at(void *ctx, uint64_t offset, uint8_t *dst, size_t len,
                                 size_t *out_filled) {
    ssize_t n = pread(*(int *)ctx, dst, len, (off_t)offset);
    if (n < 0) { *out_filled = 0; return HURRAY_ERR_TRANSPORT; }  /* moved nothing */
    *out_filled = (size_t)n;  /* short is normal; 0 means the file is truncated */
    return HURRAY_OK;
}

/* Asked once, at open: a container is located from its end. */
static HurrayStatus file_size(void *ctx, uint64_t *out_size) {
    struct stat st;
    if (fstat(*(int *)ctx, &st) != 0) { return HURRAY_ERR_TRANSPORT; }
    *out_size = (uint64_t)st.st_size;
    return HURRAY_OK;
}

int fd = open("model.hrry", O_RDONLY);
HurrayTransport transport = { .ctx = &fd, .read_at = file_read_at, .size = file_size };

HurrayFileReader *reader = NULL;
if (hurray_file_reader_new(&transport, &reader) != HURRAY_OK) return -1;

/* The names come from the footer index, so listing them reads no tensor data. */
uint64_t count = 0;
hurray_file_reader_tensor_count(reader, &count);
for (uint64_t i = 0; i < count; i++) {
    const uint8_t *name = NULL;
    size_t name_len = 0;
    hurray_file_reader_tensor_name(reader, i, &name, &name_len);  /* borrowed, no NUL */
    printf("%.*s\n", (int)name_len, (const char *)name);
}

/* One tensor by name; the bytes of every other tensor stay on disk. */
const char *want = "embeddings";
HurrayTensor *tensor = NULL;
HurrayStatus s = hurray_file_reader_read_tensor(
    reader, (const uint8_t *)want, strlen(want), &tensor);
if (s == HURRAY_ERR_NOT_FOUND) {
    /* The file is fine — that name is simply not in it, and the reader stays usable. */
} else if (s == HURRAY_OK) {
    HurrayBufferList *buffers = NULL;   /* borrowed — do not destroy */
    hurray_tensor_buffers(tensor, &buffers);
    /* ... hand the buffers to your kernels ... */
    hurray_tensor_destroy(&tensor);     /* one call frees descriptor and buffers */
}

hurray_file_reader_destroy(&reader);
close(fd);

import numpy as np
import hurray

tensor = hurray.from_numpy(np.arange(64, dtype=np.float32))
hurray.save("model.hrry", {"embeddings": tensor}, kv={"producer": "cookbook"})

# The synchronous Python surface drives the same core, with the runtime detached for the
# duration of the call.
loaded = hurray.load("model.hrry")
assert np.array_equal(np.from_dlpack(loaded["embeddings"]), np.arange(64, dtype=np.float32))
assert hurray.load_kv("model.hrry")["producer"] == "cookbook"

use hurray_io::file::{FileWriteOp, FileWriterCore};
use hurray_io::stream::WriteStep;

/// Writes every byte of one submitted item. The sink may take fewer bytes than offered.
fn drain<W: std::io::Write>(
    mut op: FileWriteOp<'_>,
    sink: &mut W,
) -> hurray_io::Result<()> {
    loop {
        match op.step()? {
            WriteStep::Output { .. } => match sink.write(op.output_src())? {
                // A zero-length write is never progress: a closed sink is its own signal.
                0 => return Err(op.report_sink_closed()),
                n => op.report_written(n)?,
            },
            WriteStep::ItemWritten => return Ok(()),
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
use hurray_core::{
    BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
    TensorDescriptor, MIN_BUFFER_ALIGNMENT,
};
use hurray_io::file::KvValue;

let handle = BufferHandle::new(
    256, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced,
)?;
let descriptor = TensorDescriptor::new(
    1, 0, ElementType::Float32, Shape::new(vec![8u64, 8])?, 0,
    LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
)?;
let payload: Vec<u8> = (0..=255u8).collect();
// Bound, not a temporary: the operation borrows the payloads until the item is written.
let buffers: [&[u8]; 1] = [&payload];

let mut file = Vec::new();
let mut core = FileWriterCore::new();
drain(core.header()?, &mut file)?;
drain(core.tensor("embeddings", &descriptor, &buffers)?, &mut file)?;
drain(
    core.finish(vec![("producer".into(), KvValue::String("cookbook".into()))])?,
    &mut file,
)?;

assert_eq!(&file[0..8], b"HRRYFILE");
assert_eq!(&file[file.len() - 4..], b"HRRY");
Ok(())
}

#include "hurray.h"

static HurrayStatus file_write(void *ctx, const uint8_t *src, size_t len,
                               size_t *out_written) {
    ssize_t n = write(*(int *)ctx, src, len);
    if (n < 0) { *out_written = 0; return HURRAY_ERR_TRANSPORT; }
    *out_written = (size_t)n;  /* never 0 on success: that means the sink is closed */
    return HURRAY_OK;
}

int fd = open("model.hrry", O_WRONLY | O_CREAT | O_TRUNC, 0644);
/* A writer needs no range reads: it emits one forward pass. */
HurrayTransport transport = { .ctx = &fd, .write = file_write };

HurrayFileWriter *writer = NULL;
/* NULL options: 4096-byte data alignment, index in write order. The 64-byte header
   goes out during this call. */
if (hurray_file_writer_new(&transport, NULL, &writer) != HURRAY_OK) return -1;

const char *name = "embeddings";
HurrayIoSlice payload = { .ptr = data, .len = data_len };
if (hurray_file_writer_write_tensor(writer, (const uint8_t *)name, strlen(name),
                                    descriptor, &payload, 1) != HURRAY_OK) {
    /* A failed file write ends the file: the writer is poisoned, and the recovery is
       to write the file again from the start, not to retry this tensor. */
    hurray_file_writer_destroy(&writer);
    return -1;
}

/* KV values are borrowed, and the same struct a reader hands back — so copying one
   file's metadata into another needs no transcription. */
const char *producer = "cookbook";
HurrayKvEntry kv[1] = {{
    .key = (const uint8_t *)"producer",
    .key_len = 8,
    .value = { .kind = HURRAY_KV_STRING,
               .data = (const uint8_t *)producer,
               .data_len = 8 },
}};

/* Until this returns, the bytes on disk are not a container: the index is written here. */
hurray_file_writer_finish(writer, kv, 1);
hurray_file_writer_destroy(&writer);
close(fd);

import numpy as np
import hurray

# `save` is the same forward pass: header, tensors, KV, index, trailer.
hurray.save(
    "two.hrry",
    {
        "embeddings": hurray.from_numpy(np.arange(64, dtype=np.float32)),
        "bias": hurray.from_numpy(np.zeros(16, dtype=np.float32)),
    },
    kv={"producer": "cookbook", "layers": 32},
)

names = list(hurray.load("two.hrry"))
assert names == ["embeddings", "bias"]

use hurray_io::file::{flatten_composite, FileCompositeNode, FileWriterCore};

fn run(
    head: hurray_core::TensorDescriptor,
    members: &[FileCompositeNode<'_>],
    core: &mut FileWriterCore,
    file: &mut Vec<u8>,
) -> hurray_io::Result<()> {
fn drain(_: hurray_io::file::FileWriteOp<'_>, _: &mut Vec<u8>) -> hurray_io::Result<()> { Ok(()) }
for node in flatten_composite("weight", &head, members)? {
    // head first, then "weight.left", "weight.right", … depth first
    drain(core.tensor(node.name, node.descriptor, node.buffers)?, file)?;
}
Ok(())
}

use hurray_io::file::{FileItem, FileReaderCore, FileStep};

fn run(core: &mut FileReaderCore, source: &mut std::io::Cursor<Vec<u8>>) -> hurray_io::Result<()> {
fn drive(_: &mut FileReaderCore, _: &mut std::io::Cursor<Vec<u8>>) -> hurray_io::Result<FileStep> {
    Ok(FileStep::Composite)
}
core.begin_composite("weight")?;
assert_eq!(drive(core, source)?, FileStep::Composite);
let composite = core.take_composite().expect("a composite is ready");

for member in &composite.members {
    match member {
        FileItem::Tensor(t) => println!("{} — {} buffer(s)", t.name, t.buffers.len()),
        FileItem::Composite(c) => println!("{} — nested, {} member(s)", c.name, c.members.len()),
    }
}
Ok(())
}

use std::io::Write;
use hurray_io::file::{FileWriteOp, FileWriterCore};
use hurray_io::stream::WriteStep;

fn drain<W: Write>(mut op: FileWriteOp<'_>, sink: &mut W) -> hurray_io::Result<()> {
    loop {
        match op.step()? {
            WriteStep::Output { .. } => match sink.write(op.output_src())? {
                0 => return Err(op.report_sink_closed()),
                n => op.report_written(n)?,
            },
            WriteStep::ItemWritten => return Ok(()),
        }
    }
}

/// A sink that fails its first write, then accepts everything.
struct FailsOnce {
    bytes: Vec<u8>,
    failed: bool,
}

impl Write for FailsOnce {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        if !self.failed {
            self.failed = true;
            return Err(std::io::Error::other("simulated"));
        }
        self.bytes.extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
use hurray_core::{
    BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
    TensorDescriptor, MIN_BUFFER_ALIGNMENT,
};
let handle = BufferHandle::new(
    256, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced,
)?;
let descriptor = TensorDescriptor::new(
    1, 0, ElementType::Float32, Shape::new(vec![8u64, 8])?, 0,
    LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
)?;
let payload = vec![0u8; 256];
let buffers: [&[u8]; 1] = [&payload];

let mut sink = FailsOnce { bytes: Vec::new(), failed: false };
let mut core = FileWriterCore::new();

// The header's first write is the one that fails, so nothing reached the sink.
assert!(drain(core.header()?, &mut sink).is_err());
assert!(!core.is_unusable());

// Nothing was emitted, so every name is still free and the file is completable.
drain(core.header()?, &mut sink)?;
drain(core.tensor("weight", &descriptor, &buffers)?, &mut sink)?;
drain(core.finish(Vec::new())?, &mut sink)?;

assert_eq!(&sink.bytes[0..8], b"HRRYFILE");
assert_eq!(&sink.bytes[sink.bytes.len() - 4..], b"HRRY");
Ok(())
}

/* Same split, through the vtable. A failure before the tensor's first byte returns the
   write callback's own status and the writer stays usable; one after it returns
   HURRAY_ERR_SINK_CLOSED and every later call fails. */
HurrayStatus st = hurray_file_writer_write_tensor(
    writer, "weight", 6, descriptor, payloads, 1);
if (st == HURRAY_ERR_SINK_CLOSED) {
    /* Torn: this file cannot be completed. Start a new one. */
} else if (st != HURRAY_OK) {
    /* Nothing was emitted — retry the same name. */
    st = hurray_file_writer_write_tensor(
        writer, "weight", 6, descriptor, payloads, 1);
}

import os

import hurray

weight = hurray.zeros((8, 8), dtype=hurray.float32)

# `hurray.save` is all-or-nothing: it owns the file it creates, so there is no partially
# written state for a caller to inspect or resume. If it raises, delete and write again.
try:
    hurray.save("weights.hrry", {"weight": weight})
except OSError:
    if os.path.exists("weights.hrry"):
        os.remove("weights.hrry")
    hurray.save("weights.hrry", {"weight": weight})

assert hurray.load("weights.hrry")["weight"].shape == (8, 8)

cargo run --example sans_io_file

cargo run --example file_vtable -p hurray-ffi