package s0fs

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/sandbox0-ai/sandbox0/storage-proxy/pkg/objectstore"
)

const (
	manifestLatestKey = "manifests/latest.json"
	manifestDir       = "manifests"
	segmentDir        = "segments"
)

const (
	DefaultSegmentTargetSizeBytes uint64 = 4 << 20
	defaultSegmentCacheMaxBytes   int64  = 64 << 20
)

var segmentCacheMaxBytes int64 = defaultSegmentCacheMaxBytes

var ErrMaterializedManifestNotFound = errors.New("materialized manifest not found")

type Manifest struct {
	Version       int            `json:"version"`
	VolumeID      string         `json:"volume_id"`
	ManifestSeq   uint64         `json:"manifest_seq"`
	CheckpointSeq uint64         `json:"checkpoint_seq"`
	CreatedAt     time.Time      `json:"created_at"`
	State         *SnapshotState `json:"state"`
}

type Materializer struct {
	volumeID             string
	store                objectstore.Store
	objectStoreForVolume ObjectStoreResolver
	headStore            HeadStore
	encryption           *EncryptionConfig
	cache                *segmentCache
	segmentTargetSize    uint64
}

func NewMaterializer(volumeID string, store objectstore.Store, headStore HeadStore, resolvers ...ObjectStoreResolver) *Materializer {
	if store == nil {
		return nil
	}
	var resolver ObjectStoreResolver
	if len(resolvers) > 0 {
		resolver = resolvers[0]
	}
	return &Materializer{
		volumeID:             volumeID,
		store:                store,
		objectStoreForVolume: resolver,
		headStore:            headStore,
		cache:                newSegmentCache(segmentCacheMaxBytes),
		segmentTargetSize:    DefaultSegmentTargetSizeBytes,
	}
}

func (m *Materializer) SetEncryption(encryption *EncryptionConfig) {
	if m != nil {
		m.encryption = encryption
	}
}

func (m *Materializer) SetSegmentTargetSize(size uint64) {
	if m == nil {
		return
	}
	if size == 0 {
		size = DefaultSegmentTargetSizeBytes
	}
	m.segmentTargetSize = size
}

func (m *Materializer) Enabled() bool {
	return m != nil && m.store != nil
}

func (m *Materializer) Materialize(ctx context.Context, state *SnapshotState, expectedManifestSeq uint64) (*Manifest, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if !m.Enabled() {
		return nil, nil
	}
	if m.volumeID == "" {
		return nil, fmt.Errorf("%w: volume id is required", ErrInvalidInput)
	}
	if state == nil {
		return nil, fmt.Errorf("%w: snapshot state is required", ErrInvalidInput)
	}

	inline := cloneStateForMaterialization(state)
	normalizeState(inline)
	defaultSegmentVolumeIDs(inline, m.volumeID)
	if inline.NextSeq <= 1 {
		inline.NextSeq = 2
	}

	nextSeq := checkpointSequence(inline)
	if nextSeq == 0 {
		return nil, fmt.Errorf("%w: manifest sequence must be non-zero", ErrInvalidInput)
	}
	if nextSeq <= expectedManifestSeq {
		return nil, fmt.Errorf("%w: manifest seq %d must advance beyond %d", ErrCommittedHeadConflict, nextSeq, expectedManifestSeq)
	}

	manifestState, segments, err := buildMaterializedState(nextSeq, m.volumeID, inline, m.segmentTargetSize)
	if err != nil {
		return nil, err
	}
	manifest := &Manifest{
		Version:       1,
		VolumeID:      m.volumeID,
		ManifestSeq:   nextSeq,
		CheckpointSeq: checkpointSequence(inline),
		CreatedAt:     time.Now().UTC(),
		State:         manifestState,
	}

	for _, segment := range segments {
		storedSegmentPayload, segmentEncryption, err := m.encryption.encryptSegment(m.volumeID, segment)
		if err != nil {
			return nil, err
		}
		segment.Encryption = segmentEncryption
		if meta := manifest.State.Segments[segment.ID]; meta != nil {
			meta.Encryption = segmentEncryption
		}
		if err := m.putBytes(ctx, segment.Key, storedSegmentPayload); err != nil {
			return nil, err
		}
		m.cache.put(segmentCacheKey(segment.VolumeID, segment.Key), segment.Payload)
	}
	if err := m.putJSON(ctx, manifestKey(nextSeq), manifest); err != nil {
		return nil, err
	}
	if m.headStore != nil {
		head := &CommittedHead{
			VolumeID:      m.volumeID,
			ManifestSeq:   manifest.ManifestSeq,
			CheckpointSeq: manifest.CheckpointSeq,
			ManifestKey:   manifestKey(manifest.ManifestSeq),
			UpdatedAt:     manifest.CreatedAt,
		}
		if err := m.headStore.CompareAndSwapCommittedHead(ctx, m.volumeID, expectedManifestSeq, head); err != nil {
			return nil, err
		}
	} else if err := m.putJSON(ctx, manifestLatestKey, manifest); err != nil {
		return nil, err
	}

	return manifest, nil
}

