// 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 (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

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

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/patrickmn/go-cache"
)

// verifiedPubKeyCacheEntry stores verification timestamp for a server public key.
type verifiedPubKeyCacheEntry struct {
	verifiedAt time.Time
}

// Client is the ModelStudio CMaaS client for confidential inference
type Client struct {
	// Configuration (immutable)
	endpoint        string
	apiKey          string
	model           string
	route           modelRoute
	reuseSession    bool
	pubKeyCacheTTL  time.Duration
	reportWriter    *reportWriter                // transparency audit report writer (nil = disabled)
	transparencySvc transparency.TranslogService // transparency log verification service
	teeVerifiers    map[string]teeVerifier       // supported TEE verifiers by type

	// Server public keys cache (for multiple servers behind load balancer)
	// Uses go-cache for automatic expiration and cleanup
	pubKeyCache *cache.Cache

	// Chat calls are serialized per client because one session owns one WebSocket reader/writer.
	chatMu sync.Mutex

	// Current active session, protected by mu.
	mu             sync.Mutex
	currentSession *Session
}

// NewClient creates a new CMaaS client (lazy loading)
func NewClient(opts *Options) (*Client, error) {
	if opts.Endpoint == "" {
		return nil, fmt.Errorf("endpoint is required")
	}
	if opts.APIKey == "" {
		return nil, fmt.Errorf("API key is required")
	}
	if opts.Model == "" {
		return nil, fmt.Errorf("model is required")
	}

	// Set public key cache TTL
	// -1: disable cache, 0: default 12h, >0: custom TTL
	ttl := opts.PubKeyCacheTTL
	if ttl == 0 {
		ttl = 12 * time.Hour // Default: 12 hours
	} else if ttl < 0 {
		ttl = -1 // Mark as disabled
	}

	// Initialize cache with automatic expiration and cleanup
	var pubKeyCache *cache.Cache
	if ttl > 0 {
		// cleanupInterval: run expired item deletion every ttl/2
		// This ensures old entries are removed even if never accessed again
		pubKeyCache = cache.New(ttl, ttl/2)
	}

	// Resolve transparency report directory and writer
	var rw *reportWriter
	if !opts.DisableTransparencyReport {
		reportDir := opts.TransparencyReportDir
		if reportDir == "" {
			if d, err := defaultReportDir(); err == nil {
				reportDir = d
			}
		}
		if reportDir != "" {
			rw = newReportWriter(reportDir)
		}
	}

	// Create transparency log verification service
	logType := opts.TransparencyLogType
	if logType == "" {
		logType = transparency.LogTypeRekorV1
	}
	transConfig := opts.TransparencyConfig
	if transConfig == nil {
		transConfig = &transparency.RekorV1ServiceConfig{}
	}
	transparencySvc, err := transparency.NewTranslogService(logType, transConfig)
	if err != nil {
		return nil, fmt.Errorf("transparency service: %w", err)
	}

	return &Client{
		endpoint:        opts.Endpoint,
		apiKey:          opts.APIKey,
		model:           opts.Model,
		route:           routeForModel(opts.Model),
		reuseSession:    opts.ReuseSession,
		pubKeyCacheTTL:  ttl,
		pubKeyCache:     pubKeyCache,
		reportWriter:    rw,
		transparencySvc: transparencySvc,
		teeVerifiers:    supportedTEEVerifiers,
	}, nil
}

// Chat performs inference request and returns ChatCompletion
// For non-streaming: completion.Content contains full response
// For streaming: read from completion.Chunks channel
func (c *Client) Chat(ctx context.Context, req *ChatCompletionRequest) (*ChatCompletion, error) {
	if req == nil {
		return nil, &ValidationError{Message: "request is required"}
	}
	if len(req.Messages) == 0 {
		return nil, &ValidationError{Message: "messages are required"}
	}

	isStream := req.Stream
	c.chatMu.Lock()
	defer func() {
		if !isStream {
			c.chatMu.Unlock()
		}
	}()

	session, err := c.prepareSession(ctx, req.RequestID)
	if err != nil {
		if isStream {
			c.chatMu.Unlock()
		}
		return nil, err
	}

	completion, err := c.doInference(ctx, session, req, func() {
		c.finishStreamChat(session)
	})

	// Retry on connection error (reuse mode only, non-reuse always creates fresh session)
	if err != nil && c.reuseSession && isConnectionError(err) {
		log.Printf("[WARN] connection lost, reconnecting task_id=%s error=%v", session.taskID, err)
		session, err = c.reconnect(ctx, req.RequestID)
		if err != nil {
			if isStream {
				c.chatMu.Unlock()
			}
			return nil, err
		}
		completion, err = c.doInference(ctx, session, req, func() {
			c.finishStreamChat(session)
		})
		if err != nil {
			if isStream {
				c.chatMu.Unlock()
			}
			return nil, err
		}
	}

	if completion != nil {
		completion.Attestation = session.attestation
		_ = c.saveTransparencyReport(req.Messages, session)
	}

	// Non-reuse mode + non-streaming: close session after use
	if !c.reuseSession && !isStream {
		c.mu.Lock()
		c.closeSession()
		c.mu.Unlock()
	}

	if isStream && (completion == nil || err != nil) {
		c.chatMu.Unlock()
	}

	return completion, err
}

