// 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 sdk

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strconv"

	"github.com/dashscope/dashscope-confidential-maas/pkg/attestation"
	"github.com/dashscope/dashscope-confidential-maas/pkg/tdx"
)

// teeVerifier verifies TEE evidence and validates the server public key binding.
type teeVerifier interface {
	VerifyEvidence(evidence json.RawMessage, serverPubKeyPEM string) error
	// ExtractDetails returns TEE-type-specific key fields as JSON (e.g. MRTD, RTMR for TDX).
	ExtractDetails(evidence json.RawMessage) (json.RawMessage, error)
}

// supportedTEEVerifiers is the package-level registry of SDK-supported TEE verifiers.
var supportedTEEVerifiers = map[string]teeVerifier{
	"TDX": &tdxVerifier{},
}

// TEEDetailsJSON returns TEE-type-specific key fields as JSON for inspection or logging.
// For TDX this includes mrtd, rtmr, td_attributes, debug_enabled.
// The transparency report uses the same ExtractDetails path to enrich stored evidence.
func (r *AttestationResult) TEEDetailsJSON() (json.RawMessage, error) {
	if r == nil {
		return nil, fmt.Errorf("nil attestation result")
	}
	v, ok := supportedTEEVerifiers[r.Bundle.TEEType]
	if !ok {
		return nil, fmt.Errorf("unsupported TEE type: %s", r.Bundle.TEEType)
	}
	return v.ExtractDetails(r.Bundle.Evidence)
}

// --- tdxVerifier: teeVerifier implementation for Intel TDX ---

type tdxVerifier struct{}

func (v *tdxVerifier) VerifyEvidence(evidence json.RawMessage, serverPubKeyPEM string) error {
	tdxEvidence, parsedQuote, quoteBytes, err := v.parseEvidence(evidence)
	if err != nil {
		return err
	}

	if _, err := v.verifyQuote(parsedQuote, quoteBytes, serverPubKeyPEM); err != nil {
		return fmt.Errorf("failed to verify TDX quote: %w", err)
	}

	if tdxEvidence.CCEventLog == "" {
		return fmt.Errorf("eventlog is required but not present in TDX evidence")
	}
	eventLog, err := base64.StdEncoding.DecodeString(tdxEvidence.CCEventLog)
	if err != nil {
		return fmt.Errorf("failed to decode eventlog: %w", err)
	}
	return v.replayEventLog(parsedQuote, eventLog)
}

// parseEvidence unmarshals evidence JSON, decodes the base64 quote and parses it.
// Returns the structured evidence, parsed quote, and raw quote bytes (needed by VerifyQuote).
func (v *tdxVerifier) parseEvidence(evidence json.RawMessage) (*attestation.TdxEvidence, *tdx.TdxEcdsaQuote, []byte, error) {
	var tdxEvidence attestation.TdxEvidence
	if err := json.Unmarshal(evidence, &tdxEvidence); err != nil {
		return nil, nil, nil, fmt.Errorf("failed to parse TDX evidence: %w", err)
	}
	if tdxEvidence.Quote == "" {
		return nil, nil, nil, fmt.Errorf("missing quote in TDX evidence")
	}
	quoteBytes, err := base64.StdEncoding.DecodeString(tdxEvidence.Quote)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to decode TDX quote: %w", err)
	}
	parsedQuote, err := tdx.ParseQuote(quoteBytes)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to parse TDX quote: %w", err)
	}
	return &tdxEvidence, parsedQuote, quoteBytes, nil
}

// verifyQuote performs cryptographic verification using raw quote bytes
// and validates the report data binds to the expected server public key.
func (v *tdxVerifier) verifyQuote(parsedQuote *tdx.TdxEcdsaQuote, quoteBytes []byte, expectedServerPubKeyPEM string) (*tdx.VerificationResult, error) {
	verifier := tdx.NewQuoteVerifier("")
	verifyResult, errType := verifier.VerifyQuote(quoteBytes)
	if errType != tdx.SGX_QL_SUCCESS {
		return nil, fmt.Errorf("quote verification failed: %v", errType)
	}
	if verifyResult == nil {
		return nil, fmt.Errorf("quote verification returned nil result")
	}

	unacceptableResults := []tdx.QvResultType{
		tdx.SGX_QL_QV_RESULT_INVALID_SIGNATURE,
		tdx.SGX_QL_QV_RESULT_REVOKED,
		tdx.SGX_QL_QV_RESULT_UNSPECIFIED,
	}
	for _, unacceptable := range unacceptableResults {
		if verifyResult.QvResult == unacceptable {
			return nil, fmt.Errorf("quote verification failed with QV result: 0x%x", verifyResult.QvResult)
		}
	}

	if err := v.validateReportData(parsedQuote, expectedServerPubKeyPEM); err != nil {
		return nil, fmt.Errorf("report data validation failed: %w", err)
	}

	return verifyResult, nil
}

