// 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 (
	"encoding/json"
	"fmt"
	"time"

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

// DashScope gateway protocol message types

// DashScopeMessage represents a complete DashScope protocol message
type DashScopeMessage struct {
	Header  DashScopeHeader  `json:"header"`
	Payload DashScopePayload `json:"payload"`
}

// DashScopeHeader contains message metadata and status
type DashScopeHeader struct {
	TaskID        string                 `json:"task_id"`
	Streaming     string                 `json:"streaming"`
	Action        string                 `json:"action"`
	Event         string                 `json:"event,omitempty"`
	ErrorCode     string                 `json:"error_code,omitempty"`
	ErrorName     string                 `json:"error_name,omitempty"`
	ErrorMessage  string                 `json:"error_message,omitempty"`
	StatusCode    int                    `json:"status_code,omitempty"`
	StatusName    string                 `json:"status_name,omitempty"`
	StatusMessage string                 `json:"status_message,omitempty"`
	Attributes    map[string]interface{} `json:"attributes,omitempty"`
}

// DashScopePayload contains the actual request/response data
type DashScopePayload struct {
	TaskGroup  string                 `json:"task_group,omitempty"`
	Task       string                 `json:"task,omitempty"`
	Function   string                 `json:"function,omitempty"`
	Model      string                 `json:"model,omitempty"`
	Input      map[string]interface{} `json:"input,omitempty"`
	Parameters map[string]interface{} `json:"parameters,omitempty"`
	Output     json.RawMessage        `json:"output,omitempty"`
	Usage      json.RawMessage        `json:"usage,omitempty"`
}

// Message represents a chat message
type Message struct {
	Role    string `json:"role"`
	Content string `json:"-"`
	// ContentParts carries OpenAI-compatible multimodal content. When set, it
	// serializes as the JSON `content` field instead of Content.
	ContentParts []ContentPart `json:"-"`
	ToolCalls    []ToolCall    `json:"tool_calls,omitempty"`
	ToolCallID   string        `json:"tool_call_id,omitempty"`
}

// ContentPart represents one OpenAI-compatible multimodal message content part.
type ContentPart struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageURL *ImageURL `json:"image_url,omitempty"`
}

// ImageURL represents OpenAI-compatible image_url payload.
type ImageURL struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"`
}

func (m Message) MarshalJSON() ([]byte, error) {
	type wireMessage struct {
		Role       string      `json:"role"`
		Content    interface{} `json:"content,omitempty"`
		ToolCalls  []ToolCall  `json:"tool_calls,omitempty"`
		ToolCallID string      `json:"tool_call_id,omitempty"`
	}
	var content interface{}
	if len(m.ContentParts) > 0 {
		content = m.ContentParts
	} else if m.Content != "" {
		content = m.Content
	}
	return json.Marshal(wireMessage{
		Role:       m.Role,
		Content:    content,
		ToolCalls:  m.ToolCalls,
		ToolCallID: m.ToolCallID,
	})
}

func (m *Message) UnmarshalJSON(data []byte) error {
	type wireMessage struct {
		Role       string          `json:"role"`
		Content    json.RawMessage `json:"content"`
		ToolCalls  []ToolCall      `json:"tool_calls,omitempty"`
		ToolCallID string          `json:"tool_call_id,omitempty"`
	}
	var wire wireMessage
	if err := json.Unmarshal(data, &wire); err != nil {
		return err
	}
	m.Role = wire.Role
	m.ToolCalls = wire.ToolCalls
	m.ToolCallID = wire.ToolCallID
	if len(wire.Content) == 0 || string(wire.Content) == "null" {
		return nil
	}
	var text string
	if err := json.Unmarshal(wire.Content, &text); err == nil {
		m.Content = text
		return nil
	}
	var parts []ContentPart
	if err := json.Unmarshal(wire.Content, &parts); err != nil {
		return fmt.Errorf("content must be string or array of content parts")
	}
	m.ContentParts = parts
	return nil
}

// ResponseMessage represents a chat message returned by the inference backend.
type ResponseMessage struct {
	Role             string     `json:"role"`
	Content          string     `json:"content"`
	ReasoningContent string     `json:"reasoning_content,omitempty"`
	Reasoning        string     `json:"reasoning,omitempty"`
	ToolCalls        []ToolCall `json:"tool_calls,omitempty"`
}

// Tool describes a function the model may invoke. Tool definitions are
// encrypted between the SDK and secure-proxy, then decrypted inside
// secure-proxy before forwarding to the inference backend.
type Tool struct {
	Type     string         `json:"type"` // always "function"
	Function FunctionSchema `json:"function"`
}

