package run

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/ardanlabs/kronk/sdk/kronk"
	"github.com/ardanlabs/kronk/sdk/kronk/model"
	"github.com/ardanlabs/kronk/sdk/tools/defaults"
	"github.com/ardanlabs/kronk/sdk/tools/libs"
	"github.com/ardanlabs/kronk/sdk/tools/models"
)

type Config struct {
	ModelName string
	BasePath  string

	// Model configuration.
	JinjaFile      string
	ContextWindow  int
	FlashAttention string
	NGpuLayers     int
	Devices        string
	MainGPU        int
	TensorSplit    string
	SplitMode      string
	CacheTypeK     string
	CacheTypeV     string
	NBatch         int
	NUBatch        int

	// Sampling parameters.
	MaxTokens        int
	Temperature      float64
	TopP             float64
	TopK             int
	MinP             float64
	RepeatPenalty    float64
	FrequencyPenalty float64
	PresencePenalty  float64
	EnableThinking   string
	ReasoningEffort  string
}

func runChat(cfg Config) error {
	mp, err := installSystem(cfg)
	if err != nil {
		return fmt.Errorf("run: unable to install system: %w", err)
	}

	krn, err := newKronk(mp, cfg)
	if err != nil {
		return fmt.Errorf("unable to init kronk: %w", err)
	}

	defer func() {
		fmt.Println("\nUnloading Kronk")
		if err := krn.Unload(context.Background()); err != nil {
			fmt.Printf("run: failed to unload model: %v\n", err)
		}
	}()

	return chat(krn, cfg)
}

func installSystem(cfg Config) (models.Path, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 25*time.Minute)
	defer cancel()

	lbs, err := libs.New(
		libs.WithBasePath(cfg.BasePath),
		libs.WithVersion(defaults.LibVersion("")),
	)
	if err != nil {
		return models.Path{}, err
	}

	if _, err := lbs.Download(ctx, kronk.FmtLogger); err != nil {
		return models.Path{}, fmt.Errorf("unable to install llama.cpp: %w", err)
	}

	mdls, err := models.NewWithPaths(cfg.BasePath)
	if err != nil {
		return models.Path{}, fmt.Errorf("unable to create models system: %w", err)
	}

	fmt.Println("Resolving model:", cfg.ModelName)

	mp, err := mdls.Download(ctx, kronk.FmtLogger, cfg.ModelName)
	if err != nil {
		return models.Path{}, fmt.Errorf("unable to install model: %w", err)
	}

	return mp, nil
}