func (v *tdxVerifier) validateReportData(parsedQuote *tdx.TdxEcdsaQuote, expectedServerPubKeyPEM string) error {
	expectedPubKeyHash := sha256.Sum256([]byte(expectedServerPubKeyPEM))

	var reportData []byte
	switch report := parsedQuote.Report.(type) {
	case tdx.EnclaveReportBody:
		reportData = report.ReportData[:]
	case tdx.TdReport10:
		reportData = report.ReportData[:]
	case tdx.TdReport15:
		reportData = report.ReportData[:]
	default:
		return fmt.Errorf("unsupported report type: %T", report)
	}

	if len(reportData) != 64 {
		return fmt.Errorf("invalid report data length: expected 64 bytes, got %d", len(reportData))
	}

	hashHex := fmt.Sprintf("%x", expectedPubKeyHash)
	if !bytes.Equal(reportData[:64], []byte(hashHex)) {
		return fmt.Errorf("quote report data mismatch: public key hash does not match")
	}

	return nil
}

func (v *tdxVerifier) replayEventLog(parsedQuote *tdx.TdxEcdsaQuote, eventLogData []byte) error {
	tdEventLog, err := tdx.NewTDEventLog(eventLogData)
	if err != nil {
		return fmt.Errorf("failed to parse eventlog: %w", err)
	}

	// Replay eventlog (includes digest verification)
	replayedRTMRs, err := tdEventLog.Replay()
	if err != nil {
		return fmt.Errorf("eventlog replay failed: %w", err)
	}

	// Extract RTMRs from quote
	var quoteRTMRs [4][48]byte
	switch report := parsedQuote.Report.(type) {
	case tdx.TdReport10:
		copy(quoteRTMRs[0][:], report.Rtmr0[:])
		copy(quoteRTMRs[1][:], report.Rtmr1[:])
		copy(quoteRTMRs[2][:], report.Rtmr2[:])
		copy(quoteRTMRs[3][:], report.Rtmr3[:])
	case tdx.TdReport15:
		copy(quoteRTMRs[0][:], report.Rtmr0[:])
		copy(quoteRTMRs[1][:], report.Rtmr1[:])
		copy(quoteRTMRs[2][:], report.Rtmr2[:])
		copy(quoteRTMRs[3][:], report.Rtmr3[:])
	default:
		return fmt.Errorf("unsupported report type for RTMR comparison: %T", report)
	}

	// Compare RTMRs
	var mismatchedRTMRs []int
	for i := 0; i < 4; i++ {
		if !bytes.Equal(quoteRTMRs[i][:], replayedRTMRs[i][:]) {
			mismatchedRTMRs = append(mismatchedRTMRs, i)
		}
	}

	if len(mismatchedRTMRs) > 0 {
		return fmt.Errorf("RTMR mismatch: RTMRs %v do not match between Quote and EventLog", mismatchedRTMRs)
	}

	return nil
}

func (v *tdxVerifier) ExtractDetails(evidence json.RawMessage) (json.RawMessage, error) {
	_, parsedQuote, _, err := v.parseEvidence(evidence)
	if err != nil {
		return nil, err
	}
	var details map[string]any
	switch r := parsedQuote.Report.(type) {
	case tdx.TdReport10:
		tdAttr := fmt.Sprintf("%x", r.TdAttributes[:])
		details = map[string]any{
			"mrtd":          fmt.Sprintf("%x", r.MrTd[:]),
			"rtmr":          []string{fmt.Sprintf("%x", r.Rtmr0[:]), fmt.Sprintf("%x", r.Rtmr1[:]), fmt.Sprintf("%x", r.Rtmr2[:]), fmt.Sprintf("%x", r.Rtmr3[:])},
			"td_attributes": tdAttr,
			"debug_enabled": debugFromTDAttributesHex(tdAttr),
		}
	case tdx.TdReport15:
		tdAttr := fmt.Sprintf("%x", r.TdAttributes[:])
		details = map[string]any{
			"mrtd":          fmt.Sprintf("%x", r.MrTd[:]),
			"rtmr":          []string{fmt.Sprintf("%x", r.Rtmr0[:]), fmt.Sprintf("%x", r.Rtmr1[:]), fmt.Sprintf("%x", r.Rtmr2[:]), fmt.Sprintf("%x", r.Rtmr3[:])},
			"td_attributes": tdAttr,
			"debug_enabled": debugFromTDAttributesHex(tdAttr),
		}
	default:
		details = map[string]any{}
	}
	return json.Marshal(details)
}

func debugFromTDAttributesHex(tdAttributesHex string) bool {
	if len(tdAttributesHex) >= 2 {
		if b, err := strconv.ParseUint(tdAttributesHex[:2], 16, 8); err == nil {
			return (b & 0x01) != 0
		}
	}
	return false
}
