package archive

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"path"
	"reflect"
	"strings"
	"time"

	"github.com/klauspost/compress/zstd"
)

var (
	// ErrChunkNotFound means the archive object implied by a thread and chunk start does not exist.
	ErrChunkNotFound = errors.New("archive chunk not found")
	// ErrChunkAlreadyExists means the archive object already exists but differs from the content
	// being written. Archive chunks are immutable; callers should treat identical chunks as success
	// and mismatches as a hard error.
	ErrChunkAlreadyExists = errors.New("archive chunk already exists")
)

// ReadDirection controls whether archive messages are returned from low seq to high seq or the reverse.
type ReadDirection int

const (
	// ReadAscending returns older messages first.
	ReadAscending ReadDirection = iota
	// ReadDescending returns newer messages first.
	ReadDescending
)

// Message is the canonical cold-history record stored inside one NDJSON archive object.
type Message struct {
	ID         string         `json:"id"`
	Seq        int64          `json:"seq"`
	SenderType string         `json:"sender_type"`
	SenderName string         `json:"sender_name"`
	Content    string         `json:"content"`
	Metadata   map[string]any `json:"metadata,omitempty"`
	Timestamp  time.Time      `json:"timestamp"`
}

// Storage reads and writes fixed-size archived message chunks.
type Storage interface {
	ReadChunk(ctx context.Context, workspaceID string, threadID string, chunkStartSeq int64) ([]Message, error)
	WriteChunk(ctx context.Context, workspaceID string, threadID string, chunkStartSeq int64, messages []Message) error
}

// ChunkStartSeq returns the aligned chunk start that contains seq.
func ChunkStartSeq(seq int64) (int64, error) {
	if seq <= 0 {
		return 0, fmt.Errorf("seq must be positive")
	}
	return ((seq-1)/ChunkMessageCount)*ChunkMessageCount + 1, nil
}

// ObjectKey builds the deterministic storage key for one archived chunk.
func ObjectKey(prefix string, workspaceID string, threadID string, chunkStartSeq int64) (string, error) {
	if err := validateObjectIdentifiers(workspaceID, threadID); err != nil {
		return "", err
	}
	if err := validateChunkStartSeq(chunkStartSeq); err != nil {
		return "", err
	}
	parts := make([]string, 0, 6)
	trimmedPrefix := strings.Trim(strings.TrimSpace(prefix), "/")
	if trimmedPrefix != "" {
		parts = append(parts, trimmedPrefix)
	}
	parts = append(parts,
		"workspaces",
		workspaceID,
		"threads",
		threadID,
		"chunks",
		fmt.Sprintf("%d.ndjson.zst", chunkStartSeq),
	)
	return path.Join(parts...), nil
}

// ReadRange loads messages from cold storage within [lowSeq, highSeq].
func ReadRange(ctx context.Context, storage Storage, workspaceID string, threadID string, lowSeq int64, highSeq int64, direction ReadDirection, limit int) ([]Message, error) {
	if storage == nil {
		return nil, fmt.Errorf("archive storage is required")
	}
	if limit <= 0 || highSeq <= 0 {
		return []Message{}, nil
	}
	if lowSeq < 1 {
		lowSeq = 1
	}
	if highSeq < lowSeq {
		return []Message{}, nil
	}

	firstChunkStart, err := ChunkStartSeq(lowSeq)
	if err != nil {
		return nil, err
	}
	lastChunkStart, err := ChunkStartSeq(highSeq)
	if err != nil {
		return nil, err
	}

	messages := make([]Message, 0, limit)
	if direction == ReadAscending {
		for chunkStartSeq := firstChunkStart; chunkStartSeq <= lastChunkStart && len(messages) < limit; chunkStartSeq += ChunkMessageCount {
			if err := ctx.Err(); err != nil {
				return nil, err
			}
			chunkMessages, err := storage.ReadChunk(ctx, workspaceID, threadID, chunkStartSeq)
			if err != nil {
				return nil, fmt.Errorf("read archive chunk %d: %w", chunkStartSeq, err)
			}
			for _, message := range chunkMessages {
				if message.Seq < lowSeq || message.Seq > highSeq {
					continue
				}
				messages = append(messages, message)
				if len(messages) == limit {
					break
				}
			}
		}
		return messages, nil
	}

	for chunkStartSeq := lastChunkStart; len(messages) < limit; {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		chunkMessages, err := storage.ReadChunk(ctx, workspaceID, threadID, chunkStartSeq)
		if err != nil {
			return nil, fmt.Errorf("read archive chunk %d: %w", chunkStartSeq, err)
		}
		for index := len(chunkMessages) - 1; index >= 0; index-- {
			message := chunkMessages[index]
			if message.Seq < lowSeq || message.Seq > highSeq {
				continue
			}
			messages = append(messages, message)
			if len(messages) == limit {
				break
			}
		}
		if chunkStartSeq == firstChunkStart {
			break
		}
		chunkStartSeq -= ChunkMessageCount
	}
	return messages, nil
}

