package blockvol

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync/atomic"
	"time"
)

// Snapshot artifact format constants.
const (
	SnapshotArtifactFormatV1 = 1
	ArtifactLayoutSingleFile = "single-file"
	ExportToolVersion        = "sw-block-cp11a4"
	exportChunkBlocks        = 256 // read 256 blocks per chunk (1MB at 4KB block size)
	exportTempSnapBase       = uint32(0xFFFF0000) // base for temp snapshot IDs
)

// exportTempSnapSeq generates unique temp snapshot IDs per export to avoid collisions.
var exportTempSnapSeq atomic.Uint32

// FlagImported is set in Superblock.Flags after a successful import.
// Used to detect non-empty volumes even when nextLSN == 1 (import bypasses WAL).
const FlagImported uint16 = 1 << 0

// SnapshotArtifactManifest describes a snapshot export artifact.
type SnapshotArtifactManifest struct {
	FormatVersion     int    `json:"format_version"`
	SourceVolume      string `json:"source_volume"`
	SourceSizeBytes   uint64 `json:"source_size_bytes"`
	SourceBlockSize   uint32 `json:"source_block_size"`
	StorageProfile    string `json:"storage_profile"`
	CreatedAt         string `json:"created_at"`
	ArtifactLayout    string `json:"artifact_layout"`
	DataObjectKey     string `json:"data_object_key"`
	DataSizeBytes     uint64 `json:"data_size_bytes"`
	SHA256            string `json:"sha256"`
	Compression       string `json:"compression"`
	ExportToolVersion string `json:"export_tool_version"`
	// BaseLSN is the exact snapshot boundary — the highest LSN whose effects
	// are included in the exported image. Added for V2 rebuild execution
	// so the receiver can verify the imported base equals the planned target.
	// Zero for manifests created before this field was added.
	BaseLSN uint64 `json:"base_lsn,omitempty"`
}

var (
	ErrUnsupportedArtifactVersion = errors.New("blockvol: unsupported artifact format version")
	ErrUnsupportedArtifactLayout  = errors.New("blockvol: unsupported artifact layout")
	ErrUnsupportedProfileExport   = errors.New("blockvol: only single profile supported for export/import")
	ErrManifestMissingField       = errors.New("blockvol: manifest missing required field")
	ErrChecksumMismatch           = errors.New("blockvol: data checksum mismatch")
	ErrImportSizeMismatch         = errors.New("blockvol: target volume size does not match manifest")
	ErrImportBlockSizeMismatch    = errors.New("blockvol: target block size does not match manifest")
	ErrImportTargetNotEmpty       = errors.New("blockvol: target volume is not empty (use AllowOverwrite)")
	ErrImportDataShort            = errors.New("blockvol: import data shorter than manifest declares")
	ErrImportActiveSnapshots      = errors.New("blockvol: cannot import with active snapshots (extent would be overwritten)")
)

// MarshalManifest encodes a manifest as indented JSON.
func MarshalManifest(m *SnapshotArtifactManifest) ([]byte, error) {
	return json.MarshalIndent(m, "", "  ")
}

// UnmarshalManifest decodes and validates a manifest from JSON.
func UnmarshalManifest(data []byte) (*SnapshotArtifactManifest, error) {
	var m SnapshotArtifactManifest
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, fmt.Errorf("blockvol: unmarshal manifest: %w", err)
	}
	if err := ValidateManifest(&m); err != nil {
		return nil, err
	}
	return &m, nil
}

