package engine

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"mm/internal/config"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"
	"github.com/firebase/genkit/go/plugins/compat_oai"
)

// openaiCompat is the Provider for every server that speaks the OpenAI
// chat-completions protocol. Docker Model Runner and llama-server are two
// parametrisations of it; Ollama's /v1, vLLM and the remote APIs would be more
// entries in init() below — the point of keeping the struct data-only.
type openaiCompat struct {
	name string

	baseURL  string // used when the config gives none
	fallback string // tried when baseURL does not answer; "" = no fallback

	// legacyBaseURLEnv keeps DMR_BASE_URL working for the `dmr` provider: the
	// variable is in every README and in the sandbox start scripts. The
	// provider-neutral AGENT_BASE_URL wins over it.
	legacyBaseURLEnv string

	// apiKeyEnv names the variable holding the key; keyRequired says whether an
	// empty one is an error (a paid API) or just means "this server does not
	// check" (llama-server without --api-key). dummyKey is sent in that case
	// because the OpenAI client wants SOMETHING in the header.
	apiKeyEnv   string
	keyRequired bool
	dummyKey    string

	// Words for Explain: how to start the server, how to get the model.
	startHint    string
	notFoundHint string // %s = model

	// probe is the server-specific part of Probe; nil = reachability only.
	probe func(ctx context.Context, b Backend) Info
}

func init() {
	register(&openaiCompat{
		name:             "dmr",
		baseURL:          "http://localhost:12434/engines/v1",
		fallback:         "http://host.docker.internal:12434/engines/v1",
		legacyBaseURLEnv: "DMR_BASE_URL",
		dummyKey:         "not-needed", // DMR ignores the Authorization header
		startHint:        "is Docker Model Runner running? (docker model status)",
		notFoundHint:     "docker model pull %s",
		probe:            probeModels,
	})
	register(&openaiCompat{
		name:    "llamacpp",
		baseURL: "http://127.0.0.1:8080/v1", // llama-server's default host:port, plus the OpenAI prefix
		// llama-server only checks a key when started with --api-key; most
		// local set-ups are not, so an unset LLAMA_API_KEY is normal.
		apiKeyEnv:    "LLAMA_API_KEY",
		dummyKey:     "not-needed",
		startHint:    "start it: llama-server -m <model.gguf> --jinja --port 8080",
		notFoundHint: "llama-server serves one model; check baseUrl ends with /v1 (model: %s)",
		probe:        probeLlamaCpp,
	})
}

func (p *openaiCompat) Name() string { return p.name }

// Resolve applies the same precedence the `dmr` package used for its URL —
// environment, then config, then fallback if the first does not answer — and
// extends it to the key and the model.
func (p *openaiCompat) Resolve(cfg config.Config) (Backend, error) {
	b := Backend{Provider: p.name, Model: cfg.Model, ContextWindow: cfg.ContextWindow}

	// An explicit environment override is taken as-is, without the fallback
	// probe: whoever set it knows where the server is.
	env := os.Getenv("AGENT_BASE_URL")
	if env == "" && p.legacyBaseURLEnv != "" {
		env = os.Getenv(p.legacyBaseURLEnv)
	}
	switch {
	case env != "":
		b.BaseURL = env
	default:
		base := cfg.BaseURL
		if base == "" {
			base = p.baseURL
		}
		// nil = the key is absent from the file, use the provider's fallback;
		// an explicit "" disables it (probe.yaml does that to pin the fake
		// engine). The distinction needs the pointer.
		fallback := p.fallback
		if cfg.Fallback != nil {
			fallback = *cfg.Fallback
		}
		b.BaseURL = base
		if fallback != "" && !reachable(base) {
			b.BaseURL = fallback
		}
	}

	keyEnv := cfg.APIKeyEnv
	if keyEnv == "" {
		keyEnv = p.apiKeyEnv
	}
	b.APIKeyEnv = keyEnv
	if keyEnv != "" {
		b.APIKey = os.Getenv(keyEnv)
	}
	if b.APIKey == "" {
		if p.keyRequired {
			return b, fmt.Errorf("%s: %s is not set (the key never goes in the YAML: export it)", p.name, keyEnv)
		}
		b.APIKey = p.dummyKey
	}
	return b, nil
}