type FunctionSchema struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description,omitempty"`
	Parameters  map[string]interface{} `json:"parameters,omitempty"`
}

// ToolCall represents a single tool/function call from the assistant.
// Name and Arguments are encrypted on the wire and decrypted by the SDK
// before being surfaced to the caller.
type ToolCall struct {
	Index    *int         `json:"index,omitempty"`
	ID       string       `json:"id"`
	Type     string       `json:"type"` // always "function"
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// ToolCallDelta is a streaming fragment of a tool call. Each fragment is
// decrypted; consumers accumulate Arguments per Index.
type ToolCallDelta struct {
	ChoiceIndex int    `json:"choice_index,omitempty"`
	Index       int    `json:"index"`
	ID          string `json:"id,omitempty"`
	Type        string `json:"type,omitempty"`
	Name        string `json:"name,omitempty"`
	Arguments   string `json:"arguments,omitempty"`
}

// ChatCompletionRequest represents an OpenAI-compatible chat completion request.
//
// JSON tags with omitempty enable automatic parameter serialization:
// nil pointer fields are omitted; Messages and Stream are excluded (handled separately).
type ChatCompletionRequest struct {
	// RequestID is the DashScope task_id for this request. It is not sent as a
	// model parameter; when empty, the SDK generates one.
	RequestID string `json:"-"`

	// Required
	Messages []Message `json:"-"` // Handled separately (encrypted before sending)
	Stream   bool      `json:"stream"`

	// Sampling parameters (nil = not set, use server default)
	Temperature         *float64 `json:"temperature,omitempty"`
	TopP                *float64 `json:"top_p,omitempty"`
	TopK                *int     `json:"top_k,omitempty"`
	MaxTokens           *int     `json:"max_tokens,omitempty"`
	FrequencyPenalty    *float64 `json:"frequency_penalty,omitempty"`
	PresencePenalty     *float64 `json:"presence_penalty,omitempty"`
	RepetitionPenalty   *float64 `json:"repetition_penalty,omitempty"`
	ThinkingTokenBudget *int     `json:"thinking_token_budget,omitempty"`
	Logprobs            *bool    `json:"logprobs,omitempty"`
	TopLogprobs         *int     `json:"top_logprobs,omitempty"`
	// Control parameters
	N                *int            `json:"n,omitempty"`
	Seed             *int            `json:"seed,omitempty"`
	Stop             []string        `json:"stop,omitempty"`
	StopTokenIDs     []int           `json:"stop_token_ids,omitempty"`
	ResponseFormat   json.RawMessage `json:"response_format,omitempty"`
	StreamOptions    *StreamOptions  `json:"stream_options,omitempty"`
	StructuredStream bool            `json:"-"`

	// Function calling
	Tools             []Tool          `json:"tools,omitempty"`
	ToolChoice        json.RawMessage `json:"tool_choice,omitempty"`
	ParallelToolCalls *bool           `json:"parallel_tool_calls,omitempty"`

	// extension: chat template kwargs (e.g. enable_thinking)
	ChatTemplateKwargs map[string]interface{} `json:"chat_template_kwargs,omitempty"`
}

// StreamOptions configures streaming behavior (only effective when Stream=true)
type StreamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

// Helpers for setting optional pointer fields inline, similar to openai.Float() / openai.Int().
func Float(v float64) *float64 { return &v }
func Int(v int) *int           { return &v }
func Bool(v bool) *bool        { return &v }

// TokenUsage represents token usage statistics
type TokenUsage struct {
	PromptTokens            int                    `json:"prompt_tokens"`
	CompletionTokens        int                    `json:"completion_tokens"`
	TotalTokens             int                    `json:"total_tokens"`
	PromptTokensDetails     map[string]interface{} `json:"prompt_tokens_details,omitempty"`
	CompletionTokensDetails map[string]interface{} `json:"completion_tokens_details,omitempty"`
}

// InferenceOutput represents the output structure for inference responses
type InferenceOutput struct {
	ProtocolVersion int               `json:"protocol_version,omitempty"`
	Choices         []InferenceChoice `json:"choices"`
}

// InferenceChoice represents a single choice in the inference response
type InferenceChoice struct {
	Index        int              `json:"index,omitempty"`
	FinishReason string           `json:"finish_reason,omitempty"`
	Message      *ResponseMessage `json:"message,omitempty"` // Non-streaming response
	Delta        *ResponseMessage `json:"delta,omitempty"`   // Streaming response
	Logprobs     json.RawMessage  `json:"logprobs,omitempty"`
}

// StreamDelta represents a structured streaming delta returned by the SDK.
type StreamDelta struct {
	Index            int             `json:"index,omitempty"`
	Content          string          `json:"content,omitempty"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	Reasoning        string          `json:"reasoning,omitempty"`
	Logprobs         json.RawMessage `json:"logprobs,omitempty"`
}

// ChatCompletionChoice represents one non-streaming completion choice.
type ChatCompletionChoice struct {
	Index            int
	Content          string
	ReasoningContent string
	Reasoning        string
	FinishReason     string
	Logprobs         json.RawMessage
	ToolCalls        []ToolCall
}

type parsedInferenceResponse struct {
	Usage   TokenUsage
	Choices []responseChoice
}

type responseChoice struct {
	Index            int
	Content          string
	ReasoningContent string
	Reasoning        string
	FinishReason     string
	Logprobs         json.RawMessage
	ToolCalls        []ToolCall
}

// KeyNegotiationOutput represents the output structure for key negotiation responses
type KeyNegotiationOutput struct {
	ProtocolVersion   int                           `json:"protocol_version,omitempty"`
	AttestationBundle attestation.AttestationBundle `json:"attestation_bundle"`
}

// ChatCompletion represents a chat completion response
type ChatCompletion struct {
	Content          string          // Full content (non-streaming only)
	ReasoningContent string          // Full reasoning content (non-streaming only)
	Reasoning        string          // Full reasoning field (non-streaming only)
	Logprobs         json.RawMessage // Full logprobs field (non-streaming only)
	Choices          []ChatCompletionChoice
	ToolCalls        []ToolCall           // Full tool calls (non-streaming only)
	Chunks           <-chan string        // Text chunk channel (streaming only)
	ToolCallChunks   <-chan ToolCallDelta // Tool-call delta channel (streaming only)
	Deltas           <-chan StreamDelta   // Structured deltas (streaming only)
	FinishReason     string               // Final stream finish_reason
	FinishLogprobs   json.RawMessage      // Final stream logprobs field
	FinishChoices    []ChatCompletionChoice
	Usage            TokenUsage         // Token usage statistics
	Error            error              // Error channel for streaming mode
	TaskID           string             // DashScope task_id used by this request/session
	Attestation      *AttestationResult // Attestation result for this chat
}

// APIError carries the (statusCode, message) tuple surfaced by the CMaaS
// service or the inference backend after secure-proxy strips any backend
// envelope. The SDK never inspects the body for OpenAI-shaped fields; the
// HTTP-facing OpenAI envelope is the caller's responsibility.
type APIError struct {
	StatusCode int
	Message    string
}

func (e *APIError) Error() string {
	if e == nil {
		return ""
	}
	return e.Message
}

// ValidationError signals that a caller-supplied argument failed the SDK's
// own sanity checks (e.g. required messages missing). Callers should map
// this to an HTTP 400.
type ValidationError struct {
	Message string
}

func (e *ValidationError) Error() string {
	if e == nil {
		return ""
	}
	return e.Message
}

// AttestationResult contains TEE attestation verification details.
// Call TEEDetailsJSON() to get TEE-specific key fields (e.g. MRTD/RTMR for TDX) as JSON.
type AttestationResult struct {
	VerifiedAt time.Time                     `json:"verified_at"`
	Bundle     attestation.AttestationBundle `json:"bundle"`
}

// Options configures the client
type Options struct {
	Endpoint       string        // WebSocket gateway endpoint
	APIKey         string        // API key for authentication
	Model          string        // Model name for gateway routing
	ReuseSession   bool          // Enable session reuse (default: false)
	PubKeyCacheTTL time.Duration // Public key cache TTL (default: 12h, -1 to disable cache)

	// TransparencyLogType specifies the transparency log type (e.g. "rekor-v1").
	// When empty, defaults to "rekor-v1".
	TransparencyLogType transparency.LogType

	// TransparencyConfig provides service-specific configuration for the transparency log.
	// For rekor-v1: pass *transparency.RekorV1ServiceConfig to override verification keys,
	// or nil to use built-in defaults.
	TransparencyConfig any

	// TransparencyReportDir overrides the directory for transparency audit records.
	// When empty, an OS-specific default is used:
	//   macOS:   ~/Library/Application Support/dashscope-confidential-maas/transparency-reports/
	//   Linux:   $XDG_DATA_HOME/dashscope-confidential-maas/transparency-reports/
	//   Windows: %APPDATA%\dashscope-confidential-maas\transparency-reports\
	TransparencyReportDir string

	// DisableTransparencyReport disables local persistence of transparency audit records.
	DisableTransparencyReport bool
}