func (m *Materializer) ReadSegmentRange(segment *Segment, off, limit int64) ([]byte, error) {
	if !m.Enabled() {
		return nil, fmt.Errorf("%w: materializer is not configured", ErrInvalidInput)
	}
	if segment == nil || segment.Key == "" {
		return nil, fmt.Errorf("%w: segment is required", ErrInvalidInput)
	}
	if limit == 0 {
		return nil, nil
	}
	if off < 0 {
		return nil, fmt.Errorf("%w: negative segment offset", ErrInvalidInput)
	}
	volumeID, store, err := m.storeForSegment(segment)
	if err != nil {
		return nil, err
	}
	cacheKey := segmentCacheKey(volumeID, segment.Key)

	if segment.Length > 0 && int64(segment.Length) <= segmentCacheMaxBytes {
		if payload, ok := m.cache.get(cacheKey); ok {
			return cloneByteRange(payload, off, limit), nil
		}
		if segment.Encryption != nil {
			payload, err := m.encryption.decryptSegmentRange(store, volumeID, segment, 0, int64(segment.Length))
			if err != nil {
				return nil, err
			}
			m.cache.put(cacheKey, payload)
			return cloneByteRange(payload, off, limit), nil
		}
		reader, err := store.Get(segment.Key, 0, int64(segment.Length))
		if err != nil {
			return nil, err
		}
		payload, readErr := io.ReadAll(reader)
		closeErr := reader.Close()
		if readErr != nil {
			return nil, readErr
		}
		if closeErr != nil {
			return nil, closeErr
		}
		m.cache.put(cacheKey, payload)
		return cloneByteRange(payload, off, limit), nil
	}

	if segment.Encryption != nil {
		return m.encryption.decryptSegmentRange(store, volumeID, segment, off, limit)
	}
	reader, err := store.Get(segment.Key, off, limit)
	if err != nil {
		return nil, err
	}
	payload, readErr := io.ReadAll(reader)
	closeErr := reader.Close()
	if readErr != nil {
		return nil, readErr
	}
	if closeErr != nil {
		return nil, closeErr
	}
	return payload, nil
}

func (m *Materializer) LoadLatestManifest(ctx context.Context) (*Manifest, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if !m.Enabled() {
		return nil, ErrMaterializedManifestNotFound
	}
	if m.headStore != nil {
		head, err := m.headStore.LoadCommittedHead(ctx, m.volumeID)
		switch {
		case err == nil:
			return m.loadManifestByKey(ctx, head.ManifestKey)
		case !errors.Is(err, ErrCommittedHeadNotFound):
			return nil, err
		}

		manifest, err := m.loadLegacyLatestManifest(ctx)
		if err != nil {
			return nil, err
		}
		head = &CommittedHead{
			VolumeID:      m.volumeID,
			ManifestSeq:   manifest.ManifestSeq,
			CheckpointSeq: manifest.CheckpointSeq,
			ManifestKey:   manifestKey(manifest.ManifestSeq),
			UpdatedAt:     manifest.CreatedAt,
		}
		if err := m.headStore.CompareAndSwapCommittedHead(ctx, m.volumeID, 0, head); err != nil {
			if errors.Is(err, ErrCommittedHeadConflict) {
				return m.LoadLatestManifest(ctx)
			}
			return nil, err
		}
		return manifest, nil
	}
	return m.loadLegacyLatestManifest(ctx)
}

func (m *Materializer) LoadLatestState(ctx context.Context) (*SnapshotState, *Manifest, error) {
	manifest, err := m.LoadLatestManifest(ctx)
	if err != nil {
		return nil, nil, err
	}
	state := cloneState(manifest.State)
	return state, manifest, nil
}

type materializedSegment struct {
	ID         string
	VolumeID   string
	Key        string
	Payload    []byte
	SHA256     string
	Encryption *SegmentEncryption
}

