// 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 provides TDX quote generation and verification functionality without CGO dependency
package tdx

import (
	"encoding/binary"
	"fmt"
)

// ParseQuote parses a TDX quote from raw bytes and returns a structured representation
func ParseQuote(rawQuote []byte) (*TdxEcdsaQuote, error) {
	if len(rawQuote) < 48 {
		return nil, fmt.Errorf("quote too short to be parsed")
	}

	// Parse QuoteHeader
	header := QuoteHeader{
		Version:            binary.LittleEndian.Uint16(rawQuote[0:2]),
		AttestationKeyType: binary.LittleEndian.Uint16(rawQuote[2:4]),
		TeeType:            binary.LittleEndian.Uint32(rawQuote[4:8]),
		QeSvn:              binary.LittleEndian.Uint16(rawQuote[8:10]),
		PceSvn:             binary.LittleEndian.Uint16(rawQuote[10:12]),
		QeVendorId:         [16]byte(rawQuote[12:28]),
		UserData:           [20]byte(rawQuote[28:48]),
	}

	// Validate header
	if header.Version < QUOTE_VERSION_3 || header.Version > QUOTE_VERSION_5 {
		return nil, fmt.Errorf("unsupported quote version: %d", header.Version)
	}

	if header.AttestationKeyType != ECDSA_256_WITH_P256_CURVE {
		return nil, fmt.Errorf("unsupported attestation key type: %d", header.AttestationKeyType)
	}

	if header.TeeType != TEE_TYPE_SGX && header.TeeType != TEE_TYPE_TDX {
		return nil, fmt.Errorf("unsupported TEE type: %d", header.TeeType)
	}

	parsedQuote := &TdxEcdsaQuote{
		Header: header,
	}

	// Parse based on version and TEE type
	offset := 48
	switch header.Version {
	case QUOTE_VERSION_3:
		if err := parseQuoteV3(rawQuote, parsedQuote, &offset); err != nil {
			return nil, err
		}
	case QUOTE_VERSION_4:
		if err := parseQuoteV4(rawQuote, parsedQuote, &offset); err != nil {
			return nil, err
		}
	case QUOTE_VERSION_5:
		if err := parseQuoteV5(rawQuote, parsedQuote, &offset); err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf("unsupported quote version: %d", header.Version)
	}

	return parsedQuote, nil
}

// parseQuoteV3 parses a v3 quote
func parseQuoteV3(rawQuote []byte, parsedQuote *TdxEcdsaQuote, offset *int) error {
	// V3 quotes are only for SGX
	if parsedQuote.Header.TeeType != TEE_TYPE_SGX {
		return fmt.Errorf("v3 quotes are only supported for SGX")
	}

	// Parse report based on TEE type
	if err := parseReport(rawQuote, parsedQuote, offset, BODY_SGX_ENCLAVE_REPORT_TYPE); err != nil {
		return err
	}

	// Parse signature data
	if err := parseSignatureData(rawQuote, parsedQuote, offset); err != nil {
		return err
	}

	return nil
}

// parseQuoteV4 parses a v4 quote
func parseQuoteV4(rawQuote []byte, parsedQuote *TdxEcdsaQuote, offset *int) error {
	// Parse report based on TEE type
	var reportType uint16
	switch parsedQuote.Header.TeeType {
	case TEE_TYPE_SGX:
		reportType = BODY_SGX_ENCLAVE_REPORT_TYPE
	case TEE_TYPE_TDX:
		reportType = BODY_TD_REPORT10_TYPE
	}

	if err := parseReport(rawQuote, parsedQuote, offset, reportType); err != nil {
		return err
	}

	// Parse signature data
	if err := parseSignatureData(rawQuote, parsedQuote, offset); err != nil {
		return err
	}

	return nil
}

// parseQuoteV5 parses a v5 quote
func parseQuoteV5(rawQuote []byte, parsedQuote *TdxEcdsaQuote, offset *int) error {
	// Parse quote body descriptor for v5 quotes (mandatory)
	if *offset+6 > len(rawQuote) {
		return fmt.Errorf("quote too short to parse body descriptor")
	}
	bodyDescriptor := &TdQuoteBodyDescriptor{
		TdQuoteBodyType: binary.LittleEndian.Uint16(rawQuote[*offset : *offset+2]),
		Size:            binary.LittleEndian.Uint32(rawQuote[*offset+2 : *offset+6]),
	}
	parsedQuote.BodyDescriptor = bodyDescriptor
	*offset += 6

	// Parse report based on TEE type and body descriptor
	var reportType uint16
	switch parsedQuote.Header.TeeType {
	case TEE_TYPE_SGX:
		reportType = BODY_SGX_ENCLAVE_REPORT_TYPE
	case TEE_TYPE_TDX:
		reportType = parsedQuote.BodyDescriptor.TdQuoteBodyType
	}

	if err := parseReport(rawQuote, parsedQuote, offset, reportType); err != nil {
		return err
	}

	// Parse signature data
	if err := parseSignatureData(rawQuote, parsedQuote, offset); err != nil {
		return err
	}

	return nil
}

