package capsule_test

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"os"
	"path/filepath"
	"reflect"
	"testing"
	"time"

	"filippo.io/age"

	"github.com/datekeys/datekeys-go/accesskey"
	"github.com/datekeys/datekeys-go/agewrap"
	"github.com/datekeys/datekeys-go/capsule"
	"github.com/datekeys/datekeys-go/internal/testkit"
	"github.com/datekeys/datekeys-go/profile"
	"github.com/datekeys/datekeys-go/provider"
)

const fixtureDir = "../testdata/fixtures"

var fixtureNames = []string{"time_only", "time_only_extensions", "time_and_key_portable", "time_and_key_recipients", "empty_payload"}

type fixture struct {
	testkit.DKCFixture
	dkc       []byte
	plaintext []byte
	release   provider.Release
	dkk       *accesskey.AccessKey
	ids       []age.Identity
}

func loadFixture(t testing.TB, name string) *fixture {
	t.Helper()
	f := &fixture{}
	if err := testkit.ReadJSON(filepath.Join(fixtureDir, name+".json"), &f.DKCFixture); err != nil {
		t.Fatal(err)
	}
	var err error
	if f.dkc, err = os.ReadFile(filepath.Join(fixtureDir, f.File)); err != nil {
		t.Fatal(err)
	}
	if f.plaintext, err = os.ReadFile(filepath.Join(fixtureDir, f.PlaintextFile)); err != nil {
		t.Fatal(err)
	}
	sig, err := hex.DecodeString(f.Release.Signature)
	if err != nil {
		t.Fatal(err)
	}
	f.release = provider.Release{Round: f.Release.Round, Signature: sig}
	if f.AccessKeyFile != "" {
		b, err := os.ReadFile(filepath.Join(fixtureDir, f.AccessKeyFile))
		if err != nil {
			t.Fatal(err)
		}
		if f.dkk, err = accesskey.Decode(bytes.NewReader(b)); err != nil {
			t.Fatal(err)
		}
	}
	for _, s := range f.Identities {
		id, err := age.ParseX25519Identity(s)
		if err != nil {
			t.Fatal(err)
		}
		f.ids = append(f.ids, id)
	}
	return f
}

func (f *fixture) unlock(t testing.TB) time.Time {
	u, err := time.Parse(time.RFC3339, f.UnlockAt)
	if err != nil {
		t.Fatal(err)
	}
	return u
}

func (f *fixture) openOptions(t testing.TB) capsule.OpenOptions {
	o := capsule.OpenOptions{Registry: testkit.Registry(), Source: testkit.NewSource(f.release), Now: testkit.Fixed(f.unlock(t))}
	if f.AccessPolicy == "time_and_key" {
		o.AccessKey = f.dkk
	}
	return o
}

