package memory

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/contentpayload"
	contentservice "github.com/tidbcloud/AgentTeam/internal/service/content"
)

const (
	memoryContentMaxVersions            = uint64(512)
	memoryContentMaxStoredBytes         = int64(1 << 30)
	memoryContentMaxPins                = uint64(64)
	memoryContentMaxFiles               = 2_000
	memoryContentMaxFileBytes           = int64(16 << 20)
	memoryContentMaxTotalBytes          = int64(64 << 20)
	memoryContentMaxManifestBytes       = int64(8 << 20)
	memoryContentOperationTimeout       = 2 * time.Minute
	memoryContentHistoryLimit           = 100
	memoryContentMaxRestoreArchiveBytes = memoryContentMaxTotalBytes + 2*memoryContentMaxManifestBytes + (4 << 20)
)

var memoryContentTreeLimits = contentpayload.Limits{
	MaxFiles:         memoryContentMaxFiles,
	MaxFileBytes:     memoryContentMaxFileBytes,
	MaxTotalBytes:    memoryContentMaxTotalBytes,
	MaxManifestBytes: memoryContentMaxManifestBytes,
}

type nativeMemoryContentService interface {
	FindContentByBinding(context.Context, string, string, string) (contentservice.Content, bool, error)
	CommitVersion(context.Context, contentservice.CommitVersionRequest) (contentservice.CommitVersionResult, error)
	ListVersions(context.Context, string, string, int) ([]contentservice.Version, error)
	Compare(context.Context, string, string, string, string) ([]contentservice.CompareEntry, error)
	OpenTreeArchive(context.Context, string, string, string, contentservice.ArchiveLimits) (io.ReadCloser, contentservice.ArchiveMeta, error)
	CreatePin(context.Context, contentservice.CreatePinRequest) (contentservice.Pin, error)
	ListPins(context.Context, string, string, int) ([]contentservice.Pin, error)
}

type NativeMemoryContentRequest struct {
	WorkspaceID string
	MemoryID    string
	ContentID   string
	VersionID   string
	Generation  uint64
	Actor       string
}

type NativeMemoryContentFile struct {
	Path       string
	Content    []byte
	Executable bool
}

type NativeMemoryCheckpointRequest struct {
	NativeMemoryContentRequest
	OperationID            string
	ReplicaID              string
	CoveredFromSequence    uint64
	CoveredThroughSequence uint64
	Files                  []NativeMemoryContentFile
}

type NativeMemoryVersionResult struct {
	ContentID       string
	VersionID       string
	ParentVersionID string
	Generation      uint64
	RootDigest      string
	SizeBytes       int64
	FileCount       int
	CreatedAt       time.Time
}

type NativeMemoryContent interface {
	Resolve(context.Context, string, string, string) (NativeMemoryContentRequest, error)
	PublishCheckpoint(context.Context, NativeMemoryCheckpointRequest) (NativeMemoryVersionResult, error)
	ListVersions(context.Context, NativeMemoryContentRequest, int) ([]contentservice.Version, error)
	Compare(context.Context, NativeMemoryContentRequest, string, string) ([]contentservice.CompareEntry, error)
	CreateBackup(context.Context, NativeMemoryContentRequest, string) (contentservice.Pin, error)
	ListBackups(context.Context, NativeMemoryContentRequest, int) ([]contentservice.Pin, error)
	Restore(context.Context, NativeMemoryContentRequest, string, string) (NativeMemoryVersionResult, error)
	OpenArchive(context.Context, NativeMemoryContentRequest, string) (io.ReadCloser, contentservice.ArchiveMeta, error)
}

// NativeMemoryContentStore owns native cloud snapshots for one exact managed
// Memory binding. It accepts no repository, branch, ref, or local path and
// therefore cannot silently fall back to AGS.
type NativeMemoryContentStore struct{ Content nativeMemoryContentService }