// Open is the former dmr.Init with the constants replaced by the Backend.
func (p *openaiCompat) Open(ctx context.Context, b Backend) (*genkit.Genkit, string, error) {
	client := &compat_oai.OpenAICompatible{
		Provider: p.name, // model prefix: "<provider>/<id>"
		BaseURL:  b.BaseURL,
		APIKey:   b.APIKey,
	}
	g := genkit.Init(ctx, genkit.WithPlugins(client))
	client.DefineModel(p.name, b.Model, ai.ModelOptions{
		Label:    b.Model,
		Supports: &compat_oai.Multimodal, // enables Tools, Multiturn, SystemRole…
	})
	return g, p.name + "/" + b.Model, nil
}

// Probe: reachability on /models — the one endpoint every OpenAI-compatible
// server has, and the one the fallback logic already relied on — then whatever
// the server-specific probe adds. The config's context window always wins over
// the server's: the operator knows what they started it with.
func (p *openaiCompat) Probe(ctx context.Context, b Backend) Info {
	var info Info
	if p.probe != nil {
		info = p.probe(ctx, b)
	} else {
		info.Reachable = reachable(b.BaseURL)
	}
	if !info.Reachable {
		info.Warnings = append(info.Warnings,
			fmt.Sprintf("%s: nothing answers at %s — %s", p.name, b.BaseURL, p.startHint))
	}
	if b.ContextWindow > 0 {
		info.ContextWindow, info.ContextSource = b.ContextWindow, "config"
	}
	return info
}

// probeClient: two seconds, like the old reachable(). A local server answers in
// milliseconds; a longer wait only delays the prompt when it is down.
var probeClient = &http.Client{Timeout: 2 * time.Second}

// reachable reports whether an OpenAI-compatible endpoint answers on /models.
func reachable(baseURL string) bool {
	resp, err := probeClient.Get(baseURL + "/models")
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == http.StatusOK
}

// getJSON fetches and decodes one endpoint; false when anything failed. Probes
// do not distinguish "down" from "not that kind of server": both mean "no info".
func getJSON(url string, out any) bool {
	resp, err := probeClient.Get(url)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return false
	}
	return json.NewDecoder(resp.Body).Decode(out) == nil
}

// probeModels (DMR): lists /models and warns when the configured model is not
// among them. Measured on DMR: asking for a model that is not pulled returns an
// error only once the request is sent — after the spinner has run a while.
func probeModels(_ context.Context, b Backend) Info {
	var list struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	info := Info{Reachable: getJSON(b.BaseURL+"/models", &list)}
	if !info.Reachable || len(list.Data) == 0 {
		return info
	}
	for _, m := range list.Data {
		if m.ID == b.Model {
			return info
		}
	}
	info.Warnings = append(info.Warnings,
		fmt.Sprintf("%s: model %q is not in the server's list — docker model pull %s", b.Provider, b.Model, b.Model))
	return info
}

// probeLlamaCpp: llama-server exposes GET /props, one level above /v1, with the
// context size it actually serves (`default_generation_settings.n_ctx`) — not
// the one the model was trained with, which is what /v1/models reports. The
// served size is the number the context-compression work needs.
func probeLlamaCpp(_ context.Context, b Backend) Info {
	var props struct {
		Settings struct {
			NCtx int `json:"n_ctx"`
		} `json:"default_generation_settings"`
		ChatTemplate string `json:"chat_template"`
	}
	root := strings.TrimSuffix(strings.TrimSuffix(b.BaseURL, "/"), "/v1")
	info := Info{Reachable: getJSON(root+"/props", &props)}
	if !info.Reachable {
		// /props missing but /models answering: some other OpenAI-compatible
		// server is on that port. Reachable, but nothing llama.cpp to report.
		info.Reachable = reachable(b.BaseURL)
		return info
	}
	if props.Settings.NCtx > 0 {
		info.ContextWindow, info.ContextSource = props.Settings.NCtx, "/props"
	}
	if props.ChatTemplate == "" {
		info.Warnings = append(info.Warnings,
			"llamacpp: the server reports no chat template; tool calls need one (and --jinja)")
	}
	return info
}