func validateObjectIdentifiers(workspaceID string, threadID string) error {
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return fmt.Errorf("workspace id is required")
	}
	if strings.Contains(trimmedWorkspaceID, "/") {
		return fmt.Errorf("workspace id %q must not contain '/'", workspaceID)
	}
	trimmedThreadID := strings.TrimSpace(threadID)
	if trimmedThreadID == "" {
		return fmt.Errorf("thread id is required")
	}
	if strings.Contains(trimmedThreadID, "/") {
		return fmt.Errorf("thread id %q must not contain '/'", threadID)
	}
	return nil
}

func validateChunkStartSeq(chunkStartSeq int64) error {
	if chunkStartSeq <= 0 {
		return fmt.Errorf("chunk start seq must be positive")
	}
	if (chunkStartSeq-1)%ChunkMessageCount != 0 {
		return fmt.Errorf("chunk start seq %d is not aligned to %d messages", chunkStartSeq, ChunkMessageCount)
	}
	return nil
}

func validateChunkMessages(chunkStartSeq int64, messages []Message) error {
	if err := validateChunkStartSeq(chunkStartSeq); err != nil {
		return err
	}
	if int64(len(messages)) != ChunkMessageCount {
		return fmt.Errorf("chunk %d must contain exactly %d messages, got %d", chunkStartSeq, ChunkMessageCount, len(messages))
	}
	for index, message := range messages {
		expectedSeq := chunkStartSeq + int64(index)
		if strings.TrimSpace(message.ID) == "" {
			return fmt.Errorf("chunk %d seq %d message id is required", chunkStartSeq, expectedSeq)
		}
		if message.Seq != expectedSeq {
			return fmt.Errorf("chunk %d expected seq %d at index %d, got %d", chunkStartSeq, expectedSeq, index, message.Seq)
		}
		if strings.TrimSpace(message.SenderType) == "" {
			return fmt.Errorf("chunk %d seq %d sender_type is required", chunkStartSeq, message.Seq)
		}
		if strings.TrimSpace(message.SenderName) == "" {
			return fmt.Errorf("chunk %d seq %d sender_name is required", chunkStartSeq, message.Seq)
		}
	}
	return nil
}

func chunkMessagesEqual(left []Message, right []Message) bool {
	if len(left) != len(right) {
		return false
	}
	for index := range left {
		if left[index].ID != right[index].ID {
			return false
		}
		if left[index].Seq != right[index].Seq {
			return false
		}
		if left[index].SenderType != right[index].SenderType {
			return false
		}
		if left[index].SenderName != right[index].SenderName {
			return false
		}
		if left[index].Content != right[index].Content {
			return false
		}
		if !reflect.DeepEqual(left[index].Metadata, right[index].Metadata) {
			return false
		}
		if !archiveTimestampEqual(left[index].Timestamp, right[index].Timestamp) {
			return false
		}
	}
	return true
}

func archiveTimestampEqual(left time.Time, right time.Time) bool {
	return left.UTC().Truncate(time.Microsecond).Equal(right.UTC().Truncate(time.Microsecond))
}

func encodeChunk(messages []Message) ([]byte, error) {
	var buffer bytes.Buffer
	encoder, err := zstd.NewWriter(&buffer)
	if err != nil {
		return nil, fmt.Errorf("create zstd encoder: %w", err)
	}
	for _, message := range messages {
		encoded, err := json.Marshal(message)
		if err != nil {
			encoder.Close()
			return nil, fmt.Errorf("marshal archive message: %w", err)
		}
		if _, err := encoder.Write(encoded); err != nil {
			encoder.Close()
			return nil, fmt.Errorf("write archive message: %w", err)
		}
		if _, err := encoder.Write([]byte{'\n'}); err != nil {
			encoder.Close()
			return nil, fmt.Errorf("write archive newline: %w", err)
		}
	}
	if err := encoder.Close(); err != nil {
		return nil, fmt.Errorf("close zstd encoder: %w", err)
	}
	return buffer.Bytes(), nil
}

func decodeChunk(reader io.Reader) ([]Message, error) {
	decoder, err := zstd.NewReader(reader)
	if err != nil {
		return nil, fmt.Errorf("create zstd decoder: %w", err)
	}
	defer decoder.Close()

	buffered := bufio.NewReader(decoder)
	messages := make([]Message, 0, ChunkMessageCount)
	for {
		line, err := buffered.ReadBytes('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			return nil, fmt.Errorf("read archive line: %w", err)
		}
		trimmedLine := bytes.TrimSpace(line)
		if len(trimmedLine) > 0 {
			var message Message
			if err := json.Unmarshal(trimmedLine, &message); err != nil {
				return nil, fmt.Errorf("decode archive message: %w", err)
			}
			messages = append(messages, message)
		}
		if errors.Is(err, io.EOF) {
			break
		}
	}
	return messages, nil
}