func (store NativeMemoryContentStore) Resolve(ctx context.Context, workspaceID, memoryID, actor string) (NativeMemoryContentRequest, error) {
	request := NativeMemoryContentRequest{
		WorkspaceID: strings.TrimSpace(workspaceID), MemoryID: strings.TrimSpace(memoryID), Actor: strings.TrimSpace(actor),
	}
	if store.Content == nil || request.WorkspaceID == "" || request.MemoryID == "" || request.Actor == "" {
		return NativeMemoryContentRequest{}, invalidArgument("native memory content is not configured or its exact scope is incomplete")
	}
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessReadMetadata))
	defer cancel()
	content, found, err := store.Content.FindContentByBinding(readCtx, request.WorkspaceID, memoryContentDomain, request.MemoryID)
	if err != nil {
		return NativeMemoryContentRequest{}, mapNativeMemoryContentError(err)
	}
	if !found || content.Model != contentservice.TreeModel {
		return NativeMemoryContentRequest{}, notFound("memory native tree lineage not found")
	}
	if err := validateMemoryContentQuotaProfile(content); err != nil {
		return NativeMemoryContentRequest{}, err
	}
	request.ContentID = content.ID
	request.VersionID = content.HeadVersionID
	request.Generation = content.Generation
	return request, nil
}

func (store NativeMemoryContentStore) PublishCheckpoint(ctx context.Context, request NativeMemoryCheckpointRequest) (NativeMemoryVersionResult, error) {
	if strings.TrimSpace(request.OperationID) == "" || strings.TrimSpace(request.ReplicaID) == "" ||
		request.CoveredFromSequence == 0 || request.CoveredThroughSequence < request.CoveredFromSequence {
		return NativeMemoryVersionResult{}, invalidArgument("memory checkpoint identity and covered sequence range are required")
	}
	return store.publish(ctx, request.NativeMemoryContentRequest, request.Files, checkpointOperationKey(request), true)
}

func (store NativeMemoryContentStore) ListVersions(ctx context.Context, request NativeMemoryContentRequest, limit int) ([]contentservice.Version, error) {
	if limit <= 0 || limit > memoryContentHistoryLimit {
		return nil, invalidArgument("memory history limit must be between 1 and 100")
	}
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return nil, err
	}
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessReadHistory))
	defer cancel()
	return store.Content.ListVersions(readCtx, request.WorkspaceID, content.ID, limit)
}

func (store NativeMemoryContentStore) Compare(ctx context.Context, request NativeMemoryContentRequest, beforeVersionID, afterVersionID string) ([]contentservice.CompareEntry, error) {
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return nil, err
	}
	beforeVersionID = strings.TrimSpace(beforeVersionID)
	afterVersionID = strings.TrimSpace(afterVersionID)
	if beforeVersionID == "" || afterVersionID == "" {
		return nil, invalidArgument("memory compare requires two exact versions")
	}
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessCompare))
	defer cancel()
	return store.Content.Compare(readCtx, request.WorkspaceID, content.ID, beforeVersionID, afterVersionID)
}

func (store NativeMemoryContentStore) CreateBackup(ctx context.Context, request NativeMemoryContentRequest, label string) (contentservice.Pin, error) {
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return contentservice.Pin{}, err
	}
	if strings.TrimSpace(request.VersionID) == "" || strings.TrimSpace(request.VersionID) != content.HeadVersionID {
		return contentservice.Pin{}, conflict("memory backup requires the exact cloud head")
	}
	writeCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessManagePin))
	defer cancel()
	pins, err := store.Content.ListPins(writeCtx, request.WorkspaceID, content.ID, int(memoryContentMaxPins))
	if err != nil {
		return contentservice.Pin{}, err
	}
	if len(pins) >= int(memoryContentMaxPins) {
		return contentservice.Pin{}, conflict("memory backup quota exceeded")
	}
	return store.Content.CreatePin(writeCtx, contentservice.CreatePinRequest{
		WorkspaceID: request.WorkspaceID, ContentID: content.ID, VersionID: request.VersionID,
		Label: strings.TrimSpace(label), Actor: request.Actor,
	})
}

func (store NativeMemoryContentStore) ListBackups(ctx context.Context, request NativeMemoryContentRequest, limit int) ([]contentservice.Pin, error) {
	if limit <= 0 || limit > int(memoryContentMaxPins) {
		return nil, invalidArgument("memory backup limit must be between 1 and 64")
	}
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return nil, err
	}
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessManagePin))
	defer cancel()
	return store.Content.ListPins(readCtx, request.WorkspaceID, content.ID, limit)
}

