// 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 (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"time"

	"github.com/dashscope/dashscope-confidential-maas/pkg/protocol"
	"github.com/dashscope/dashscope-confidential-maas/pkg/transparency"
	"github.com/dashscope/dashscope-confidential-maas/pkg/utils"

	"github.com/gorilla/websocket"
)

// performKeyNegotiation performs ECDHE key negotiation with secure proxy.
func (c *Client) performKeyNegotiation(conn *websocket.Conn, taskID string) ([]byte, *AttestationResult, error) {
	clientPrivKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate client ECDH key pair: %w", err)
	}

	clientPubKeyPEM, err := utils.ConvertECDSAPubKeyToPEM(&clientPrivKey.PublicKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to convert client public key to PEM: %w", err)
	}

	clientPubKey := utils.MarshalECDHPublicKey(&clientPrivKey.PublicKey)

	keyNegReq := &DashScopeMessage{
		Header: DashScopeHeader{
			TaskID:     taskID,
			Streaming:  "duplex",
			Action:     "run-task",
			Attributes: sdkHeaderAttributes(),
		},
		Payload: DashScopePayload{
			TaskGroup: c.route.taskGroup,
			Task:      c.route.task,
			Function:  c.route.function,
			Model:     c.model,
			Input: map[string]interface{}{
				"client_public_key": clientPubKeyPEM,
				"action":            "key_negotiation",
				"protocol_version":  protocol.CurrentProtocolVersion,
			},
		},
	}

	if err := conn.WriteJSON(keyNegReq); err != nil {
		return nil, nil, fmt.Errorf("failed to send key negotiation request: %w", err)
	}

	var taskStartedResp DashScopeMessage
	if err := conn.ReadJSON(&taskStartedResp); err != nil {
		return nil, nil, fmt.Errorf("failed to read task-started response: %w", err)
	}

	if taskStartedResp.Header.Event == "task-failed" || taskStartedResp.Header.ErrorName != "" {
		return nil, nil, fmt.Errorf("task failed: [%s] %s",
			taskStartedResp.Header.ErrorName,
			taskStartedResp.Header.ErrorMessage)
	}

	if taskStartedResp.Header.Event != "task-started" {
		return nil, nil, fmt.Errorf("unexpected first response event: %s (expected: task-started)", taskStartedResp.Header.Event)
	}

	var keyNegResp DashScopeMessage
	if err := conn.ReadJSON(&keyNegResp); err != nil {
		return nil, nil, fmt.Errorf("failed to read key negotiation result: %w", err)
	}

	if keyNegResp.Header.Event == "task-failed" || keyNegResp.Header.ErrorName != "" {
		return nil, nil, fmt.Errorf("key negotiation failed: [%s] %s",
			keyNegResp.Header.ErrorName,
			keyNegResp.Header.ErrorMessage)
	}

	var keyNegOutput KeyNegotiationOutput
	if err := json.Unmarshal(keyNegResp.Payload.Output, &keyNegOutput); err != nil {
		return nil, nil, fmt.Errorf("failed to parse key negotiation output: %w", err)
	}
	if keyNegOutput.ProtocolVersion != 0 && keyNegOutput.ProtocolVersion != protocol.CurrentProtocolVersion {
		return nil, nil, fmt.Errorf("unsupported protocol_version in key negotiation response: %d", keyNegOutput.ProtocolVersion)
	}

	bundle := keyNegOutput.AttestationBundle
	if bundle.ServerPublicKey == "" {
		return nil, nil, fmt.Errorf("missing server_public_key in attestation bundle")
	}
	if bundle.EphemeralKeyShare == "" {
		return nil, nil, fmt.Errorf("missing ephemeral_key_share in attestation bundle")
	}
	if bundle.Signature == "" {
		return nil, nil, fmt.Errorf("missing signature in attestation bundle")
	}

	serverPubKeyPEM := bundle.ServerPublicKey

	// Cache hit: reuse previously verified result, derive new session key
	if cached := c.findPubKeyCache(serverPubKeyPEM); cached != nil {
		sessionKey, err := deriveSessionKeyFromEphemeral(
			clientPrivKey, clientPubKey,
			bundle.EphemeralKeyShare, bundle.Signature, serverPubKeyPEM,
		)
		if err != nil {
			return nil, nil, err
		}

		result := &AttestationResult{
			VerifiedAt: cached.verifiedAt,
			Bundle:     bundle,
		}
		return sessionKey, result, nil
	}

	// Full verification path
	if len(bundle.Evidence) == 0 {
		return nil, nil, fmt.Errorf("missing evidence in attestation bundle")
	}

	verifier, ok := c.teeVerifiers[bundle.TEEType]
	if !ok {
		return nil, nil, fmt.Errorf("unsupported TEE type: %s", bundle.TEEType)
	}
	if err := verifier.VerifyEvidence(bundle.Evidence, serverPubKeyPEM); err != nil {
		return nil, nil, fmt.Errorf("TEE verification failed: %w", err)
	}

	// Verify transparency log
	if bundle.Transparency == nil {
		return nil, nil, fmt.Errorf("missing transparency data in attestation bundle")
	}
	if len(bundle.Transparency.Proofs) == 0 {
		return nil, nil, fmt.Errorf("no transparency proof entries in attestation bundle")
	}
	for _, proof := range bundle.Transparency.Proofs {
		if err := c.transparencySvc.VerifyProofEntry(
			&bundle.Transparency.ReleaseManifest,
			&proof,
		); err != nil {
			return nil, nil, fmt.Errorf("transparency log verification failed: %w", err)
		}
	}
	if err := transparency.VerifyTEEMeasurements(
		&bundle.Transparency.ReleaseManifest,
		bundle.Evidence,
		bundle.TEEType,
	); err != nil {
		return nil, nil, fmt.Errorf("TEE measurement cross-validation failed: %w", err)
	}

	// Derive session key
	sessionKey, err := deriveSessionKeyFromEphemeral(
		clientPrivKey, clientPubKey,
		bundle.EphemeralKeyShare, bundle.Signature, serverPubKeyPEM,
	)
	if err != nil {
		return nil, nil, err
	}

	result := &AttestationResult{
		VerifiedAt: time.Now(),
		Bundle:     bundle,
	}

	c.addPubKeyCache(serverPubKeyPEM, result)

	return sessionKey, result, nil
}

