package litellm

// Model implements google.golang.org/adk/model.LLM against any
// OpenAI-compatible endpoint (LiteLLM, Ollama, vLLM, etc.).
//
// The official Go ADK v1.x ships only Gemini adapters. This adapter
// implements the official interface directly via net/http — no extra deps.

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"iter"
	"net/http"

	adkmodel "google.golang.org/adk/model"
	"google.golang.org/genai"
)

// Model is an ADK-compatible LLM backed by an OpenAI-compatible endpoint.
type Model struct {
	name    string
	baseURL string
	apiKey  string
	client  *http.Client
}

// New creates an OpenAI-compatible ADK model.
// name is the model identifier sent in requests (e.g. "berget/llama-3.3-70b").
// baseURL is the API base without path (e.g. "https://llm-api.d-ma.be/v1").
func New(name, baseURL, apiKey string) *Model {
	return &Model{name: name, baseURL: baseURL, apiKey: apiKey, client: &http.Client{}}
}

func (m *Model) Name() string { return m.name }

// --- OpenAI wire types (minimal subset ADK uses) ---

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

type oaiToolCall struct {
	ID       string          `json:"id"`
	Type     string          `json:"type"`
	Function oaiFunctionCall `json:"function"`
}

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

type oaiTool struct {
	Type     string         `json:"type"`
	Function oaiFunctionDef `json:"function"`
}

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

type oaiRequest struct {
	Model    string       `json:"model"`
	Messages []oaiMessage `json:"messages"`
	Tools    []oaiTool    `json:"tools,omitempty"`
}

type oaiChoice struct {
	Message      oaiMessage `json:"message"`
	FinishReason string     `json:"finish_reason"`
}

type oaiResponse struct {
	Choices []oaiChoice `json:"choices"`
	Error   *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

func (m *Model) GenerateContent(ctx context.Context, req *adkmodel.LLMRequest, _ bool) iter.Seq2[*adkmodel.LLMResponse, error] {
	return func(yield func(*adkmodel.LLMResponse, error) bool) {
		msgs := contentsToMessages(req.Contents)
		tools := adk2oaiTools(req)

		oaiReq := oaiRequest{Model: m.name, Messages: msgs, Tools: tools}

		body, err := json.Marshal(oaiReq)
		if err != nil {
			yield(nil, fmt.Errorf("marshal: %w", err))
			return
		}

		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
			m.baseURL+"/chat/completions", bytes.NewReader(body))
		if err != nil {
			yield(nil, fmt.Errorf("new request: %w", err))
			return
		}
		httpReq.Header.Set("Content-Type", "application/json")
		httpReq.Header.Set("Authorization", "Bearer "+m.apiKey)

		resp, err := m.client.Do(httpReq)
		if err != nil {
			yield(nil, fmt.Errorf("http: %w", err))
			return
		}
		defer resp.Body.Close()

		raw, err := io.ReadAll(resp.Body)
		if err != nil {
			yield(nil, fmt.Errorf("read body: %w", err))
			return
		}
		if resp.StatusCode != http.StatusOK {
			yield(nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(raw)))
			return
		}

		var oaiResp oaiResponse
		if err := json.Unmarshal(raw, &oaiResp); err != nil {
			yield(nil, fmt.Errorf("unmarshal: %w", err))
			return
		}
		if oaiResp.Error != nil {
			yield(nil, fmt.Errorf("api error: %s", oaiResp.Error.Message))
			return
		}
		if len(oaiResp.Choices) == 0 {
			yield(nil, fmt.Errorf("no choices in response"))
			return
		}

		content := oaiChoiceToContent(oaiResp.Choices[0])
		yield(&adkmodel.LLMResponse{Content: content, TurnComplete: true}, nil)
	}
}

func contentsToMessages(contents []*genai.Content) []oaiMessage {
	var msgs []oaiMessage
	for _, c := range contents {
		if c == nil {
			continue
		}
		var textBuf bytes.Buffer
		var toolCalls []oaiToolCall

		for _, p := range c.Parts {
			if p == nil {
				continue
			}
			if p.Text != "" {
				textBuf.WriteString(p.Text)
			}
			if p.FunctionCall != nil {
				argBytes, _ := json.Marshal(p.FunctionCall.Args)
				toolCalls = append(toolCalls, oaiToolCall{
					ID:   p.FunctionCall.ID,
					Type: "function",
					Function: oaiFunctionCall{
						Name:      p.FunctionCall.Name,
						Arguments: string(argBytes),
					},
				})
			}
			if p.FunctionResponse != nil {
				respBytes, _ := json.Marshal(p.FunctionResponse.Response)
				msgs = append(msgs, oaiMessage{
					Role:       "tool",
					Content:    string(respBytes),
					ToolCallID: p.FunctionResponse.ID,
				})
			}
		}

		if len(toolCalls) > 0 || textBuf.Len() > 0 {
			msg := oaiMessage{Role: c.Role}
			if c.Role == "model" {
				msg.Role = "assistant"
			}
			msg.Content = textBuf.String()
			msg.ToolCalls = toolCalls
			msgs = append(msgs, msg)
		}
	}
	return msgs
}

func adk2oaiTools(req *adkmodel.LLMRequest) []oaiTool {
	if len(req.Tools) == 0 {
		return nil
	}
	var tools []oaiTool
	for name, def := range req.Tools {
		raw, _ := json.Marshal(def)
		var m map[string]json.RawMessage
		_ = json.Unmarshal(raw, &m)
		var desc string
		if d, ok := m["description"]; ok {
			_ = json.Unmarshal(d, &desc)
		}
		params := m["parameters"]
		// Some endpoints (e.g. Berget) reject null parameters for zero-arg tools.
		if len(params) == 0 || string(params) == "null" {
			params = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		tools = append(tools, oaiTool{
			Type: "function",
			Function: oaiFunctionDef{
				Name:        name,
				Description: desc,
				Parameters:  params,
			},
		})
	}
	return tools
}

func oaiChoiceToContent(choice oaiChoice) *genai.Content {
	msg := choice.Message
	var parts []*genai.Part

	if msg.Content != "" {
		parts = append(parts, &genai.Part{Text: msg.Content})
	}
	for _, tc := range msg.ToolCalls {
		var args map[string]any
		_ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
		parts = append(parts, &genai.Part{
			FunctionCall: &genai.FunctionCall{
				ID:   tc.ID,
				Name: tc.Function.Name,
				Args: args,
			},
		})
	}

	role := msg.Role
	if role == "assistant" {
		role = "model"
	}
	return &genai.Content{Role: role, Parts: parts}
}