func (store NativeMemoryContentStore) Restore(ctx context.Context, request NativeMemoryContentRequest, targetVersionID, operationID string) (NativeMemoryVersionResult, error) {
	if strings.TrimSpace(targetVersionID) == "" || strings.TrimSpace(operationID) == "" {
		return NativeMemoryVersionResult{}, invalidArgument("memory restore target and operation are required")
	}
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return NativeMemoryVersionResult{}, err
	}
	files, err := store.snapshotFiles(ctx, request, content.ID, targetVersionID)
	if err != nil {
		return NativeMemoryVersionResult{}, err
	}
	key := memoryRestoreOperationKey(request, targetVersionID, operationID)
	return store.publish(ctx, request, files, key, true)
}

func (store NativeMemoryContentStore) OpenArchive(ctx context.Context, request NativeMemoryContentRequest, versionID string) (io.ReadCloser, contentservice.ArchiveMeta, error) {
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return nil, contentservice.ArchiveMeta{}, err
	}
	versionID = strings.TrimSpace(versionID)
	if versionID == "" {
		return nil, contentservice.ArchiveMeta{}, invalidArgument("memory archive requires one exact version")
	}
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessReadArchive))
	body, meta, err := store.Content.OpenTreeArchive(readCtx, request.WorkspaceID, content.ID, versionID, contentservice.ArchiveLimits{
		MaxFiles: memoryContentMaxFiles, MaxFileBytes: memoryContentMaxFileBytes, MaxTotalBytes: memoryContentMaxTotalBytes,
	})
	if err != nil {
		cancel()
		return nil, contentservice.ArchiveMeta{}, mapNativeMemoryContentError(err)
	}
	return &memoryContentReadCloser{ReadCloser: body, cancel: cancel}, meta, nil
}

func (store NativeMemoryContentStore) publish(ctx context.Context, request NativeMemoryContentRequest, files []NativeMemoryContentFile, operationKey string, checkpoint bool) (NativeMemoryVersionResult, error) {
	content, err := store.exactContent(ctx, request)
	if err != nil {
		return NativeMemoryVersionResult{}, err
	}
	treeFiles, size, err := memoryTreeFiles(files)
	if err != nil {
		return NativeMemoryVersionResult{}, err
	}
	// Only preflight a live-head request. A retry after a lost acknowledgement
	// intentionally carries the prior head; the Content Store operation receipt
	// must get the chance to replay before its transactional head/quota checks.
	if strings.TrimSpace(request.VersionID) == content.HeadVersionID &&
		(content.RetainedVersionCount >= memoryContentMaxVersions || size > memoryContentMaxStoredBytes-content.StoredPayloadBytes) {
		return NativeMemoryVersionResult{}, conflict("memory checkpoint quota exceeded")
	}
	writeCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessPublish))
	defer cancel()
	result, err := store.Content.CommitVersion(writeCtx, contentservice.CommitVersionRequest{
		WorkspaceID: request.WorkspaceID, ContentID: content.ID,
		ExpectedHeadVersionID: request.VersionID, ExpectedGeneration: request.Generation,
		Actor: request.Actor, Model: contentservice.TreeModel, Checkpoint: checkpoint,
		Files: treeFiles, TreeLimits: memoryContentTreeLimits,
		OperationKey: operationKey, OperationNamespace: "memory.content", OperationVersion: 1,
	})
	if err != nil {
		return NativeMemoryVersionResult{}, mapNativeMemoryContentError(err)
	}
	return nativeMemoryVersionResult(result.Version), nil
}