// deriveSessionKeyFromEphemeral computes shared secret and derives session key
func deriveSessionKeyFromEphemeral(
	clientPrivKey *ecdsa.PrivateKey,
	clientPubKey []byte,
	ephemeralKeySharePEM string,
	signatureB64 string,
	serverPubKeyPEM string,
) ([]byte, error) {
	// Parse ephemeral key share
	ephemeralKey, err := utils.ParseECDSAPubKeyFromPEM(ephemeralKeySharePEM)
	if err != nil {
		return nil, fmt.Errorf("failed to parse ephemeral key share: %w", err)
	}

	// Decode signature
	signature, err := base64.StdEncoding.DecodeString(signatureB64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode signature: %w", err)
	}

	// Parse server public key
	serverPubKey, err := utils.ParseECDSAPubKeyFromPEM(serverPubKeyPEM)
	if err != nil {
		return nil, fmt.Errorf("failed to parse server public key: %w", err)
	}

	// Verify ephemeral key share signature
	if err := utils.VerifyECDSAPublicKeySignature(ephemeralKey, signature, serverPubKey); err != nil {
		return nil, fmt.Errorf("ephemeral key share signature verification failed: %w", err)
	}

	// Convert ephemeral ECDSA public key to raw format for ECDH
	ephemeralECDHPublic := utils.MarshalECDHPublicKey(ephemeralKey)

	// Compute ECDH shared secret
	sharedSecret, err := utils.ComputeECDHSharedSecret(clientPrivKey, ephemeralECDHPublic)
	if err != nil {
		return nil, fmt.Errorf("failed to compute ECDH shared secret: %w", err)
	}

	// Derive session key
	sessionKey, err := utils.DeriveSessionKey(sharedSecret, clientPubKey, ephemeralECDHPublic)
	// Zero shared secret after key derivation
	clear(sharedSecret)
	if err != nil {
		return nil, fmt.Errorf("failed to derive session key: %w", err)
	}

	return sessionKey, nil
}