func newKronk(mp models.Path, runCfg Config) (*kronk.Kronk, error) {
	if err := kronk.Init(); err != nil {
		return nil, fmt.Errorf("unable to init kronk: %w", err)
	}

	cfg := model.Config{
		ModelFiles: mp.ModelFiles,
	}

	if runCfg.JinjaFile != "" {
		cfg.JinjaFile = runCfg.JinjaFile
	}

	if runCfg.ContextWindow > 0 {
		cfg.PtrContextWindow = new(runCfg.ContextWindow)
	}

	if runCfg.NBatch > 0 {
		cfg.PtrNBatch = new(runCfg.NBatch)
	}

	if runCfg.NUBatch > 0 {
		cfg.PtrNUBatch = new(runCfg.NUBatch)
	}

	if runCfg.NGpuLayers != 0 {
		v := runCfg.NGpuLayers
		cfg.PtrNGpuLayers = &v
	}

	if runCfg.Devices != "" {
		cfg.Devices = strings.Split(runCfg.Devices, ",")
	}

	if runCfg.MainGPU >= 0 {
		cfg.PtrMainGPU = &runCfg.MainGPU
	}

	if runCfg.TensorSplit != "" {
		parts := strings.Split(runCfg.TensorSplit, ",")
		splits := make([]float32, len(parts))
		for i, p := range parts {
			v, err := strconv.ParseFloat(strings.TrimSpace(p), 32)
			if err != nil {
				return nil, fmt.Errorf("invalid tensor-split value %q: %w", p, err)
			}
			splits[i] = float32(v)
		}
		cfg.TensorSplit = splits
	}

	if runCfg.SplitMode != "" {
		sm, err := model.ParseSplitMode(runCfg.SplitMode)
		if err != nil {
			return nil, fmt.Errorf("invalid split-mode: %w", err)
		}
		cfg.PtrSplitMode = &sm
	}

	if runCfg.CacheTypeK != "" {
		ct, err := model.ParseGGMLType(runCfg.CacheTypeK)
		if err != nil {
			return nil, fmt.Errorf("invalid cache-type-k: %w", err)
		}
		cfg.CacheTypeK = ct
	}

	if runCfg.CacheTypeV != "" {
		ct, err := model.ParseGGMLType(runCfg.CacheTypeV)
		if err != nil {
			return nil, fmt.Errorf("invalid cache-type-v: %w", err)
		}
		cfg.CacheTypeV = ct
	}

	switch runCfg.FlashAttention {
	case "on":
		cfg.FlashAttention = model.FlashAttentionEnabled
	case "off":
		cfg.FlashAttention = model.FlashAttentionDisabled
	case "auto":
		cfg.FlashAttention = model.FlashAttentionAuto
	}

	krn, err := kronk.New(model.WithConfig(cfg))
	if err != nil {
		return nil, fmt.Errorf("unable to create inference model: %w", err)
	}

	fmt.Print("- system info:\n\t")
	for k, v := range krn.SystemInfo() {
		fmt.Printf("%s:%v, ", k, v)
	}
	fmt.Println()

	fmt.Println("- contextWindow  :", krn.ModelConfig().ContextWindow())
	fmt.Printf("- k/v            : %s/%s\n", krn.ModelConfig().CacheTypeK, krn.ModelConfig().CacheTypeV)
	fmt.Println("- flashAttention :", krn.ModelConfig().FlashAttention)
	fmt.Println("- nBatch         :", krn.ModelConfig().NBatch())
	fmt.Println("- nuBatch        :", krn.ModelConfig().NUBatch())
	fmt.Println("- modelType      :", krn.ModelInfo().Type)
	fmt.Println("- embeddings     :", krn.ModelInfo().IsEmbedModel)
	fmt.Println("- isGPT          :", krn.ModelInfo().IsGPTModel)
	fmt.Println("- template       :", krn.ModelInfo().Template.FileName)
	fmt.Println("- grammar        :", krn.ModelConfig().DefaultParams.Grammar != "")
	fmt.Println("- nSeqMax        :", krn.ModelConfig().NSeqMax())
	fmt.Println("- vramTotal      :", krn.ModelInfo().VRAMTotal/(1024*1024), "MiB")
	fmt.Println("- slotMemory     :", krn.ModelInfo().SlotMemory/(1024*1024), "MiB")
	fmt.Println("- modelSize      :", krn.ModelInfo().Size/(1000*1000), "MB")
	fmt.Println("- imc            :", krn.ModelConfig().IncrementalCache())
	if n := krn.ModelConfig().PtrNGpuLayers; n != nil {
		fmt.Println("- nGPULayers     :", *n)
	} else {
		fmt.Println("- nGPULayers     : all")
	}

	if runCfg.JinjaFile != "" {
		fmt.Println("- jinjaFile      :", runCfg.JinjaFile)
	}
	if runCfg.Devices != "" {
		fmt.Println("- devices        :", runCfg.Devices)
	}
	if runCfg.MainGPU >= 0 {
		fmt.Println("- mainGPU        :", runCfg.MainGPU)
	}
	if runCfg.TensorSplit != "" {
		fmt.Println("- tensorSplit    :", runCfg.TensorSplit)
	}
	if runCfg.SplitMode != "" {
		fmt.Println("- splitMode      :", runCfg.SplitMode)
	}
	if runCfg.MaxTokens != 0 {
		fmt.Println("- maxTokens      :", runCfg.MaxTokens)
	}
	if runCfg.Temperature != 0 {
		fmt.Println("- temperature    :", runCfg.Temperature)
	}
	if runCfg.TopP != 0 {
		fmt.Println("- topP           :", runCfg.TopP)
	}
	if runCfg.TopK != 0 {
		fmt.Println("- topK           :", runCfg.TopK)
	}
	if runCfg.MinP != 0 {
		fmt.Println("- minP           :", runCfg.MinP)
	}
	if runCfg.RepeatPenalty != 0 {
		fmt.Println("- repeatPen      :", runCfg.RepeatPenalty)
	}
	if runCfg.FrequencyPenalty != 0 {
		fmt.Println("- freqPen        :", runCfg.FrequencyPenalty)
	}
	if runCfg.PresencePenalty != 0 {
		fmt.Println("- presPen        :", runCfg.PresencePenalty)
	}
	if runCfg.EnableThinking != "" {
		fmt.Println("- thinking       :", runCfg.EnableThinking)
	}
	if runCfg.ReasoningEffort != "" {
		fmt.Println("- reasonEffort   :", runCfg.ReasoningEffort)
	}

	return krn, nil
}