func (store NativeMemoryContentStore) exactContent(ctx context.Context, request NativeMemoryContentRequest) (contentservice.Content, error) {
	if store.Content == nil || strings.TrimSpace(request.WorkspaceID) == "" || strings.TrimSpace(request.MemoryID) == "" ||
		strings.TrimSpace(request.ContentID) == "" || strings.TrimSpace(request.Actor) == "" || request.Generation == 0 {
		return contentservice.Content{}, invalidArgument("native memory content is not configured or its exact scope is incomplete")
	}
	lookup := request
	lookup.ContentID = ""
	lookup.Generation = 0
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, lookup, contentservice.AccessReadMetadata))
	defer cancel()
	content, found, err := store.Content.FindContentByBinding(readCtx, request.WorkspaceID, memoryContentDomain, request.MemoryID)
	if err != nil {
		return contentservice.Content{}, mapNativeMemoryContentError(err)
	}
	if !found || content.ID != strings.TrimSpace(request.ContentID) || content.Model != contentservice.TreeModel || content.Generation != request.Generation {
		return contentservice.Content{}, conflict("memory content binding does not match the exact tree lineage")
	}
	// Memory's finite quota is part of its product contract. Until the shared
	// Content Store can create a lineage with this exact profile, refusing a
	// wider generic tree quota is safer than relying on racy domain preflights.
	if err := validateMemoryContentQuotaProfile(content); err != nil {
		return contentservice.Content{}, err
	}
	return content, nil
}

func validateMemoryContentQuotaProfile(content contentservice.Content) error {
	if content.MaxRetainedVersions != memoryContentMaxVersions || content.MaxStoredPayloadBytes != memoryContentMaxStoredBytes || content.MaxPins != memoryContentMaxPins {
		return fmt.Errorf("memory content quota profile is unavailable")
	}
	return nil
}

func (store NativeMemoryContentStore) snapshotFiles(ctx context.Context, request NativeMemoryContentRequest, contentID, versionID string) ([]NativeMemoryContentFile, error) {
	readCtx, cancel := boundedMemoryContentContext(memoryContentContext(ctx, request, contentservice.AccessReadArchive))
	defer cancel()
	versionID = strings.TrimSpace(versionID)
	body, meta, err := store.Content.OpenTreeArchive(readCtx, request.WorkspaceID, contentID, versionID, contentservice.ArchiveLimits{
		MaxFiles: memoryContentMaxFiles, MaxFileBytes: memoryContentMaxFileBytes, MaxTotalBytes: memoryContentMaxTotalBytes,
	})
	if err != nil {
		return nil, mapNativeMemoryContentError(err)
	}
	encoded, readErr := io.ReadAll(io.LimitReader(body, memoryContentMaxRestoreArchiveBytes+1))
	closeErr := body.Close()
	if readErr != nil || closeErr != nil {
		return nil, errors.Join(readErr, closeErr)
	}
	if int64(len(encoded)) > memoryContentMaxRestoreArchiveBytes {
		return nil, conflict("memory restore archive exceeds bounds")
	}
	archive, err := zip.NewReader(bytes.NewReader(encoded), int64(len(encoded)))
	if err != nil || len(archive.File) != meta.FileCount || len(archive.File) > memoryContentMaxFiles {
		return nil, conflict("memory restore archive does not match exact tree")
	}
	files := make([]NativeMemoryContentFile, 0, len(archive.File))
	seen := make(map[string]struct{}, len(archive.File))
	var total int64
	for _, item := range archive.File {
		path := strings.TrimSpace(item.Name)
		size := int64(item.UncompressedSize64)
		if item.Method != zip.Store || !item.Mode().IsRegular() || !validMemoryMigrationPath(path) || path != item.Name {
			return nil, conflict("memory restore archive contains an invalid entry")
		}
		if _, duplicate := seen[path]; duplicate {
			return nil, conflict("memory restore archive contains a duplicate path")
		}
		seen[path] = struct{}{}
		if size < 0 || size > memoryContentMaxFileBytes || total > memoryContentMaxTotalBytes-size {
			return nil, conflict("memory restore snapshot exceeds bounds")
		}
		entry, err := item.Open()
		if err != nil {
			return nil, conflict("memory restore archive entry cannot be opened")
		}
		data, readErr := io.ReadAll(io.LimitReader(entry, size+1))
		closeErr := entry.Close()
		if readErr != nil || closeErr != nil {
			return nil, errors.Join(readErr, closeErr)
		}
		if int64(len(data)) != size {
			return nil, conflict("memory restore archive entry size is inconsistent")
		}
		total += size
		files = append(files, NativeMemoryContentFile{Path: path, Content: data, Executable: item.Mode().Perm()&0o111 != 0})
	}
	if total != meta.SizeBytes {
		return nil, conflict("memory restore archive size does not match exact tree")
	}
	return files, nil
}

