package ollama

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
)

type generateRequest struct {
	Model  string `json:"model"`
	Prompt string `json:"prompt"`
	Stream bool   `json:"stream"`
	Think  *bool  `json:"think,omitempty"`
}

type generateResponse struct {
	Response string `json:"response"`
}

type listModelsResponse struct {
	Models []struct {
		Name string `json:"name"`
	} `json:"models"`
}

// Client implements the LLM driven port with the Ollama generate API.
type Client struct {
	baseURL string
	http    *http.Client
}

// New constructs an Ollama client with OTel-instrumented transport.
func New(baseURL string, httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = &http.Client{Transport: otelhttp.NewTransport(http.DefaultTransport)}
	}
	if httpClient.Transport == nil {
		httpClient.Transport = otelhttp.NewTransport(http.DefaultTransport)
	}
	return &Client{baseURL: baseURL, http: httpClient}
}

// Generate sends one non-streaming prompt to Ollama.
func (c *Client) Generate(ctx context.Context, model, prompt string) (string, error) {
	reqBody := generateRequest{
		Model:  model,
		Prompt: prompt,
		Stream: false,
	}
	if isThinkingModel(model) {
		disabled := false
		reqBody.Think = &disabled
	}

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

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/generate", bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("build ollama request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return "", fmt.Errorf("call ollama: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("ollama returned status %s", resp.Status)
	}

	var out generateResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", fmt.Errorf("decode ollama response: %w", err)
	}
	return out.Response, nil
}

// ListModels returns the installed model names from Ollama.
func (c *Client) ListModels(ctx context.Context) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/api/tags", nil)
	if err != nil {
		return nil, fmt.Errorf("build ollama list models request: %w", err)
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("list ollama models: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("ollama returned status %s", resp.Status)
	}

	var out listModelsResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, fmt.Errorf("decode ollama list models response: %w", err)
	}

	names := make([]string, 0, len(out.Models))
	for _, model := range out.Models {
		if model.Name != "" {
			names = append(names, model.Name)
		}
	}
	return names, nil
}

func isThinkingModel(name string) bool {
	return strings.HasPrefix(name, "qwen3")
}