func chat(krn *kronk.Kronk, cfg Config) error {
	messages := model.DocumentArray()

	var systemPrompt = `
		You are a helpful AI assistant. You are designed to help users answer
		questions, create content, and provide information in a helpful and
		accurate manner. Always follow the user's instructions carefully and
		respond with clear, concise, and well-structured answers. You are a
		helpful AI assistant. You are designed to help users answer questions,
		create content, and provide information in a helpful and accurate manner.
		Always follow the user's instructions carefully and respond with clear,
		concise, and well-structured answers. You are a helpful AI assistant.
		You are designed to help users answer questions, create content, and
		provide information in a helpful and accurate manner. Always follow the
		user's instructions carefully and respond with clear, concise, and
		well-structured answers.`

	messages = append(messages,
		model.TextMessage(model.RoleSystem, systemPrompt),
	)

	for {
		var err error
		messages, err = userInput(messages)
		if err != nil {
			if errors.Is(err, io.EOF) {
				return nil
			}
			return fmt.Errorf("run:user input: %w", err)
		}

		messages, err = func() ([]model.D, error) {
			ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
			defer cancel()

			d := model.D{
				"messages": messages,
				"tools":    toolDocuments(),
			}

			if cfg.MaxTokens != 0 {
				d["max_tokens"] = cfg.MaxTokens
			}
			if cfg.Temperature != 0 {
				d["temperature"] = cfg.Temperature
			}
			if cfg.TopP != 0 {
				d["top_p"] = cfg.TopP
			}
			if cfg.TopK != 0 {
				d["top_k"] = cfg.TopK
			}
			if cfg.MinP != 0 {
				d["min_p"] = cfg.MinP
			}
			if cfg.RepeatPenalty != 0 {
				d["repeat_penalty"] = cfg.RepeatPenalty
			}
			if cfg.FrequencyPenalty != 0 {
				d["frequency_penalty"] = cfg.FrequencyPenalty
			}
			if cfg.PresencePenalty != 0 {
				d["presence_penalty"] = cfg.PresencePenalty
			}
			if cfg.EnableThinking != "" {
				d["enable_thinking"] = cfg.EnableThinking
			}
			if cfg.ReasoningEffort != "" {
				d["reasoning_effort"] = cfg.ReasoningEffort
			}

			ch, err := performChat(ctx, krn, d)
			if err != nil {
				return nil, fmt.Errorf("run: unable to perform chat: %w", err)
			}

			messages, err = modelResponse(krn, messages, ch)
			if err != nil {
				return nil, fmt.Errorf("run: model response: %w", err)
			}

			return messages, nil
		}()

		if err != nil {
			return fmt.Errorf("run: unable to perform chat: %w", err)
		}
	}
}

func userInput(messages []model.D) ([]model.D, error) {
	fmt.Print("\nUSER> ")

	reader := bufio.NewReader(os.Stdin)

	userInput, err := reader.ReadString('\n')
	if err != nil {
		return messages, fmt.Errorf("unable to read user input: %w", err)
	}

	if userInput == "quit\n" {
		return nil, io.EOF
	}

	messages = append(messages,
		model.TextMessage(model.RoleUser, userInput),
	)

	return messages, nil
}