func (c *Client) finishStreamChat(session *Session) {
	defer c.chatMu.Unlock()

	if c.reuseSession {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.currentSession == session {
		_ = c.closeSession()
		return
	}
	_ = session.Close()
}

// prepareSession returns a ready-to-use session
// - Reuse mode: return existing session or create new one
// - Non-reuse mode: always create fresh session
func (c *Client) prepareSession(ctx context.Context, taskID string) (*Session, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.reuseSession || (taskID != "" && c.currentSession != nil && c.currentSession.taskID != taskID) {
		c.closeSession()
	}
	if c.currentSession == nil {
		if err := c.establishSession(ctx, taskID); err != nil {
			return nil, err
		}
	}
	return c.currentSession, nil
}

// reconnect forces a new session by closing current and establishing new one
func (c *Client) reconnect(ctx context.Context, taskID string) (*Session, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.closeSession()
	if err := c.establishSession(ctx, taskID); err != nil {
		return nil, fmt.Errorf("reconnect failed: %w", err)
	}
	return c.currentSession, nil
}

// doInference executes inference (streaming or non-streaming)
func (c *Client) doInference(ctx context.Context, session *Session, req *ChatCompletionRequest, finishStreamChat func()) (*ChatCompletion, error) {
	if req.Stream {
		return c.performInferenceStream(ctx, session, req, finishStreamChat)
	}
	return c.performInference(ctx, session, req)
}

// Close closes the client connection
func (c *Client) Close() error {
	c.chatMu.Lock()
	defer c.chatMu.Unlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	err := c.closeSession()
	if c.reportWriter != nil {
		if cerr := c.reportWriter.Close(); cerr != nil && err == nil {
			err = cerr
		}
	}
	return err
}

// CurrentTaskID returns the active DashScope task_id, if a session exists.
func (c *Client) CurrentTaskID() string {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.currentSession == nil {
		return ""
	}
	return c.currentSession.taskID
}

// establishSession establishes WebSocket connection and performs key negotiation
func (c *Client) establishSession(ctx context.Context, taskID string) error {
	// 1. Establish WebSocket
	conn, err := dialWebSocket(ctx, c.endpoint, c.apiKey)
	if err != nil {
		return err
	}

	if taskID == "" {
		taskID = uuid.New().String()
	}

	// 3. Key negotiation (check cache or full verification)
	sessionKey, result, err := c.performKeyNegotiation(conn, taskID)
	if err != nil {
		conn.Close()
		return err
	}

	// 4. Create session. Keep the gateway connection active during the request
	// even when ReuseSession=false; non-reuse still closes the session after use.
	c.currentSession = newSession(conn, sessionKey, result, result.Bundle.ServerPublicKey, taskID)

	return nil
}

func (c *Client) findPubKeyCache(pubKey string) *verifiedPubKeyCacheEntry {
	if c.pubKeyCache == nil {
		return nil
	}
	if val, found := c.pubKeyCache.Get(pubKey); found {
		entry := val.(*verifiedPubKeyCacheEntry)
		return &verifiedPubKeyCacheEntry{
			verifiedAt: entry.verifiedAt,
		}
	}
	return nil
}

func (c *Client) addPubKeyCache(pubKey string, result *AttestationResult) {
	if c.pubKeyCache == nil {
		return
	}
	c.pubKeyCache.Set(pubKey, &verifiedPubKeyCacheEntry{
		verifiedAt: result.VerifiedAt,
	}, cache.DefaultExpiration)
}

// closeSession closes current session
func (c *Client) closeSession() error {
	if c.currentSession != nil {
		err := c.currentSession.Close()
		c.currentSession = nil
		return err
	}
	return nil
}

// isConnectionError checks if error is a connection error
func isConnectionError(err error) bool {
	if err == nil {
		return false
	}
	// Check for WebSocket close errors
	if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
		return true
	}
	// Check for DashScope Gateway specific errors that require reconnection
	errStr := err.Error()
	// DashScope Gateway closes connection after 60s inactivity, returns "inference failed: [] Response timeout."
	if errStr == "inference failed: [] Response timeout." {
		return true
	}
	return false
}
