package capsule

import (
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"slices"
	"strings"
	"time"
	"unicode/utf8"

	"g.activething.com/go/DateKeys/internal/pathrule"
)

// Source is a file that EncryptFiles writes into a format 3 capsule.
type Source struct {
	// Path is the path of the file in the capsule, relative, with '/'
	// between its segments (spec §29.5). It is stored as given: EncryptFiles
	// rejects a path that breaks a rule, with a message that names the rule
	// and the character, and never corrects it (spec §62.1 rule 15).
	Path string
	// Size is the number of bytes of the file. EncryptFiles checks it in
	// each of its two readings.
	Size int64
	// ModTime is the modification time of the file at its source, taken
	// when it is loaded, as os.FileInfo.ModTime gives it, or the zero Time
	// when unknown. It is stored in seconds when it falls from 1970-01-01 to
	// 9999-12-31T23:59:59Z, and omitted otherwise, never clipped (spec
	// §62.1 rule 16). It is informative: it proves nothing.
	ModTime time.Time
	// Open returns a reader of the file from its start. EncryptFiles calls
	// it twice, and closes each reader.
	Open func() (io.ReadCloser, error)
}

// EncryptFiles writes a format 3 .dkc holding the files of sources and the
// comment and declared author of opts (spec §29.2 to §29.6, §61, §62,
// §62.1). It needs no network: the round is resolved locally and tlock uses
// only the pinned public key.
//
// It reads each file twice, and writes nothing to dst before the second
// reading. First it checks the paths and the texts with the rules of the
// reader, measures L with a head whose salt and SHA-256 are zero, as long as
// the final one, and hashes each file. Then it seals the control, with L,
// and streams PAYLOAD_AGE, reading each file again: a file whose size or
// SHA-256 has changed makes it fail (spec §62.1 rule 18), and dst then holds
// a partial capsule that must be discarded and never presented as a capsule
// (rule 9). The files go in the byte order of their paths, whatever the
// order of sources (R8), without the empty folders, which a path cannot
// name.
//
// The head, the control and the security area are decoded with the rules of
// the reader before anything is written (spec §62.1 rule 17), and the
// self-checks of Encrypt apply too. The security area is in an area of
// AreaLen bytes, or LargeAreaLen with opts.LargeArea (rule 13): empty, or
// with the signature of opts.AuthorKey, made before anything is written
// with the commitments of the final control and the head (spec v0.11, §29.8)
// and checked with the strict profile (rule 19).
//
// opts is as for Encrypt, with the head in Comment, Author and the head
// extensions, and with Length 0: L is the length of BODY.
func EncryptFiles(dst io.Writer, sources []Source, opts EncryptOptions) (*Result, error) {
	if opts.Length != 0 {
		return nil, errors.New("capsule: EncryptOptions.Length is for Encrypt: EncryptFiles computes L from the files")
	}
	common := uint32(AreaLen)
	switch t := opts.TestAreaLen; {
	case t == 0:
	case !opts.TestVectors:
		return nil, errors.New("capsule: EncryptOptions.TestAreaLen is for generators of test vectors: it needs TestVectors (spec §62.1 rule 13)")
	case t%AreaUnit != 0 || t > MaxAreaLen || opts.LargeArea:
		return nil, fmt.Errorf("capsule: a test area of %d bytes: a multiple of %d up to %d, without LargeArea", t, AreaUnit, MaxAreaLen)
	default:
		common = t
	}
	s, err := newSealer(opts, 0)
	if err != nil {
		return nil, err
	}
	h, order, err := newHead(sources, opts)
	if err != nil {
		return nil, err
	}

	// Step 2 of spec §61: L, with a head as long as the final one.
	measured, err := EncodeHead(h)
	if err != nil {
		return nil, err
	}
	if len(measured) > MaxHeadLen {
		return nil, fmt.Errorf("capsule: the head is %d bytes, more than %d: fewer files or shorter paths", len(measured), MaxHeadLen)
	}
	if err := selfCheckHead(measured); err != nil {
		return nil, err
	}
	var content uint64
	if n := len(h.Files); n > 0 {
		content = h.Files[n-1].End
	}
	// newHead bounds content by MaxPayloadLength: the sum does not overflow.
	// This L is the first one, with the common area: the signature decides
	// whether the area grows (spec §29.8), and prepare returns the final L.
	area := common
	length := BodyFrameSize + uint64(area) + uint64(len(measured)) + content
	if _, err := PaddedLength(length, s.code); err != nil {
		return nil, err
	}

	// Step 3: the first reading, for the SHA-256 of each file.
	for i := range h.Files {
		if h.Files[i].SHA256, err = readSource(nil, sources[order[i]], h.Files[i].Size, false); err != nil {
			return nil, err
		}
	}

	// Step 12: the head, with a fresh salt, and SECURITY_CBOR, decoded with
	// the rules of the reader; write decodes CONTROL_CBOR.
	_, _ = rand.Read(h.Salt[:]) // never fails since Go 1.24
	head, err := EncodeHead(h)
	if err != nil {
		return nil, err
	}
	if len(head) != len(measured) {
		return nil, fmt.Errorf("capsule: internal error: the head is %d bytes, measured %d", len(head), len(measured))
	}
	if err := selfCheckHead(head); err != nil {
		return nil, err
	}

	// SECURITY_CBOR and the frame are final once the control is: the
	// signature commits to it (spec v0.11, §29.8). prepare builds them
	// before anything is written, and write decodes CONTROL_CBOR.
	var security []byte
	var fb [BodyFrameSize]byte
	prepare := func(c *Control) (uint64, error) {
		var err error
		if security, err = s.security(c, head); err != nil {
			return 0, err
		}
		// The area is the common one, or the large one only when what was
		// signed does not fit and LargeArea allows it (§62.1 rule 13).
		switch {
		case len(security) <= int(common):
			area = common
		case !opts.LargeArea:
			return 0, fmt.Errorf("capsule: SECURITY_CBOR of %d bytes does not fit in the area of %d bytes: LargeArea lets the writer widen it to %d", len(security), common, LargeAreaLen)
		case len(security) <= LargeAreaLen:
			area = LargeAreaLen
		default:
			return 0, fmt.Errorf("capsule: SECURITY_CBOR of %d bytes does not fit in the area of %d bytes, the largest", len(security), LargeAreaLen)
		}
		final := BodyFrameSize + uint64(area) + uint64(len(head)) + content
		if _, err := PaddedLength(final, s.code); err != nil {
			return 0, err
		}
		frame := BodyFrame{AreaLen: area, SecurityLen: uint32(len(security)), HeadLen: uint32(len(head))}
		fb = frame.Bytes()
		if _, err := ParseBodyFrame(fb[:], final); err != nil {
			return 0, fmt.Errorf("capsule: self-check: %w", err)
		}
		if err := CheckHeadEnd(h, frame.ContentLength(final)); err != nil {
			return 0, fmt.Errorf("capsule: self-check: %w", err)
		}
		return final, nil
	}

	// Step 16: BODY, and the second reading of each file.
	res, err := s.write(dst, Format3, length, prepare, func(w io.Writer) error {
		for _, b := range [][]byte{fb[:], security, make([]byte, int(area)-len(security)), head} {
			if _, err := w.Write(b); err != nil {
				return err
			}
		}
		for i := range h.Files {
			f := &h.Files[i]
			sum, err := readSource(w, sources[order[i]], f.Size, true)
			if err != nil {
				return err
			}
			if sum != f.SHA256 {
				return fmt.Errorf("capsule: file %q changed after its first reading: its SHA-256 is another", f.Path)
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	res.Head, res.Security = h, s.verdicts
	return res, nil
}

// security returns SECURITY_CBOR for the capsule whose final control is c and
// whose head is head: empty, or with the signature of opts.AuthorKey or
// opts.CMSSigner, and the seal of opts.Sealer (spec v0.11, §29.3, §29.8 to
// §29.11). The signature is made first and the seal after it, which seals it.
// It decodes and evaluates what it returns with the rules of the reader, in
// the context of this capsule (§62.1 rules 17, 19 and 21). The caller decides
// the area from its length.
func (s *sealer) security(c *Control, head []byte) ([]byte, error) {
	o := s.opts
	sc := &SecurityContext{HeadDigest: HeadDigest(head), RoundTime: s.unlock}
	if o.AuthorKey == nil && o.CMSSigner == nil && o.Sealer == nil {
		security := EncodeSecurity()
		v := EvaluateSecurityIn(security, sc)
		if v != (Verdicts{Signature: VerdictNoSignature, Seal: VerdictNoSeal}) {
			return nil, fmt.Errorf("capsule: self-check: the reader finds the verdicts %s and %s in this security area", v.Signature, v.Seal)
		}
		s.verdicts = v
		return security, nil
	}
	var err error
	if sc.ControlCommit, err = ControlCommit(c, Format3); err != nil {
		return nil, err
	}

	var signature, seal []byte
	wantSig := VerdictNoSignature
	var key [32]byte
	switch {
	case o.AuthorKey != nil:
		pub := o.AuthorKey.Public()
		if len(pub) != 32 {
			return nil, fmt.Errorf("capsule: the author key is %d bytes, not 32", len(pub))
		}
		msg := AuthorMessage(sc.ControlCommit, sc.HeadDigest, SignersDigest(AlgEd25519, nil))
		if signature, err = EncodeAuthorSignature(AlgEd25519, pub, o.AuthorKey.Sign(msg)); err != nil {
			return nil, err
		}
		wantSig = VerdictSignedOther
		copy(key[:], pub)
	case o.CMSSigner != nil:
		list, err := EncodeSigners(o.CMSSigner.Signers())
		if err != nil {
			return nil, err
		}
		// AUTHOR_MESSAGE is what the person sees and signs elsewhere: the
		// callback may take as long as she needs.
		msg := AuthorMessage(sc.ControlCommit, sc.HeadDigest, SignersDigest(AlgCMS, list))
		der, err := o.CMSSigner.Sign(msg)
		if err != nil {
			return nil, fmt.Errorf("capsule: signing: %w", err)
		}
		if signature, err = EncodeAuthorSignature(AlgCMS, list, der); err != nil {
			return nil, err
		}
		wantSig = VerdictSignedComplete
	}
	wantSeal := VerdictNoSeal
	if o.Sealer != nil {
		subject := SealSubject(sc.ControlCommit, sc.HeadDigest, SigPart(signature))
		token, err := o.Sealer.Seal(subject)
		if err != nil {
			return nil, fmt.Errorf("capsule: sealing: %w", err)
		}
		if seal, err = EncodeSeal(SealTypeRFC3161, token); err != nil {
			return nil, err
		}
		wantSeal = VerdictSealed
	}
	security, err := EncodeSecurityWith(signature, seal)
	if err != nil {
		return nil, err
	}
	v := EvaluateSecurityIn(security, sc)
	// A seal that proves nothing before the round time (S5) is still a seal
	// that verifies: the writer's clock and the authority's may differ, or
	// the token may carry no accuracy. Result.Security lets the caller warn
	// of it (§62.1 rule 19).
	sealOK := v.Seal == wantSeal || wantSeal == VerdictSealed && v.Seal == VerdictSealedLate
	if v.Signature != wantSig || !sealOK || v.AuthorKey != key {
		return nil, fmt.Errorf("capsule: self-check: the reader finds the verdicts %s and %s in this security area, not %s and %s%s", v.Signature, v.Seal, wantSig, wantSeal, detailText(v.Detail))
	}
	s.verdicts = v
	return security, nil
}

// detailText names the signers that failed, for the error of a writer.
func detailText(d *Detail) string {
	if d == nil {
		return ""
	}
	var parts []string
	for _, l := range d.Signers {
		if l.Result != "valid" {
			parts = append(parts, l.Holder+": "+l.Result)
		}
	}
	if len(parts) == 0 {
		return ""
	}
	return " (" + strings.Join(parts, "; ") + ")"
}

// newHead checks the files and the texts of opts with the rules of spec
// §29.4 to §29.6, in the words of a writer (spec §62.1 rule 15), and returns
// the head with the files in the byte order of their paths, their layout and
// mtime, and a zero salt and zero SHA-256; order[i] is the source of entry
// i. The comment has its CR LF, and any lone CR, turned into LF (§29.6).
func newHead(sources []Source, opts EncryptOptions) (*Head, []int, error) {
	comment := strings.ReplaceAll(strings.ReplaceAll(opts.Comment, "\r\n", "\n"), "\r", "\n")
	switch {
	case len(sources) == 0 && comment == "":
		return nil, nil, errors.New("capsule: a format 3 capsule holds at least one file or a comment (spec §62.1 rule 14)")
	case len(sources) > MaxFiles:
		return nil, nil, fmt.Errorf("capsule: %d files, more than %d", len(sources), MaxFiles)
	}
	if err := checkHeadText("comment", comment, MaxCommentLen, pathrule.CheckComment); err != nil {
		return nil, nil, err
	}
	if err := checkHeadText("declared author", opts.Author, MaxAuthorLen, pathrule.CheckAuthor); err != nil {
		return nil, nil, err
	}

	order := make([]int, len(sources))
	for i := range order {
		order[i] = i
	}
	// R8: the byte order of the paths, which is the order of Go strings.
	slices.SortStableFunc(order, func(a, b int) int { return strings.Compare(sources[a].Path, sources[b].Path) })
	h := &Head{Comment: comment, Author: opts.Author, Critical: opts.HeadCritical, Noncritical: opts.HeadNoncritical}
	paths := make([]string, len(order))
	var end uint64
	for i, j := range order {
		src, p := sources[j], sources[j].Path
		switch {
		case i > 0 && p == paths[i-1]:
			return nil, nil, fmt.Errorf("capsule: path %q given twice", p)
		case !utf8.ValidString(p):
			return nil, nil, fmt.Errorf("capsule: path %q: R1: not valid UTF-8", p)
		case len(p) == 0 || len(p) > MaxPathLen:
			return nil, nil, fmt.Errorf("capsule: path %q: R1: %d bytes, not 1 to %d", p, len(p), MaxPathLen)
		case src.Size < 0:
			return nil, nil, fmt.Errorf("capsule: file %q: negative size %d", p, src.Size)
		case src.Open == nil:
			return nil, nil, fmt.Errorf("capsule: file %q: Source.Open is nil", p)
		case uint64(src.Size) > MaxPayloadLength-end:
			return nil, nil, fmt.Errorf("capsule: the files add up to more than %d bytes, the maximum of L", uint64(MaxPayloadLength))
		}
		if err := pathrule.CheckPath(p); err != nil {
			return nil, nil, fmt.Errorf("capsule: path %q: %w", p, err)
		}
		f := File{Path: p, Size: uint64(src.Size), Start: end, End: end + uint64(src.Size)}
		if t := src.ModTime; !t.IsZero() {
			if u := t.Unix(); u >= 0 && u <= MaxMTime {
				f.MTime, f.HasMTime = uint64(u), true
			}
		}
		h.Files = append(h.Files, f)
		paths[i], end = p, f.End
	}
	if err := pathrule.CheckTree(paths); err != nil {
		var e *pathrule.Error
		if errors.As(err, &e) && e.Paths[0] > 0 {
			return nil, nil, fmt.Errorf("capsule: paths %q and %q: %w", paths[e.Paths[1]-1], paths[e.Paths[0]-1], err)
		}
		return nil, nil, fmt.Errorf("capsule: paths: %w", err)
	}
	return h, order, nil
}

// checkHeadText checks the comment or the declared author, when present:
// valid UTF-8, at most max bytes, and the characters of spec §29.6.
func checkHeadText(what, s string, max int, check func(string) error) error {
	switch {
	case s == "":
		return nil
	case !utf8.ValidString(s):
		return fmt.Errorf("capsule: %s: not valid UTF-8", what)
	case len(s) > max:
		return fmt.Errorf("capsule: %s: %d bytes, more than %d", what, len(s), max)
	}
	if err := check(s); err != nil {
		return fmt.Errorf("capsule: %s: %w", what, err)
	}
	return nil
}

// selfCheckHead decodes HEAD_CBOR with the rules of the reader, but for the
// knowledge of its critical extensions, which depends on the reader, as
// selfCheckControl does with the control (spec §62.1 rule 17). A head that
// the reader rejects would only be found after the date.
func selfCheckHead(b []byte) error {
	if _, err := decodeHead(b); err != nil {
		return fmt.Errorf("capsule: self-check: the reader rejects this head: %w", err)
	}
	return nil
}

// readSource reads the file of src, which must be exactly size bytes, and
// returns its SHA-256; w, when not nil, receives its bytes. In the second
// reading, a file whose size differs has changed (spec §62.1 rule 18).
func readSource(w io.Writer, src Source, size uint64, second bool) ([32]byte, error) {
	var sum [32]byte
	mismatch := func(format string, args ...any) error {
		if second {
			return fmt.Errorf("capsule: file %q changed after its first reading: %s", src.Path, fmt.Sprintf(format, args...))
		}
		return fmt.Errorf("capsule: file %q: %s", src.Path, fmt.Sprintf(format, args...))
	}
	rc, err := src.Open()
	if err != nil {
		return sum, fmt.Errorf("capsule: file %q: %w", src.Path, err)
	}
	defer rc.Close()
	h := sha256.New()
	buf := make([]byte, 32<<10)
	defer clear(buf)
	var n uint64
	for {
		k, err := rc.Read(buf)
		if uint64(k) > size-n {
			return sum, mismatch("more than its size of %d bytes", size)
		}
		h.Write(buf[:k])
		if w != nil && k > 0 {
			if _, err := w.Write(buf[:k]); err != nil {
				return sum, err
			}
		}
		n += uint64(k)
		if err == io.EOF {
			break
		}
		if err != nil {
			return sum, fmt.Errorf("capsule: file %q: %w", src.Path, err)
		}
	}
	if n != size {
		return sum, mismatch("%d bytes, not its size of %d", n, size)
	}
	h.Sum(sum[:0])
	return sum, nil
}