func buildMaterializedState(manifestSeq uint64, volumeID string, state *SnapshotState, targetSize uint64) (*SnapshotState, []*materializedSegment, error) {
	manifestState := &SnapshotState{
		NextSeq:   state.NextSeq,
		NextInode: state.NextInode,
		Nodes:     cloneNodeMap(state.Nodes),
		Children:  cloneChildrenMap(state.Children),
		Data:      make(map[uint64][]byte),
		ColdFiles: make(map[uint64][]FileExtent),
		Segments:  make(map[string]*Segment),
	}

	builder := newSegmentBuilder(manifestSeq, volumeID, targetSize)
	inodes := make([]uint64, 0, len(state.Data)+len(state.ColdFiles))
	seen := make(map[uint64]struct{}, len(state.Data)+len(state.ColdFiles))
	for inode := range state.Data {
		seen[inode] = struct{}{}
		inodes = append(inodes, inode)
	}
	for inode, extents := range state.ColdFiles {
		if len(extents) == 0 {
			continue
		}
		if _, ok := seen[inode]; !ok {
			seen[inode] = struct{}{}
			inodes = append(inodes, inode)
		}
	}
	sort.Slice(inodes, func(i, j int) bool { return inodes[i] < inodes[j] })

	for _, inode := range inodes {
		node := state.Nodes[inode]
		if node == nil || node.Type == TypeDirectory {
			continue
		}
		var fileExtents []FileExtent
		if payload, ok := state.Data[inode]; ok && len(payload) > 0 {
			extents, err := builder.append(payload)
			if err != nil {
				return nil, nil, err
			}
			fileExtents = append(fileExtents, extents...)
		} else {
			extents, err := materializeFileExtents(builder, state, manifestState, inode)
			if err != nil {
				return nil, nil, err
			}
			fileExtents = append(fileExtents, extents...)
		}
		fileExtents = coalesceExtents(fileExtents)
		if len(fileExtents) > 0 {
			manifestState.ColdFiles[inode] = fileExtents
		}
	}

	segments := builder.finish()
	for _, segment := range segments {
		manifestState.Segments[segment.ID] = &Segment{
			ID:       segment.ID,
			VolumeID: segment.VolumeID,
			Key:      segment.Key,
			Length:   uint64(len(segment.Payload)),
			SHA256:   segment.SHA256,
		}
	}
	return manifestState, segments, nil
}

func buildSegment(manifestSeq uint64, volumeID string, state *SnapshotState) (*materializedSegment, map[uint64][]FileExtent, error) {
	manifestState, segments, err := buildMaterializedState(manifestSeq, volumeID, state, 0)
	if err != nil {
		return nil, nil, err
	}
	segment := &materializedSegment{
		ID:       fmt.Sprintf("%020d-0", manifestSeq),
		VolumeID: volumeID,
		Key:      fmt.Sprintf("%s/%020d-0.bin", segmentDir, manifestSeq),
	}
	if len(segments) > 0 {
		segment = segments[0]
	}
	return segment, manifestState.ColdFiles, nil
}

func materializeFileExtents(builder *segmentBuilder, state, manifestState *SnapshotState, inode uint64) ([]FileExtent, error) {
	var out []FileExtent
	for _, extent := range state.ColdFiles[inode] {
		if extent.Length == 0 {
			continue
		}
		if extent.SegmentID == "" {
			out = append(out, extent)
			continue
		}
		existing := state.Segments[extent.SegmentID]
		if existing == nil {
			return nil, fmt.Errorf("%w: missing retained segment %s", ErrInvalidInput, extent.SegmentID)
		}
		if isInlineSegment(existing) {
			payload, err := inlineSegmentRange(existing, extent.Offset, extent.Length)
			if err != nil {
				return nil, err
			}
			extents, err := builder.append(payload)
			if err != nil {
				return nil, err
			}
			out = append(out, extents...)
			continue
		}
		out = append(out, extent)
		manifestState.Segments[extent.SegmentID] = cloneSegment(existing)
	}
	return out, nil
}

type segmentBuilder struct {
	manifestSeq uint64
	volumeID    string
	targetSize  uint64
	nextIndex   int
	current     *materializedSegment
	segments    []*materializedSegment
}

func newSegmentBuilder(manifestSeq uint64, volumeID string, targetSize uint64) *segmentBuilder {
	if targetSize == 0 {
		targetSize = DefaultSegmentTargetSizeBytes
	}
	return &segmentBuilder{
		manifestSeq: manifestSeq,
		volumeID:    volumeID,
		targetSize:  targetSize,
	}
}

