// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tdx

import (
	"bytes"
	"crypto/sha512"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestReplayEventLogOVMF(t *testing.T) {
	// Load OVMF eventlog data
	testDataPath := filepath.Join("testdata", "quote", "tdx_CCEL_data_ovmf")
	eventLogData, err := os.ReadFile(testDataPath)
	require.NoError(t, err, "failed to read OVMF eventlog data: %s", testDataPath)

	// Replay eventlog to calculate RTMRs
	rtmrs, err := ReplayEventLog(eventLogData)

	// Verify replay succeeded
	require.NoError(t, err, "replay should succeed with valid OVMF eventlog data")
	require.NotNil(t, rtmrs, "rtmrs should not be nil when replay succeeds")

	// Verify RTMR structure
	assert.Len(t, rtmrs, 4, "should have 4 RTMR values")
	for i, rtmr := range rtmrs {
		assert.Len(t, rtmr, 48, "RTMR[%d] should be 48 bytes", i)
	}

	t.Logf("OVMF eventlog replay succeeded")
}

func TestReplayEventLogGn8vTee(t *testing.T) {
	// Load eventlog data
	testDataPath := filepath.Join("testdata", "quote", "tdx_CCEL_data")
	eventLogData, err := os.ReadFile(testDataPath)
	require.NoError(t, err, "failed to read eventlog data: %s", testDataPath)

	// Replay eventlog to calculate RTMRs
	rtmrs, err := ReplayEventLog(eventLogData)

	// Verify replay succeeded
	require.NoError(t, err, "replay should succeed with valid eventlog data")
	require.NotNil(t, rtmrs, "rtmrs should not be nil when replay succeeds")

	// Verify RTMR structure
	assert.Len(t, rtmrs, 4, "should have 4 RTMR values")
	for i, rtmr := range rtmrs {
		assert.Len(t, rtmr, 48, "RTMR[%d] should be 48 bytes", i)
	}

	t.Logf("eventlog replay succeeded")
}

func TestReplayEventLogUKI(t *testing.T) {
	// Load UKI eventlog data
	testDataPath := filepath.Join("testdata", "quote", "tdx_CCEL_data_uki")
	eventLogData, err := os.ReadFile(testDataPath)
	require.NoError(t, err, "failed to read UKI eventlog data: %s", testDataPath)

	// Replay eventlog to calculate RTMRs
	rtmrs, err := ReplayEventLog(eventLogData)

	// Verify replay succeeded
	require.NoError(t, err, "replay should succeed with valid UKI eventlog data")
	require.NotNil(t, rtmrs, "rtmrs should not be nil when replay succeeds")

	// Verify RTMR structure
	assert.Len(t, rtmrs, 4, "should have 4 RTMR values")
	for i, rtmr := range rtmrs {
		assert.Len(t, rtmr, 48, "RTMR[%d] should be 48 bytes", i)
	}

	t.Logf("UKI eventlog replay succeeded")
}

func TestReplayEventLogPai(t *testing.T) {
	testDataPath := filepath.Join("testdata", "quote", "tdx_CCEL_data_pai")
	eventLogData, err := os.ReadFile(testDataPath)
	require.NoError(t, err, "failed to read PAI eventlog data: %s", testDataPath)

	rtmrs, err := ReplayEventLog(eventLogData)
	require.NoError(t, err, "replay should succeed with valid PAI eventlog data")
	require.NotNil(t, rtmrs, "rtmrs should not be nil when replay succeeds")

	assert.Len(t, rtmrs, 4, "should have 4 RTMR values")
	for i, rtmr := range rtmrs {
		assert.Len(t, rtmr, 48, "RTMR[%d] should be 48 bytes", i)
	}

	t.Logf("PAI eventlog replay succeeded")
}

// TestDumpEventLogPai walks every entry of the PAI eventlog and prints
// per-event details + RTMR extend trace.
// Run with `-v`. Useful for eyeballing event types, AAEL payload format,
// and verifying per-event digests alongside the final RTMR.
func TestDumpEventLogPai(t *testing.T) {
	testDataPath := filepath.Join("testdata", "quote", "tdx_CCEL_data_pai")
	data, err := os.ReadFile(testDataPath)
	require.NoError(t, err)

	tdEventLog, err := NewTDEventLog(data)
	require.NoError(t, err)
	require.NotNil(t, tdEventLog)

	// Set EVENT_DUMP_FILE=<path> to write the dump to that path (relative
	// or absolute, in the test's working directory) without testing.T
	// line-number prefixes. Without the env var, output goes through
	// t.Logf and is only visible with `go test -v`.
	logf := t.Logf
	if path := os.Getenv("EVENT_DUMP_FILE"); path != "" {
		f, err := os.Create(path)
		require.NoError(t, err)
		defer f.Close()
		logf = func(format string, a ...any) {
			fmt.Fprintf(f, format+"\n", a...)
		}
		t.Logf("dumping to %s", path)
	}

	logf("========== PAI EVENTLOG REPLAY DETAILS ==========")
	logf("Parsed successfully: %d events found", len(tdEventLog.Entries))

	var rtmrs [TDX_RTMR_COUNT][TDX_RTMR_LENGTH_BY_BYTES]byte
	logf("Initial RTMR values (all zeros):")
	for i := range rtmrs {
		logf("  RTMR%d: %x", i, rtmrs[i])
	}

	for i, e := range tdEventLog.Entries {
		rtmrIdx := int(e.TDRegisterIndex) - 1

		logf("")
		logf("[EVENT #%d]", i)
		logf("  Event Type:  0x%x (%s)", e.EventType, paiEventTypeName(e.EventType))
		logf("  Target RTMR: RTMR%d", rtmrIdx)
		logf("  Event Size:  %d bytes", e.EventSize)
		logf("  Event Data (HEX):   %s", hex.EncodeToString(e.Event))
		logf("  Event Data (ASCII): %s", asciiSafe(e.Event))

		recorded := e.Digests[0]
		logf("  Digest (recorded):  %x", recorded)

		computed, ok := paiComputeEventDigest(e)
		if ok {
			logf("  Digest (computed):  %x", computed)
			if bytes.Equal(computed, recorded) {
				logf("  [VERIFY] Digest MATCH")
			} else {
				logf("  [VERIFY] Digest MISMATCH")
			}
		} else {
			logf("  [VERIFY] SKIP (no compute logic for this event type)")
		}

		// Extend RTMR: rtmr = sha384(rtmr || digest)
		before := rtmrs[rtmrIdx]
		h := sha512.New384()
		h.Write(before[:])
		h.Write(recorded)
		copy(rtmrs[rtmrIdx][:], h.Sum(nil))
		logf("  RTMR%d BEFORE: %x", rtmrIdx, before)
		logf("  RTMR%d AFTER:  %x", rtmrIdx, rtmrs[rtmrIdx])

		if e.EventType == EV_EVENT_TAG && len(e.Event) > AAEL_TAGGED_EVENT_HEADER_SIZE {
			payload := e.Event[AAEL_TAGGED_EVENT_HEADER_SIZE:]
			logf("  AAEL payload:       %s",
				strings.TrimRight(string(payload), "\x00"))
		}
	}

	logf("")
	logf("========== FINAL RTMRs ==========")
	for i := range rtmrs {
		logf("  RTMR%d: %x", i, rtmrs[i])
	}

	// Cross-check against package Replay path.
	got, err := ReplayEventLog(data)
	require.NoError(t, err)
	for i := range got {
		assert.Equal(t, rtmrs[i], got[i], "RTMR%d divergence between local and ReplayEventLog", i)
	}

	containers, err := ExtractAAELContainerEvents(data)
	require.NoError(t, err)
	logf("")
	logf("AAEL container/image-pull events: %d", len(containers))
	for i, c := range containers {
		logf("  [%d] name=%s digest=%s", i, c.Name, c.Digest)
	}
}

func paiComputeEventDigest(e *TDEventLogEntry) ([]byte, bool) {
	switch e.EventType {
	case EV_EFI_VARIABLE_BOOT:
		v := NewVariableLog(e.Event)
		h := sha512.New384()
		h.Write(v.VariableData)
		return h.Sum(nil), true
	case EV_EFI_VARIABLE_DRIVER_CONFIG, EV_SEPARATOR, EV_EFI_ACTION,
		EV_EFI_GPT_EVENT, EV_EFI_VARIABLE_AUTHORITY:
		h := sha512.New384()
		h.Write(e.Event)
		return h.Sum(nil), true
	default:
		return nil, false
	}
}

func asciiSafe(b []byte) string {
	var sb strings.Builder
	for _, c := range b {
		if c >= 0x20 && c < 0x7f {
			sb.WriteByte(c)
		} else {
			sb.WriteByte('.')
		}
	}
	return sb.String()
}

func paiEventTypeName(t uint32) string {
	switch t {
	case EV_PREBOOT_CERT:
		return "EV_PREBOOT_CERT"
	case EV_POST_CODE:
		return "EV_POST_CODE"
	case EV_NO_ACTION:
		return "EV_NO_ACTION"
	case EV_SEPARATOR:
		return "EV_SEPARATOR"
	case EV_ACTION:
		return "EV_ACTION"
	case EV_EVENT_TAG:
		return "EV_EVENT_TAG"
	case EV_PLATFORM_CONFIG_FLAGS:
		return "EV_PLATFORM_CONFIG_FLAGS"
	case EV_IPL:
		return "EV_IPL"
	case EV_EFI_VARIABLE_DRIVER_CONFIG:
		return "EV_EFI_VARIABLE_DRIVER_CONFIG"
	case EV_EFI_VARIABLE_BOOT:
		return "EV_EFI_VARIABLE_BOOT"
	case EV_EFI_BOOT_SERVICES_APPLICATION:
		return "EV_EFI_BOOT_SERVICES_APPLICATION"
	case EV_EFI_BOOT_SERVICES_DRIVER:
		return "EV_EFI_BOOT_SERVICES_DRIVER"
	case EV_EFI_RUNTIME_SERVICES_DRIVER:
		return "EV_EFI_RUNTIME_SERVICES_DRIVER"
	case EV_EFI_GPT_EVENT:
		return "EV_EFI_GPT_EVENT"
	case EV_EFI_ACTION:
		return "EV_EFI_ACTION"
	case EV_EFI_PLATFORM_FIRMWARE_BLOB2:
		return "EV_EFI_PLATFORM_FIRMWARE_BLOB2"
	case EV_EFI_HANDOFF_TABLES2:
		return "EV_EFI_HANDOFF_TABLES2"
	case EV_EFI_VARIABLE_AUTHORITY:
		return "EV_EFI_VARIABLE_AUTHORITY"
	default:
		return fmt.Sprintf("UNKNOWN(0x%X)", t)
	}
}

func TestExtractKernelDigestPai(t *testing.T) {
	data, err := os.ReadFile(filepath.Join("testdata", "quote", "tdx_CCEL_data_pai"))
	require.NoError(t, err)

	got, err := ExtractKernelDigest(data)
	require.NoError(t, err)
	// Sourced from TestDumpEventLogPai output (event #3, td_payload).
	const want = "85619a9146b5c1409bf034bdabd6d93df26a5a23b41b43f92a7a4025760a1c4f110194ae75a2b9625b0de97f2b00fc8f"
	assert.Equal(t, want, got)
}

func TestExtractKernelCmdlineDigestPai(t *testing.T) {
	data, err := os.ReadFile(filepath.Join("testdata", "quote", "tdx_CCEL_data_pai"))
	require.NoError(t, err)

	got, err := ExtractKernelCmdlineDigest(data)
	require.NoError(t, err)
	// Sourced from TestDumpEventLogPai output (event #4, td_payload_info).
	const want = "b3af693977dff21afd58c649a454b550d20c71ac2f02188217817e1a63dd5c2aedb671665d44ffa24487048c226dedb2"
	assert.Equal(t, want, got)
}

func TestExtractAAELContainerEventsPai(t *testing.T) {
	data, err := os.ReadFile(filepath.Join("testdata", "quote", "tdx_CCEL_data_pai"))
	require.NoError(t, err)

	events, err := ExtractAAELContainerEvents(data)
	require.NoError(t, err)
	require.Len(t, events, 1, "expected one kangaroo/pull-image event in PAI eventlog")

	e := events[0]
	assert.Empty(t, e.Name, "PAI form should not populate Name")
	assert.Equal(t, "pai-registry.cn-wulanchabu-acdr-1.cr.aliyuncs.com/default/busybox:2025052701", e.Reference)
	assert.Equal(t, "sha256:16ece118a36d152ed7562347741470b1bb0b0f7561ff1fe0049857a03d4fec5d", e.Digest)
}