// parseReport parses the report section of a quote
func parseReport(rawQuote []byte, parsedQuote *TdxEcdsaQuote, offset *int, reportType uint16) error {
	switch reportType {
	case BODY_SGX_ENCLAVE_REPORT_TYPE:
		if *offset+384 > len(rawQuote) {
			return fmt.Errorf("quote too short to parse enclave report")
		}
		report := EnclaveReportBody{
			CpuSvn:     [16]byte(rawQuote[*offset : *offset+16]),
			MiscSelect: binary.LittleEndian.Uint32(rawQuote[*offset+16 : *offset+20]),
			Reserved1:  [28]byte(rawQuote[*offset+20 : *offset+48]),
			Attributes: [16]byte(rawQuote[*offset+48 : *offset+64]),
			MrEnclave:  [32]byte(rawQuote[*offset+64 : *offset+96]),
			Reserved2:  [32]byte(rawQuote[*offset+96 : *offset+128]),
			MrSigner:   [32]byte(rawQuote[*offset+128 : *offset+160]),
			Reserved3:  [96]byte(rawQuote[*offset+160 : *offset+256]),
			IsvProdID:  binary.LittleEndian.Uint16(rawQuote[*offset+256 : *offset+258]),
			IsvSvn:     binary.LittleEndian.Uint16(rawQuote[*offset+258 : *offset+260]),
			Reserved4:  [60]byte(rawQuote[*offset+260 : *offset+320]),
			ReportData: [64]byte(rawQuote[*offset+320 : *offset+384]),
		}
		parsedQuote.Report = report
		*offset += 384
	case BODY_TD_REPORT10_TYPE:
		if *offset+584 > len(rawQuote) {
			return fmt.Errorf("quote too short to parse TD report 1.0")
		}
		report := TdReport10{
			TeeTcbSvn:      [16]byte(rawQuote[*offset : *offset+16]),
			MrSeam:         [48]byte(rawQuote[*offset+16 : *offset+64]),
			MrSignerSeam:   [48]byte(rawQuote[*offset+64 : *offset+112]),
			SeamAttributes: [8]byte(rawQuote[*offset+112 : *offset+120]),
			TdAttributes:   [8]byte(rawQuote[*offset+120 : *offset+128]),
			XFAM:           [8]byte(rawQuote[*offset+128 : *offset+136]),
			MrTd:           [48]byte(rawQuote[*offset+136 : *offset+184]),
			MrConfigId:     [48]byte(rawQuote[*offset+184 : *offset+232]),
			MrOwner:        [48]byte(rawQuote[*offset+232 : *offset+280]),
			MrOwnerConfig:  [48]byte(rawQuote[*offset+280 : *offset+328]),
			Rtmr0:          [48]byte(rawQuote[*offset+328 : *offset+376]),
			Rtmr1:          [48]byte(rawQuote[*offset+376 : *offset+424]),
			Rtmr2:          [48]byte(rawQuote[*offset+424 : *offset+472]),
			Rtmr3:          [48]byte(rawQuote[*offset+472 : *offset+520]),
			ReportData:     [64]byte(rawQuote[*offset+520 : *offset+584]),
		}
		parsedQuote.Report = report
		*offset += 584
	case BODY_TD_REPORT15_TYPE:
		if *offset+648 > len(rawQuote) {
			return fmt.Errorf("quote too short to parse TD report 1.5")
		}
		report := TdReport15{
			TeeTcbSvn:      [16]byte(rawQuote[*offset : *offset+16]),
			MrSeam:         [48]byte(rawQuote[*offset+16 : *offset+64]),
			MrSignerSeam:   [48]byte(rawQuote[*offset+64 : *offset+112]),
			SeamAttributes: [8]byte(rawQuote[*offset+112 : *offset+120]),
			TdAttributes:   [8]byte(rawQuote[*offset+120 : *offset+128]),
			XFAM:           [8]byte(rawQuote[*offset+128 : *offset+136]),
			MrTd:           [48]byte(rawQuote[*offset+136 : *offset+184]),
			MrConfigId:     [48]byte(rawQuote[*offset+184 : *offset+232]),
			MrOwner:        [48]byte(rawQuote[*offset+232 : *offset+280]),
			MrOwnerConfig:  [48]byte(rawQuote[*offset+280 : *offset+328]),
			Rtmr0:          [48]byte(rawQuote[*offset+328 : *offset+376]),
			Rtmr1:          [48]byte(rawQuote[*offset+376 : *offset+424]),
			Rtmr2:          [48]byte(rawQuote[*offset+424 : *offset+472]),
			Rtmr3:          [48]byte(rawQuote[*offset+472 : *offset+520]),
			ReportData:     [64]byte(rawQuote[*offset+520 : *offset+584]),
			TeeTcbSvn2:     [16]byte(rawQuote[*offset+584 : *offset+600]),
			MrServiceTd:    [48]byte(rawQuote[*offset+600 : *offset+648]),
		}
		parsedQuote.Report = report
		*offset += 648
	default:
		return fmt.Errorf("unsupported report type: %d", reportType)
	}

	return nil
}