func (b *segmentBuilder) append(payload []byte) ([]FileExtent, error) {
	if len(payload) == 0 {
		return nil, nil
	}
	var out []FileExtent
	for len(payload) > 0 {
		segment := b.ensureCurrent()
		space := b.targetSize - uint64(len(segment.Payload))
		if space == 0 {
			b.current = nil
			continue
		}
		n := len(payload)
		if uint64(n) > space {
			n = int(space)
		}
		offset := uint64(len(segment.Payload))
		segment.Payload = append(segment.Payload, payload[:n]...)
		out = append(out, FileExtent{
			SegmentID: segment.ID,
			Offset:    offset,
			Length:    uint64(n),
		})
		payload = payload[n:]
	}
	return out, nil
}

func (b *segmentBuilder) ensureCurrent() *materializedSegment {
	if b.current != nil && uint64(len(b.current.Payload)) < b.targetSize {
		return b.current
	}
	segmentID := fmt.Sprintf("%020d-%d", b.manifestSeq, b.nextIndex)
	b.nextIndex++
	b.current = &materializedSegment{
		ID:       segmentID,
		VolumeID: b.volumeID,
		Key:      fmt.Sprintf("%s/%s.bin", segmentDir, segmentID),
		Payload:  make([]byte, 0, int(b.targetSize)),
	}
	b.segments = append(b.segments, b.current)
	return b.current
}

func (b *segmentBuilder) finish() []*materializedSegment {
	for _, segment := range b.segments {
		sum := sha256.Sum256(segment.Payload)
		segment.SHA256 = hex.EncodeToString(sum[:])
	}
	return b.segments
}

func manifestKey(seq uint64) string {
	return fmt.Sprintf("%s/%020d.json", manifestDir, seq)
}

func checkpointSequence(state *SnapshotState) uint64 {
	if state == nil || state.NextSeq == 0 {
		return 0
	}
	return state.NextSeq - 1
}

func (m *Materializer) putJSON(ctx context.Context, key string, value any) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	payload, err := json.Marshal(value)
	if err != nil {
		return fmt.Errorf("marshal %s: %w", key, err)
	}
	payload, err = m.encryption.encryptBlob(payload, stateBlobAAD(m.volumeID, "object:"+key))
	if err != nil {
		return fmt.Errorf("encrypt %s: %w", key, err)
	}
	return m.putBytes(ctx, key, payload)
}

func (m *Materializer) getJSON(ctx context.Context, key string, value any) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	reader, err := m.store.Get(key, 0, -1)
	if err != nil {
		return fmt.Errorf("get %s: %w", key, err)
	}
	defer reader.Close()

	payload, err := io.ReadAll(reader)
	if err != nil {
		return fmt.Errorf("read %s: %w", key, err)
	}
	if plaintext, encrypted, err := m.encryption.decryptBlobIfEncrypted(payload, stateBlobAAD(m.volumeID, "object:"+key)); encrypted || err != nil {
		if err != nil {
			return err
		}
		payload = plaintext
	}
	if err := json.Unmarshal(payload, value); err != nil {
		return fmt.Errorf("decode %s: %w", key, err)
	}
	return nil
}

func (m *Materializer) loadManifestByKey(ctx context.Context, key string) (*Manifest, error) {
	if strings.TrimSpace(key) == "" {
		return nil, ErrMaterializedManifestNotFound
	}
	var manifest Manifest
	if err := m.getJSON(ctx, key, &manifest); err != nil {
		return nil, err
	}
	if manifest.State == nil {
		return nil, fmt.Errorf("materialized manifest %s has no state", key)
	}
	normalizeState(manifest.State)
	sourceVolumeID := manifest.VolumeID
	if strings.TrimSpace(sourceVolumeID) == "" {
		sourceVolumeID = m.volumeID
	}
	defaultSegmentVolumeIDs(manifest.State, sourceVolumeID)
	return &manifest, nil
}

func (m *Materializer) loadLegacyLatestManifest(ctx context.Context) (*Manifest, error) {
	if _, err := m.store.Head(manifestLatestKey); err != nil {
		return nil, ErrMaterializedManifestNotFound
	}
	return m.loadManifestByKey(ctx, manifestLatestKey)
}

func (m *Materializer) putBytes(ctx context.Context, key string, payload []byte) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := m.store.Put(key, bytes.NewReader(payload)); err != nil {
		return fmt.Errorf("put %s: %w", key, err)
	}
	return nil
}

func cloneByteRange(payload []byte, off, limit int64) []byte {
	if off > int64(len(payload)) {
		off = int64(len(payload))
	}
	end := int64(len(payload))
	if limit >= 0 && off+limit < end {
		end = off + limit
	}
	return append([]byte(nil), payload[off:end]...)
}

