package taskrun

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

const ManifestSchemaVersion = 1

const (
	MemberKindDocument      = "document"
	MemberKindAssistantTurn = "assistant_turn"
	MemberKindRunManifest   = "run_manifest"
	MemberKindAttachment    = "attachment"
	MemberKindEphemeral     = "ephemeral"
)

// ManifestMember is a retained output of a task run. Evidence/Gaps are not members.
type ManifestMember struct {
	LogicalID string `json:"logical_id"`
	Digest    string `json:"digest"`
	Kind      string `json:"kind"`
}

// RunManifest is the typed terminal protocol. Workspace copies are rebuildable, not the source of truth.
type RunManifest struct {
	SchemaVersion   int                        `json:"schema_version"`
	RunID           string                     `json:"run_id"`
	WorkspaceID     string                     `json:"workspace_id"`
	StorageIdentity string                     `json:"storage_identity"`
	Members         []ManifestMember           `json:"members"`
	Extensions      map[string]json.RawMessage `json:"extensions,omitempty"`
	CanonicalBytes  []byte                     `json:"-"`
	SHA256          string                     `json:"-"`
}

type canonicalManifest struct {
	SchemaVersion   int              `json:"schema_version"`
	RunID           string           `json:"run_id"`
	WorkspaceID     string           `json:"workspace_id"`
	StorageIdentity string           `json:"storage_identity"`
	Members         []ManifestMember `json:"members"`
	Extensions      json.RawMessage  `json:"extensions,omitempty"`
}

// Freeze canonicalizes members, encodes stable JSON, and fills SHA-256.
func (m *RunManifest) Freeze() error {
	if m == nil {
		return fmt.Errorf("taskrun: manifest required")
	}
	if strings.TrimSpace(m.RunID) == "" {
		return fmt.Errorf("taskrun: manifest run_id required")
	}
	if m.SchemaVersion == 0 {
		m.SchemaVersion = ManifestSchemaVersion
	}
	if m.SchemaVersion != ManifestSchemaVersion {
		return fmt.Errorf("taskrun: unsupported manifest schema %d", m.SchemaVersion)
	}
	members := append([]ManifestMember(nil), m.Members...)
	sort.Slice(members, func(i, j int) bool {
		if members[i].LogicalID == members[j].LogicalID {
			return members[i].Kind < members[j].Kind
		}
		return members[i].LogicalID < members[j].LogicalID
	})
	for i := range members {
		members[i].LogicalID = strings.TrimSpace(members[i].LogicalID)
		members[i].Digest = strings.ToLower(strings.TrimSpace(members[i].Digest))
		members[i].Kind = strings.TrimSpace(members[i].Kind)
		if members[i].LogicalID == "" || members[i].Digest == "" || members[i].Kind == "" {
			return fmt.Errorf("taskrun: manifest member incomplete")
		}
	}
	m.Members = members

	canon := canonicalManifest{
		SchemaVersion:   m.SchemaVersion,
		RunID:           m.RunID,
		WorkspaceID:     m.WorkspaceID,
		StorageIdentity: m.StorageIdentity,
		Members:         members,
		Extensions:      marshalSortedExtensions(m.Extensions),
	}
	raw, err := marshalCanonical(canon)
	if err != nil {
		return err
	}
	sum := sha256.Sum256(raw)
	m.CanonicalBytes = raw
	m.SHA256 = hex.EncodeToString(sum[:])
	return nil
}

func marshalCanonical(v any) ([]byte, error) {
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(v); err != nil {
		return nil, err
	}
	return bytes.TrimSpace(buf.Bytes()), nil
}

func marshalSortedExtensions(ext map[string]json.RawMessage) json.RawMessage {
	if len(ext) == 0 {
		return nil
	}
	keys := make([]string, 0, len(ext))
	for k := range ext {
		if strings.TrimSpace(k) == "" || !strings.Contains(k, "/") {
			continue
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	if len(keys) == 0 {
		return nil
	}
	var buf bytes.Buffer
	buf.WriteByte('{')
	for i, k := range keys {
		if i > 0 {
			buf.WriteByte(',')
		}
		kb, _ := json.Marshal(k)
		buf.Write(kb)
		buf.WriteByte(':')
		if len(ext[k]) == 0 {
			buf.WriteString("null")
		} else {
			buf.Write(ext[k])
		}
	}
	buf.WriteByte('}')
	return buf.Bytes()
}

func memberNeedsRetention(kind string) bool {
	switch kind {
	case MemberKindDocument, MemberKindAssistantTurn, MemberKindRunManifest, MemberKindAttachment:
		return true
	default:
		return false
	}
}