// parseSignatureData parses the signature data section of a quote
func parseSignatureData(rawQuote []byte, parsedQuote *TdxEcdsaQuote, offset *int) error {
	// Parse signature data size
	if *offset+4 > len(rawQuote) {
		return fmt.Errorf("quote too short to parse signature data size")
	}
	signatureDataSize := binary.LittleEndian.Uint32(rawQuote[*offset : *offset+4])
	parsedQuote.SignatureDataLen = signatureDataSize
	*offset += 4

	// Check if we have enough data for at least signature + key (128 bytes)
	if int(signatureDataSize) < 128 {
		return fmt.Errorf("signature data size %d is too small, expected at least 128 bytes", signatureDataSize)
	}
	if *offset+int(signatureDataSize) > len(rawQuote) {
		return fmt.Errorf("quote too short to parse signature data")
	}

	// Parse ECDSA signature (64 bytes)
	if *offset+64 > len(rawQuote) {
		return fmt.Errorf("quote too short to parse ECDSA signature")
	}
	copy(parsedQuote.QuoteSignature[:], rawQuote[*offset:*offset+64])
	*offset += 64

	// Parse ECDSA attestation key (64 bytes)
	if *offset+64 > len(rawQuote) {
		return fmt.Errorf("quote too short to parse ECDSA attestation key")
	}
	copy(parsedQuote.EcdsaAttestationKey[:], rawQuote[*offset:*offset+64])
	*offset += 64

	// Parse QE certification data based on TEE type
	switch parsedQuote.Header.TeeType {
	case TEE_TYPE_SGX:
		// SGX quotes: EcdsaAttestationKey is directly followed by QeReportCertificationData
		remainingData := rawQuote[*offset:]
		qeReportCertData, err := parseQeReportCertificationData(remainingData)
		if err != nil {
			return fmt.Errorf("failed to parse SGX QE Report Certification Data: %w", err)
		}

		// For SGX, we wrap the QeReportCertificationData in a QeCertificationData structure
		// to maintain consistency with the unified data model
		parsedQuote.QeCertificationData = QeCertificationData{
			CertificationDataType: 6, // SGX also uses type 6 for QE Report Certification Data
			Size:                  uint32(len(remainingData)),
			CertificationData:     qeReportCertData,
		}

		// Update offset to end of quote
		*offset = len(rawQuote)

	case TEE_TYPE_TDX:
		// TDX quotes: EcdsaAttestationKey is followed by QeCertificationData (outer wrapper)
		certData, err := parseQeCertificationData(rawQuote, offset)
		if err != nil {
			return err
		}
		parsedQuote.QeCertificationData = *certData

		// Validate that the outer certification data type is 6
		if parsedQuote.QeCertificationData.CertificationDataType != 6 {
			return fmt.Errorf("TDX quote must have certification data type 6 (QE Report Certification Data), got type %d",
				parsedQuote.QeCertificationData.CertificationDataType)
		}

	default:
		return fmt.Errorf("unsupported TEE type for signature data parsing: %d", parsedQuote.Header.TeeType)
	}

	return nil
}

