// Package backend — stateless one-shot operations.
// These functions provide request/response LLM access without session state.
package backend

import (
	"context"
	"fmt"
	"os"
	"strings"
)

// GenerateRequest is the input for a stateless one-shot generation.
type GenerateRequest struct {
	Prompt  string `json:"prompt"`
	System  string `json:"system"`
	Backend string `json:"backend"`
	Model   string `json:"model"`
}

// Generate performs a stateless one-shot LLM generation with no session context.
// Returns the full model response text.
func Generate(ctx context.Context, req GenerateRequest) (string, error) {
	if req.Prompt == "" {
		return "", fmt.Errorf("empty prompt")
	}
	be, err := NewWithName(req.Backend)
	if err != nil {
		return "", err
	}
	if req.Model != "" {
		be.SetModel(req.Model)
	}
	var msgs []Message
	if req.System != "" {
		msgs = append(msgs, Message{Role: "system", Content: req.System})
	}
	msgs = append(msgs, Message{Role: "user", Content: req.Prompt})
	ch, err := be.ChatStream(ctx, msgs, nil, GenerationParams{})
	if err != nil {
		return "", err
	}
	var sb strings.Builder
	for ev := range ch {
		sb.WriteString(ev.Content)
	}
	return sb.String(), nil
}

// RouteRequest is the input for model routing.
type RouteRequest struct {
	Task    string `json:"task"`
	Backend string `json:"backend"` // optional: only consider models from this backend
}

// RouteResult is the output of model routing.
type RouteResult struct {
	Backend string
	Model   string
}

// Route classifies a task description and selects the most appropriate model
// from the actually available models. Fetches the model list from configured
// backends, presents them to a cheap classifier, and returns the selection.
func Route(ctx context.Context, req RouteRequest) (RouteResult, error) {
	if req.Task == "" {
		return RouteResult{}, fmt.Errorf("provide a task description")
	}

	// Classifier backend/model
	classifyBackend := os.Getenv("OLLIE_ROUTE_BACKEND")
	if classifyBackend == "" {
		classifyBackend = "ollama"
	}
	classifyModel := os.Getenv("OLLIE_ROUTE_MODEL")
	if classifyModel == "" {
		classifyModel = "qwen3:8b"
	}
	cbe, err := NewWithName(classifyBackend)
	if err != nil {
		return RouteResult{}, fmt.Errorf("route backend: %w", err)
	}
	cbe.SetModel(classifyModel)

	// Fetch available models
	type modelEntry struct {
		backend string
		model   string
	}
	var entries []modelEntry
	var modelList strings.Builder

	backends := Backends()
	if req.Backend != "" {
		backends = []string{req.Backend}
	}
	for _, name := range backends {
		b, err := NewWithName(name)
		if err != nil {
			continue
		}
		for _, m := range b.Models(ctx) {
			entries = append(entries, modelEntry{backend: name, model: m})
			fmt.Fprintf(&modelList, "%d. %s/%s\n", len(entries), name, m)
		}
	}
	if len(entries) == 0 {
		return RouteResult{}, fmt.Errorf("no models available")
	}

	classifyPrompt := fmt.Sprintf(`You are a task router. Given a task and a numbered list of available models, select the single best model for the task.

Available models:
%s
Reply with ONLY the number of the best model. Nothing else.

Task: %s`, modelList.String(), req.Task)

	msgs := []Message{
		{Role: "user", Content: classifyPrompt},
	}
	ch, err := cbe.ChatStream(ctx, msgs, nil, GenerationParams{})
	if err != nil {
		return RouteResult{}, fmt.Errorf("route classification failed: %w", err)
	}
	var sb strings.Builder
	for ev := range ch {
		sb.WriteString(ev.Content)
	}
	result := strings.TrimSpace(sb.String())

	// Parse the number from the response
	var choice int
	fmt.Sscanf(result, "%d", &choice)
	if choice < 1 || choice > len(entries) {
		// Default to first model if parsing fails
		choice = 1
	}
	picked := entries[choice-1]

	return RouteResult{Backend: picked.backend, Model: picked.model}, nil
}