func decryptAge(t *testing.T, file []byte, id age.Identity) []byte {
	t.Helper()
	r, err := age.Decrypt(bytes.NewReader(file), id)
	if err != nil {
		t.Fatal(err)
	}
	b, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func stanzaList(in []capsule.StanzaInfo) []testkit.FixtureStanza {
	out := make([]testkit.FixtureStanza, len(in))
	for i, s := range in {
		out[i] = testkit.FixtureStanza{Type: s.Type, Args: s.Args}
	}
	return out
}

// Spec §67: conformance is shown by decrypting and verifying each official
// fixture and comparing every intermediate value and the plaintext.
func TestConformanceFixtures(t *testing.T) {
	p := profile.Quicknet()
	for _, name := range fixtureNames {
		t.Run(name, func(t *testing.T) {
			f := loadFixture(t, name)
			if sum := sha256.Sum256(f.dkc); hex.EncodeToString(sum[:]) != f.SHA256 {
				t.Fatal("fixture bytes changed")
			}
			if sum := sha256.Sum256(f.plaintext); hex.EncodeToString(sum[:]) != f.PlaintextSHA256 {
				t.Fatal("plaintext file changed")
			}
			if err := provider.Verify(p, provider.Condition{Round: f.release.Round}, f.release); err != nil {
				t.Fatalf("embedded release: %v", err)
			}

			// PRELUDE, PUBLIC_HEADER and the pre-unlock view (steps 1 to 8).
			parts, err := testkit.Split(f.dkc)
			if err != nil {
				t.Fatal(err)
			}
			if hex.EncodeToString(parts.Prelude) != f.Prelude || hex.EncodeToString(parts.Header) != f.PublicHeader {
				t.Fatal("PRELUDE or PUBLIC_HEADER differ")
			}
			in, err := capsule.Inspect(bytes.NewReader(f.dkc), capsule.InspectOptions{Registry: testkit.Registry()})
			if err != nil {
				t.Fatal(err)
			}
			if in.Header.DateKey.Compact() != f.DateKey || in.Header.CapsuleIDHex() != f.CapsuleID ||
				in.Header.Policy.String() != f.AccessPolicy || in.UnlockAt.Format(time.RFC3339) != f.UnlockAt {
				t.Fatalf("inspection differs: %+v", in.Header)
			}
			if !reflect.DeepEqual(stanzaList(in.OuterStanzas), f.OuterStanzas) || !reflect.DeepEqual(stanzaList(in.PayloadStanzas), f.PayloadStanzas) {
				t.Fatal("stanzas differ")
			}
			for _, c := range in.Checks {
				if !c.OK || c.Step > 8 {
					t.Fatalf("unexpected inspection check %+v", c)
				}
			}
			var pre [capsule.PreludeSize]byte
			copy(pre[:], parts.Prelude)
			if b := capsule.HeaderBinding(pre, parts.Header); hex.EncodeToString(b[:]) != f.HeaderBinding {
				t.Fatal("header_binding differs")
			}
			h, err := capsule.DecodeHeader(parts.Header)
			if err != nil {
				t.Fatal(err)
			}
			if re, _ := capsule.EncodeHeader(h); !bytes.Equal(re, parts.Header) {
				t.Fatal("EncodeHeader(DecodeHeader(x)) != x")
			}
			if len(h.Critical)+len(h.Noncritical) != len(f.HeaderExtensions) {
				t.Fatal("header extensions differ")
			}

			// Opening layer by layer: OUTER_TIME_AGE, INNER_ACCESS_AGE, CONTROL_CBOR.
			timeID, _ := agewrap.NewTimeIdentity(p, f.release.Round, f.release)
			inner := decryptAge(t, parts.Sealed, timeID)
			control := inner
			if f.Structure == "time_and_key" {
				st, err := agewrap.Stanzas(bytes.NewReader(inner))
				if err != nil {
					t.Fatal(err)
				}
				got := make([]testkit.FixtureStanza, len(st))
				for i, s := range st {
					got[i] = testkit.FixtureStanza{Type: s.Type, Args: s.Args}
				}
				if !reflect.DeepEqual(got, f.InnerStanzas) {
					t.Fatal("INNER_ACCESS_AGE stanzas differ")
				}
				kid, err := f.dkk.Identity()
				if err != nil {
					t.Fatal(err)
				}
				acc, _ := agewrap.NewAccessIdentity(kid)
				control = decryptAge(t, inner, acc)
			}
			if hex.EncodeToString(control) != f.ControlCBOR {
				t.Fatal("CONTROL_CBOR differs")
			}
			ctrl, err := capsule.DecodeControl(control)
			if err != nil {
				t.Fatal(err)
			}
			if hex.EncodeToString(ctrl.PayloadIdentity[:]) != f.PayloadIdentity || hex.EncodeToString(ctrl.HeaderBinding[:]) != f.HeaderBinding {
				t.Fatal("I_PAYLOAD or header_binding in CONTROL_CBOR differ")
			}
			if re, _ := capsule.EncodeControl(ctrl); !bytes.Equal(re, control) {
				t.Fatal("EncodeControl(DecodeControl(x)) != x")
			}
			if len(ctrl.Critical)+len(ctrl.Noncritical) != len(f.ControlExt) {
				t.Fatal("control extensions differ")
			}
			payloadID, _ := agewrap.NewPayloadIdentity(ctrl.PayloadIdentity[:])
			if got := decryptAge(t, parts.Payload, payloadID); !bytes.Equal(got, f.plaintext) {
				t.Fatal("PAYLOAD_AGE plaintext differs")
			}

			// The complete flow through the public API, stage by stage.
			var out bytes.Buffer
			opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), f.openOptions(t))
			if err != nil {
				t.Fatal(err)
			}
			if !bytes.Equal(out.Bytes(), f.plaintext) {
				t.Fatal("plaintext differs")
			}
			var stages []testkit.FixtureStage
			for _, c := range opened.Inspection.Checks {
				stages = append(stages, testkit.FixtureStage{Step: c.Step, Name: c.Name, OK: c.OK, Error: c.Error})
			}
			if !reflect.DeepEqual(stages, f.Stages) {
				t.Fatalf("stages differ:\n got %+v\nwant %+v", stages, f.Stages)
			}
			if len(opened.ControlNoncritical) != len(ctrl.Noncritical) {
				t.Fatal("Open does not report the control extensions")
			}
		})
	}
}

// Every known recipient of a multi-recipient fixture opens it on its own.
func TestFixtureRecipients(t *testing.T) {
	f := loadFixture(t, "time_and_key_recipients")
	if len(f.ids) != 2 {
		t.Fatal("fixture should have two known recipients")
	}
	for i, id := range f.ids {
		o := f.openOptions(t)
		o.AccessKey = nil
		o.Identities = []age.Identity{id}
		var out bytes.Buffer
		if _, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), o); err != nil || !bytes.Equal(out.Bytes(), f.plaintext) {
			t.Fatalf("recipient %d: %v", i, err)
		}
	}
}

// A non-seekable reader works too: only the capsule_digest shortcut is skipped.
func TestOpenFromPlainReader(t *testing.T) {
	for _, name := range []string{"time_only", "time_and_key_portable"} {
		f := loadFixture(t, name)
		var out bytes.Buffer
		r := struct{ io.Reader }{bytes.NewReader(f.dkc)}
		if _, err := capsule.Open(context.Background(), &out, r, f.openOptions(t)); err != nil || !bytes.Equal(out.Bytes(), f.plaintext) {
			t.Fatalf("%s: %v", name, err)
		}
	}
}