package llm

import (
	"context"
	"encoding/json"
)

const (
	RoleSystem    = "system"
	RoleUser      = "user"
	RoleAssistant = "assistant"
	RoleTool      = "tool"
)

type Provider interface {
	ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error
}

type ChatRequest struct {
	Model      string         `json:"model"`
	Messages   []Message      `json:"messages"`
	Tools      []Tool         `json:"tools,omitempty"`
	ToolChoice string         `json:"tool_choice,omitempty"`
	Stream     bool           `json:"stream"`
	Extra      map[string]any `json:"-"`
}

func (r ChatRequest) MarshalJSON() ([]byte, error) {
	payload := map[string]any{
		"model":    r.Model,
		"messages": r.Messages,
		"stream":   r.Stream,
	}
	if len(r.Tools) > 0 {
		payload["tools"] = r.Tools
	}
	if r.ToolChoice != "" {
		payload["tool_choice"] = r.ToolChoice
	}
	for key, value := range r.Extra {
		if _, exists := payload[key]; exists {
			continue
		}
		payload[key] = value
	}
	return json.Marshal(payload)
}

type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content,omitempty"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id,omitempty"`
	Type     string       `json:"type,omitempty"`
	Function FunctionCall `json:"function,omitempty"`
}

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

type Tool struct {
	Type     string       `json:"type"`
	Function ToolFunction `json:"function"`
}

type ToolFunction struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"`
}

type StreamEvent struct {
	Content      string
	ToolCalls    []ToolCallDelta
	FinishReason string
}

type ToolCallDelta struct {
	Index     int
	ID        string
	Type      string
	Name      string
	Arguments string
}