// parseQeCertificationData parses QE certification data.
// This function can be reused for both outer and inner QE certification data parsing
func parseQeCertificationData(data []byte, offset *int) (*QeCertificationData, error) {
	// Parse certification data type (2 bytes)
	if *offset+2 > len(data) {
		return nil, fmt.Errorf("data too short to parse certification data type")
	}
	certDataType := binary.LittleEndian.Uint16(data[*offset : *offset+2])
	*offset += 2

	// Parse certification data size (4 bytes)
	if *offset+4 > len(data) {
		return nil, fmt.Errorf("data too short to parse certification data size")
	}
	certDataSize := binary.LittleEndian.Uint32(data[*offset : *offset+4])
	*offset += 4

	// Validate we have enough data
	if *offset+int(certDataSize) > len(data) {
		return nil, fmt.Errorf("data too short to parse certification data")
	}

	var certificationData interface{}

	// Parse certification data based on type
	if certDataType == 6 {
		// Type 6: QE Report Certification Data (expected for TDX quotes)
		certDataBytes := make([]byte, certDataSize)
		copy(certDataBytes, data[*offset:*offset+int(certDataSize)])

		qeReportCertData, err := parseQeReportCertificationData(certDataBytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse QE Report Certification Data: %w", err)
		}

		certificationData = qeReportCertData
	} else {
		// Other types: store as raw bytes (for compatibility or non-standard quotes)
		certDataBytes := make([]byte, certDataSize)
		copy(certDataBytes, data[*offset:*offset+int(certDataSize)])
		certificationData = certDataBytes
	}

	*offset += int(certDataSize)

	return &QeCertificationData{
		CertificationDataType: certDataType,
		Size:                  certDataSize,
		CertificationData:     certificationData,
	}, nil
}

// parseQeReportCertificationData parses QE Report Certification Data from raw bytes
func parseQeReportCertificationData(data []byte) (*QeReportCertificationData, error) {
	if len(data) < 384+64+2 { // QE report + signature + auth data size
		return nil, fmt.Errorf("data too short for QE Report Certification Data")
	}

	qeReportCertData := &QeReportCertificationData{}
	offset := 0

	// Parse QE report (384 bytes)
	qeReportCertData.QeReport = EnclaveReportBody{
		CpuSvn:     [16]byte(data[offset : offset+16]),
		MiscSelect: binary.LittleEndian.Uint32(data[offset+16 : offset+20]),
		Reserved1:  [28]byte(data[offset+20 : offset+48]),
		Attributes: [16]byte(data[offset+48 : offset+64]),
		MrEnclave:  [32]byte(data[offset+64 : offset+96]),
		Reserved2:  [32]byte(data[offset+96 : offset+128]),
		MrSigner:   [32]byte(data[offset+128 : offset+160]),
		Reserved3:  [96]byte(data[offset+160 : offset+256]),
		IsvProdID:  binary.LittleEndian.Uint16(data[offset+256 : offset+258]),
		IsvSvn:     binary.LittleEndian.Uint16(data[offset+258 : offset+260]),
		Reserved4:  [60]byte(data[offset+260 : offset+320]),
		ReportData: [64]byte(data[offset+320 : offset+384]),
	}
	offset += 384

	// Parse QE report signature (64 bytes)
	if offset+64 > len(data) {
		return nil, fmt.Errorf("data too short to parse QE report signature")
	}
	copy(qeReportCertData.QeReportSignature[:], data[offset:offset+64])
	offset += 64

	// Parse QE authentication data size (2 bytes)
	if offset+2 > len(data) {
		return nil, fmt.Errorf("data too short to parse QE auth data size")
	}
	qeAuthDataSize := binary.LittleEndian.Uint16(data[offset : offset+2])
	offset += 2

	// Parse QE authentication data
	if offset+int(qeAuthDataSize) > len(data) {
		return nil, fmt.Errorf("data too short to parse QE authentication data")
	}
	qeReportCertData.QeAuthenticationData = make([]byte, qeAuthDataSize)
	copy(qeReportCertData.QeAuthenticationData, data[offset:offset+int(qeAuthDataSize)])
	offset += int(qeAuthDataSize)

	// Parse inner certification data using the reusable function
	tempOffset := 0
	innerCertData, err := parseQeCertificationData(data[offset:], &tempOffset)
	if err != nil {
		return nil, fmt.Errorf("failed to parse inner certification data: %w", err)
	}
	qeReportCertData.QeCertificationData = *innerCertData

	return qeReportCertData, nil
}

// ExtractMRTD parses a TDX quote and returns the MRTD (MrTd) as a hex string.
func ExtractMRTD(quoteData []byte) (string, error) {
	q, err := ParseQuote(quoteData)
	if err != nil {
		return "", fmt.Errorf("parse TDX quote: %w", err)
	}
	switch r := q.Report.(type) {
	case TdReport10:
		return fmt.Sprintf("%x", r.MrTd[:]), nil
	case TdReport15:
		return fmt.Sprintf("%x", r.MrTd[:]), nil
	default:
		return "", fmt.Errorf("unexpected report type: %T", q.Report)
	}
}