func (m *Materializer) storeForSegment(segment *Segment) (string, objectstore.Store, error) {
	if segment == nil {
		return "", nil, fmt.Errorf("%w: segment is required", ErrInvalidInput)
	}
	volumeID := strings.TrimSpace(segment.VolumeID)
	if volumeID == "" {
		volumeID = m.volumeID
	}
	if volumeID == "" || volumeID == m.volumeID {
		return m.volumeID, m.store, nil
	}
	if m.objectStoreForVolume == nil {
		return "", nil, fmt.Errorf("%w: segment %s belongs to volume %s but no object store resolver is configured", ErrInvalidInput, segment.ID, volumeID)
	}
	store, err := m.objectStoreForVolume(volumeID)
	if err != nil {
		return "", nil, err
	}
	if store == nil {
		return "", nil, fmt.Errorf("%w: object store resolver returned nil for volume %s", ErrInvalidInput, volumeID)
	}
	return volumeID, store, nil
}

func segmentCacheKey(volumeID, key string) string {
	return volumeID + "\x00" + key
}

func defaultSegmentVolumeIDs(state *SnapshotState, volumeID string) {
	if state == nil {
		return
	}
	for _, segment := range state.Segments {
		if segment != nil && strings.TrimSpace(segment.VolumeID) == "" {
			segment.VolumeID = volumeID
		}
	}
}

type segmentCache struct {
	mu       sync.Mutex
	maxBytes int64
	size     int64
	entries  map[string][]byte
	order    []string
}

func newSegmentCache(maxBytes int64) *segmentCache {
	if maxBytes <= 0 {
		maxBytes = defaultSegmentCacheMaxBytes
	}
	return &segmentCache{
		maxBytes: maxBytes,
		entries:  make(map[string][]byte),
	}
}

func (c *segmentCache) get(key string) ([]byte, bool) {
	if c == nil || key == "" {
		return nil, false
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	payload, ok := c.entries[key]
	return payload, ok
}

func (c *segmentCache) put(key string, payload []byte) {
	if c == nil || key == "" || int64(len(payload)) > c.maxBytes {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()

	if existing, ok := c.entries[key]; ok {
		c.size -= int64(len(existing))
	} else {
		c.order = append(c.order, key)
	}
	// The cache takes ownership of payload. ReadSegmentRange returns cloned
	// ranges to callers, so cached segment bytes remain immutable after put.
	c.entries[key] = payload
	c.size += int64(len(payload))

	for c.size > c.maxBytes && len(c.order) > 0 {
		evict := c.order[0]
		c.order = c.order[1:]
		if evicted, ok := c.entries[evict]; ok {
			delete(c.entries, evict)
			c.size -= int64(len(evicted))
		}
	}
}

func cloneNodeMap(nodes map[uint64]*Node) map[uint64]*Node {
	cloned := make(map[uint64]*Node, len(nodes))
	for inode, node := range nodes {
		cloned[inode] = cloneNode(node)
	}
	return cloned
}

func cloneChildrenMap(children map[uint64]map[string]uint64) map[uint64]map[string]uint64 {
	cloned := make(map[uint64]map[string]uint64, len(children))
	for inode, entries := range children {
		entryClone := make(map[string]uint64, len(entries))
		for name, childInode := range entries {
			entryClone[name] = childInode
		}
		cloned[inode] = entryClone
	}
	return cloned
}

func cloneSegment(segment *Segment) *Segment {
	if segment == nil {
		return nil
	}
	copy := *segment
	if segment.Encryption != nil {
		enc := *segment.Encryption
		enc.WrappedKey = append([]byte(nil), segment.Encryption.WrappedKey...)
		enc.NoncePrefix = append([]byte(nil), segment.Encryption.NoncePrefix...)
		copy.Encryption = &enc
	}
	copy.InlineData = append([]byte(nil), segment.InlineData...)
	return &copy
}

func cloneSegmentForMaterialization(segment *Segment) *Segment {
	if segment == nil {
		return nil
	}
	clone := *segment
	if segment.Encryption != nil {
		enc := *segment.Encryption
		enc.WrappedKey = append([]byte(nil), segment.Encryption.WrappedKey...)
		enc.NoncePrefix = append([]byte(nil), segment.Encryption.NoncePrefix...)
		clone.Encryption = &enc
	}
	if isInlineSegment(segment) {
		// Inline segment payloads are immutable after creation; materialization
		// snapshots can share the backing bytes while the engine lock is released.
		clone.InlineData = segment.InlineData
	} else {
		clone.InlineData = append([]byte(nil), segment.InlineData...)
	}
	return &clone
}