// ValidateManifest checks that all required fields are present and supported.
func ValidateManifest(m *SnapshotArtifactManifest) error {
	if m.FormatVersion != SnapshotArtifactFormatV1 {
		return fmt.Errorf("%w: %d", ErrUnsupportedArtifactVersion, m.FormatVersion)
	}
	if m.ArtifactLayout != ArtifactLayoutSingleFile {
		return fmt.Errorf("%w: %s", ErrUnsupportedArtifactLayout, m.ArtifactLayout)
	}
	if m.StorageProfile != "single" {
		return fmt.Errorf("%w: %s", ErrUnsupportedProfileExport, m.StorageProfile)
	}
	if m.SourceSizeBytes == 0 {
		return fmt.Errorf("%w: source_size_bytes", ErrManifestMissingField)
	}
	if m.SourceBlockSize == 0 {
		return fmt.Errorf("%w: source_block_size", ErrManifestMissingField)
	}
	if m.SHA256 == "" {
		return fmt.Errorf("%w: sha256", ErrManifestMissingField)
	}
	if m.DataSizeBytes == 0 {
		return fmt.Errorf("%w: data_size_bytes", ErrManifestMissingField)
	}
	return nil
}

// ExportOptions configures a snapshot export.
type ExportOptions struct {
	DataObjectKey string // S3 object key for the data object (stored in manifest)
	SnapshotID    uint32 // if > 0, export from existing snapshot; if 0, create+delete temp
}

// ExportSnapshot exports a crash-consistent snapshot of the volume to w.
// If opts.SnapshotID is 0, a temporary snapshot is created and deleted after export.
// If opts.SnapshotID is > 0, the existing snapshot is used (not deleted).
// The full logical volume image is streamed to w with SHA-256 computed inline.
func (v *BlockVol) ExportSnapshot(ctx context.Context, w io.Writer, opts ExportOptions) (*SnapshotArtifactManifest, error) {
	if err := v.beginOp(); err != nil {
		return nil, err
	}
	defer v.endOp()

	if v.Profile() != ProfileSingle {
		return nil, ErrUnsupportedProfileExport
	}

	info := v.Info()
	snapID := opts.SnapshotID
	deleteSnap := false

	if snapID == 0 {
		// Generate unique temp snapshot ID to avoid collisions with concurrent exports
		// and user-created snapshots.
		snapID = exportTempSnapBase + exportTempSnapSeq.Add(1)
		if err := v.CreateSnapshot(snapID); err != nil {
			return nil, fmt.Errorf("blockvol: export create temp snapshot: %w", err)
		}
		deleteSnap = true
		defer func() {
			if deleteSnap {
				v.DeleteSnapshot(snapID)
			}
		}()
	} else {
		// Verify snapshot exists.
		v.snapMu.RLock()
		_, ok := v.snapshots[snapID]
		v.snapMu.RUnlock()
		if !ok {
			return nil, ErrSnapshotNotFound
		}
	}

	// Read snapshot BaseLSN for the manifest.
	var baseLSN uint64
	v.snapMu.RLock()
	if snap, ok := v.snapshots[snapID]; ok {
		baseLSN = snap.header.BaseLSN
	}
	v.snapMu.RUnlock()

	h := sha256.New()
	mw := io.MultiWriter(w, h)

	totalBlocks := info.VolumeSize / uint64(info.BlockSize)
	chunkBlocks := uint64(exportChunkBlocks)
	var totalWritten uint64

	for startBlock := uint64(0); startBlock < totalBlocks; startBlock += chunkBlocks {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		remaining := totalBlocks - startBlock
		n := chunkBlocks
		if n > remaining {
			n = remaining
		}
		readBytes := uint32(n * uint64(info.BlockSize))

		data, err := v.ReadSnapshot(snapID, startBlock, readBytes)
		if err != nil {
			return nil, fmt.Errorf("blockvol: export read block %d: %w", startBlock, err)
		}

		if _, err := mw.Write(data); err != nil {
			return nil, fmt.Errorf("blockvol: export write: %w", err)
		}
		totalWritten += uint64(len(data))
	}

	manifest := &SnapshotArtifactManifest{
		FormatVersion:     SnapshotArtifactFormatV1,
		SourceVolume:      v.Path(),
		SourceSizeBytes:   info.VolumeSize,
		SourceBlockSize:   info.BlockSize,
		StorageProfile:    "single",
		CreatedAt:         time.Now().UTC().Format(time.RFC3339),
		ArtifactLayout:    ArtifactLayoutSingleFile,
		DataObjectKey:     opts.DataObjectKey,
		DataSizeBytes:     totalWritten,
		SHA256:            hex.EncodeToString(h.Sum(nil)),
		Compression:       "none",
		ExportToolVersion: ExportToolVersion,
		BaseLSN:           baseLSN,
	}

	return manifest, nil
}