func toolDocuments() []model.D {
	return model.DocumentArray(
		model.D{
			"type": "function",
			"function": model.D{
				"name":        "get_weather",
				"description": "Get the current weather for a location",
				"parameters": model.D{
					"type": "object",
					"properties": model.D{
						"location": model.D{
							"type":        "string",
							"description": "The location to get the weather for, e.g. San Francisco, CA",
						},
					},
					"required": []any{"location"},
				},
			},
		},
	)
}

func performChat(ctx context.Context, krn *kronk.Kronk, d model.D) (<-chan model.ChatResponse, error) {
	ch, err := krn.ChatStreaming(ctx, d)
	if err != nil {
		return nil, fmt.Errorf("chat streaming: %w", err)
	}

	return ch, nil
}

func modelResponse(krn *kronk.Kronk, messages []model.D, ch <-chan model.ChatResponse) ([]model.D, error) {
	fmt.Print("\nMODEL> ")

	var reasoning bool
	var lr model.ChatResponse
	var content strings.Builder

loop:
	for resp := range ch {
		lr = resp

		if len(resp.Choices) == 0 {
			continue
		}

		switch resp.Choices[0].FinishReason() {
		case model.FinishReasonError:
			return messages, fmt.Errorf("error from model: %s", resp.Choices[0].Delta.Content)

		case model.FinishReasonStop:
			break loop

		case model.FinishReasonTool:
			fmt.Println()
			if krn.ModelInfo().IsGPTModel {
				fmt.Println()
			}

			fmt.Printf("\u001b[92mModel Asking For Tool Calls:\n\u001b[0m")

			var toolCallDocs []model.D
			for _, tool := range resp.Choices[0].Delta.ToolCalls {
				fmt.Printf("\u001b[92mToolID[%s]: %s(%s)\n\u001b[0m",
					tool.ID,
					tool.Function.Name,
					tool.Function.Arguments,
				)

				argsJSON, _ := json.Marshal(tool.Function.Arguments)
				toolCallDocs = append(toolCallDocs, model.D{
					"id":   tool.ID,
					"type": "function",
					"function": model.D{
						"name":      tool.Function.Name,
						"arguments": string(argsJSON),
					},
				})
			}

			messages = append(messages, model.D{
				"role":       "assistant",
				"tool_calls": toolCallDocs,
			})

			for _, tool := range resp.Choices[0].Delta.ToolCalls {
				messages = append(messages, model.D{
					"role":         "tool",
					"tool_call_id": tool.ID,
					"name":         tool.Function.Name,
					"content":      `{"temperature": "72°F", "condition": "sunny"}`,
				})
			}

			break loop

		default:
			if resp.Choices[0].Delta.Reasoning != "" {
				fmt.Printf("\u001b[91m%s\u001b[0m", resp.Choices[0].Delta.Reasoning)
				reasoning = true
				continue
			}

			if reasoning {
				reasoning = false

				fmt.Println()
				if krn.ModelInfo().IsGPTModel {
					fmt.Println()
				}
			}

			content.WriteString(resp.Choices[0].Delta.Content)
			fmt.Printf("%s", resp.Choices[0].Delta.Content)
		}
	}

	// Append the assistant's response to conversation history.
	if content.Len() > 0 {
		messages = append(messages, model.TextMessage(model.RoleAssistant, content.String()))
	}

	// -------------------------------------------------------------------------

	contextTokens := lr.Usage.PromptTokens + lr.Usage.CompletionTokens
	contextWindow := krn.ModelConfig().ContextWindow()
	percentage := (float64(contextTokens) / float64(contextWindow)) * 100
	of := float32(contextWindow) / float32(1024)

	fmt.Printf("\n\n\u001b[90mInput: %d  Reasoning: %d  Completion: %d  Output: %d  Window: %d (%.0f%% of %.0fK) TPS: %.2f\u001b[0m\n",
		lr.Usage.PromptTokens, lr.Usage.ReasoningTokens, lr.Usage.CompletionTokens, lr.Usage.OutputTokens, contextTokens, percentage, of, lr.Usage.TokensPerSecond)

	return messages, nil
}