func memoryTreeFiles(files []NativeMemoryContentFile) ([]contentpayload.File, int64, error) {
	if len(files) > memoryContentMaxFiles {
		return nil, 0, conflict("memory checkpoint exceeds file limit")
	}
	result := make([]contentpayload.File, len(files))
	var total int64
	for index, file := range files {
		size := int64(len(file.Content))
		if size > memoryContentMaxFileBytes || total > memoryContentMaxTotalBytes-size {
			return nil, 0, conflict("memory checkpoint exceeds byte limits")
		}
		total += size
		result[index] = contentpayload.File{Path: file.Path, Size: size, Executable: file.Executable, Content: bytes.NewReader(file.Content)}
	}
	return result, total, nil
}

func memoryContentContext(ctx context.Context, request NativeMemoryContentRequest, action contentservice.AccessAction) context.Context {
	return contentservice.WithAccessPrincipal(ctx, contentservice.AccessPrincipal{
		Subject: strings.TrimSpace(request.Actor), WorkspaceID: strings.TrimSpace(request.WorkspaceID),
		Domain: memoryContentDomain, BindingKey: strings.TrimSpace(request.MemoryID), ContentID: strings.TrimSpace(request.ContentID),
		Generation: request.Generation, Actions: []contentservice.AccessAction{action},
	})
}

func boundedMemoryContentContext(ctx context.Context) (context.Context, context.CancelFunc) {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithTimeout(ctx, memoryContentOperationTimeout)
}

func checkpointOperationKey(request NativeMemoryCheckpointRequest) string {
	hash := sha256.New()
	for _, part := range []string{request.MemoryID, request.OperationID, request.VersionID, request.ReplicaID} {
		writeMemoryOperationPart(hash, []byte(strings.TrimSpace(part)))
	}
	writeMemoryOperationUint64(hash, request.Generation)
	writeMemoryOperationUint64(hash, request.CoveredFromSequence)
	writeMemoryOperationUint64(hash, request.CoveredThroughSequence)
	return "memory.content:" + hex.EncodeToString(hash.Sum(nil))
}

func memoryRestoreOperationKey(request NativeMemoryContentRequest, targetVersionID, operationID string) string {
	hash := sha256.New()
	for _, value := range []string{request.MemoryID, operationID, request.VersionID, targetVersionID} {
		writeMemoryOperationPart(hash, []byte(strings.TrimSpace(value)))
	}
	writeMemoryOperationUint64(hash, request.Generation)
	return "memory.content:" + hex.EncodeToString(hash.Sum(nil))
}

type memoryContentReadCloser struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (reader *memoryContentReadCloser) Close() error {
	reader.cancel()
	return reader.ReadCloser.Close()
}

func writeMemoryOperationPart(writer io.Writer, value []byte) {
	var size [8]byte
	binary.BigEndian.PutUint64(size[:], uint64(len(value)))
	_, _ = writer.Write(size[:])
	_, _ = writer.Write(value)
}

func writeMemoryOperationUint64(writer io.Writer, value uint64) {
	var encoded [8]byte
	binary.BigEndian.PutUint64(encoded[:], value)
	_, _ = writer.Write(encoded[:])
}

func nativeMemoryVersionResult(version contentservice.Version) NativeMemoryVersionResult {
	return NativeMemoryVersionResult{
		ContentID: version.ContentID, VersionID: version.ID, ParentVersionID: version.ParentVersionID,
		Generation: version.Generation, RootDigest: version.RootDigest, SizeBytes: version.SizeBytes,
		FileCount: version.FileCount, CreatedAt: version.CreatedAt,
	}
}

func mapNativeMemoryContentError(err error) error {
	switch {
	case err == nil:
		return nil
	case errors.Is(err, contentservice.ErrAccessDenied):
		return forbidden("memory content access denied")
	case contentservice.IsNotFound(err):
		return notFound("memory content version not found")
	case contentservice.IsConflict(err):
		return conflict("memory cloud head, generation, or operation is stale")
	case contentservice.IsQuotaExceeded(err):
		return conflict("memory checkpoint quota exceeded")
	default:
		return err
	}
}

var _ NativeMemoryContent = NativeMemoryContentStore{}