// ImportOptions configures a snapshot import.
type ImportOptions struct {
	AllowOverwrite bool // if true, allow import into a non-empty volume
}

// ImportSnapshot imports a snapshot artifact into this volume.
// The volume must match the manifest's size and block size.
// Data is read from r and written directly to the extent region, bypassing the WAL.
// SHA-256 is verified against the manifest after the full read.
//
// IMPORTANT: Import is a destructive, non-atomic operation. If the import fails
// mid-stream (reader error, context cancellation), the extent region will contain
// a mix of old and new data. The volume remains operational but its data is
// inconsistent. On failure, the volume should be discarded or re-imported with
// AllowOverwrite. A future enhancement may add pre-import snapshot for rollback.
//
// Active snapshots are rejected because import overwrites the extent region that
// non-CoW'd snapshot blocks read from, which would silently corrupt snapshot reads.
func (v *BlockVol) ImportSnapshot(ctx context.Context, manifest *SnapshotArtifactManifest, r io.Reader, opts ImportOptions) error {
	if err := v.beginOp(); err != nil {
		return err
	}
	defer v.endOp()

	// Exclusive lock: drains all in-flight I/O before modifying extent/WAL/dirtyMap.
	v.ioMu.Lock()
	defer v.ioMu.Unlock()

	if err := ValidateManifest(manifest); err != nil {
		return err
	}

	info := v.Info()
	if info.VolumeSize != manifest.SourceSizeBytes {
		return fmt.Errorf("%w: target=%d manifest=%d", ErrImportSizeMismatch, info.VolumeSize, manifest.SourceSizeBytes)
	}
	if info.BlockSize != manifest.SourceBlockSize {
		return fmt.Errorf("%w: target=%d manifest=%d", ErrImportBlockSizeMismatch, info.BlockSize, manifest.SourceBlockSize)
	}

	// Reject import if active snapshots exist. Import overwrites the extent
	// region directly; non-CoW'd snapshot blocks read from extent and would
	// return import data instead of snapshot-time data (BUG-CP11A4-1).
	v.snapMu.RLock()
	snapCount := len(v.snapshots)
	v.snapMu.RUnlock()
	if snapCount > 0 {
		return ErrImportActiveSnapshots
	}

	// Empty check: nextLSN > 1 means WAL writes occurred; FlagImported means a
	// previous import wrote directly to extent (bypassing WAL, so nextLSN stays 1).
	if !opts.AllowOverwrite && (v.nextLSN.Load() > 1 || v.super.Flags&FlagImported != 0) {
		return ErrImportTargetNotEmpty
	}

	// Pause flusher — we write directly to extent.
	if err := v.flusher.PauseAndFlush(); err != nil {
		v.flusher.Resume()
		return fmt.Errorf("blockvol: import flush: %w", err)
	}
	defer v.flusher.Resume()

	// Stream data to extent, computing SHA-256.
	h := sha256.New()
	tr := io.TeeReader(r, h)

	extentStart := v.super.WALOffset + v.super.WALSize
	chunkSize := uint64(exportChunkBlocks) * uint64(info.BlockSize)
	buf := make([]byte, chunkSize)
	var totalRead uint64

	for totalRead < manifest.DataSizeBytes {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		remaining := manifest.DataSizeBytes - totalRead
		readSize := chunkSize
		if readSize > remaining {
			readSize = remaining
		}

		n, err := io.ReadFull(tr, buf[:readSize])
		if n > 0 {
			writeOff := int64(extentStart) + int64(totalRead)
			if _, werr := v.fd.WriteAt(buf[:n], writeOff); werr != nil {
				return fmt.Errorf("blockvol: import write at offset %d: %w", writeOff, werr)
			}
			totalRead += uint64(n)
		}
		if err != nil {
			if err == io.EOF || err == io.ErrUnexpectedEOF {
				break
			}
			return fmt.Errorf("blockvol: import read: %w", err)
		}
	}

	if totalRead != manifest.DataSizeBytes {
		return fmt.Errorf("%w: read %d bytes, manifest declares %d", ErrImportDataShort, totalRead, manifest.DataSizeBytes)
	}

	// Verify checksum.
	got := hex.EncodeToString(h.Sum(nil))
	if got != manifest.SHA256 {
		return fmt.Errorf("%w: got %s, want %s", ErrChecksumMismatch, got, manifest.SHA256)
	}

	// Fsync extent.
	if err := v.fd.Sync(); err != nil {
		return fmt.Errorf("blockvol: import sync: %w", err)
	}

	// Reset WAL and dirty map for clean state.
	v.dirtyMap.Clear()
	v.wal.Reset()
	v.super.WALHead = 0
	v.super.WALTail = 0
	v.super.Flags |= FlagImported
	if err := v.persistSuperblock(); err != nil {
		return fmt.Errorf("blockvol: import persist superblock: %w", err)
	}

	return nil
}

// ImportSnapshotForRebuild imports a snapshot artifact and converges all
// local runtime state to the exact snapshot boundary (baseLSN). This is
// the rebuild-oriented import primitive for V2 snapshot_tail execution.
//
// Unlike generic ImportSnapshot, this method:
//   - requires manifest.BaseLSN > 0 (exact boundary must be explicit)
//   - verifies the imported base equals the requested snapshotLSN
//   - converges WALCheckpointLSN, nextLSN, flusher, and receiver to baseLSN
//
// After this call, the volume is at exactly snapshotLSN. The caller can
// then replay WAL tail entries from snapshotLSN+1 to targetLSN.
func (v *BlockVol) ImportSnapshotForRebuild(ctx context.Context, manifest *SnapshotArtifactManifest, r io.Reader, snapshotLSN uint64) error {
	// Validate: manifest must carry explicit BaseLSN matching the requested boundary.
	if manifest.BaseLSN == 0 {
		return fmt.Errorf("blockvol: rebuild import requires explicit BaseLSN in manifest")
	}
	if manifest.BaseLSN != snapshotLSN {
		return fmt.Errorf("blockvol: rebuild import boundary mismatch: manifest BaseLSN=%d != requested snapshotLSN=%d",
			manifest.BaseLSN, snapshotLSN)
	}

	// Use generic import with AllowOverwrite (rebuild target may have stale data).
	if err := v.ImportSnapshot(ctx, manifest, r, ImportOptions{AllowOverwrite: true}); err != nil {
		return err
	}

	// Converge all runtime state to the exact snapshot boundary.
	// This is the same state handoff as RebuildInstaller.Commit but
	// for the snapshot_tail path, where the boundary is exact.
	v.mu.Lock()
	v.super.WALCheckpointLSN = snapshotLSN
	if _, err := v.fd.Seek(0, 0); err != nil {
		v.mu.Unlock()
		return fmt.Errorf("blockvol: rebuild import seek superblock: %w", err)
	}
	if _, err := v.super.WriteTo(v.fd); err != nil {
		v.mu.Unlock()
		return fmt.Errorf("blockvol: rebuild import write superblock: %w", err)
	}
	if err := v.fd.Sync(); err != nil {
		v.mu.Unlock()
		return fmt.Errorf("blockvol: rebuild import sync superblock: %w", err)
	}
	v.mu.Unlock()

	if v.flusher != nil {
		v.flusher.SetCheckpointLSN(snapshotLSN)
	}

	// Set nextLSN (unconditional, not monotonic — rebuild replaces truth).
	v.nextLSN.Store(snapshotLSN + 1)

	// Align receiver progress.
	if v.replRecv != nil {
		v.replRecv.mu.Lock()
		v.replRecv.receivedLSN = snapshotLSN
		v.replRecv.mu.Unlock()
	}

	return nil
}