package models

import (
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"github.com/ardanlabs/kronk/sdk/tools/backend"
	"go.yaml.in/yaml/v2"
)

// File provides information about a model.
type File struct {
	ID                   string
	OwnedBy              string
	ModelFamily          string
	TokenizerFingerprint string
	Size                 int64
	Modified             time.Time
	Validated            bool
	HasProjection        bool
}

// Files returns all the models in the model directory.
func (m *Models) Files() ([]File, error) {
	var list []File

	index := m.loadIndex()

	for modelID, mp := range index {
		if len(mp.ModelFiles) == 0 {
			continue
		}

		var totalSize int64
		var modified time.Time

		for _, f := range mp.ModelFiles {
			info, err := os.Stat(f)
			if err != nil {
				return nil, fmt.Errorf("stat: %w", err)
			}

			totalSize += info.Size()
			if info.ModTime().After(modified) {
				modified = info.ModTime()
			}
		}

		modelPath := strings.TrimPrefix(mp.ModelFiles[0], m.modelsPath)
		modelPath = strings.TrimPrefix(modelPath, string(filepath.Separator))
		parts := strings.Split(modelPath, string(filepath.Separator))

		var ownedBy string
		var modelFamily string

		if len(parts) > 0 {
			ownedBy = parts[0]
		}

		if len(parts) > 1 {
			modelFamily = parts[1]
		}

		mf := File{
			ID:                   modelID,
			OwnedBy:              ownedBy,
			ModelFamily:          modelFamily,
			TokenizerFingerprint: mp.TokenizerFingerprint,
			Size:                 totalSize,
			Modified:             modified,
			Validated:            mp.Validated,
			HasProjection:        mp.ProjFile != "",
		}

		list = append(list, mf)
	}

	slices.SortFunc(list, func(a, b File) int {
		return strings.Compare(strings.ToLower(a.ID), strings.ToLower(b.ID))
	})

	return list, nil
}

// retrieveFile finds the model and returns the model file information.
func (m *Models) retrieveFile(modelID string) (File, error) {
	if modelID == "" {
		return File{}, fmt.Errorf("retrieve-file: missing model id")
	}

	mp, err := m.FullPath(modelID)
	if err != nil {
		return File{}, fmt.Errorf("retrieve-file: unable to retrieve path: %w", err)
	}

	if len(mp.ModelFiles) == 0 {
		return File{}, fmt.Errorf("retrieve-file: no model files found")
	}

	var totalSize int64
	var modified time.Time

	for _, f := range mp.ModelFiles {
		info, err := os.Stat(f)
		if err != nil {
			return File{}, fmt.Errorf("stat: %w", err)
		}

		totalSize += info.Size()
		if info.ModTime().After(modified) {
			modified = info.ModTime()
		}
	}

	modelPath := strings.TrimPrefix(mp.ModelFiles[0], m.modelsPath)
	modelPath = strings.TrimPrefix(modelPath, string(filepath.Separator))
	parts := strings.Split(modelPath, string(filepath.Separator))

	var ownedBy string
	var modelFamily string

	if len(parts) > 0 {
		ownedBy = parts[0]
	}

	if len(parts) > 1 {
		modelFamily = parts[1]
	}

	mf := File{
		ID:          modelID,
		OwnedBy:     ownedBy,
		ModelFamily: modelFamily,
		Size:        totalSize,
		Modified:    modified,
	}

	return mf, nil
}

// =============================================================================

// FileInfo provides all the model details.
type FileInfo struct {
	ID          string
	Object      string
	ModelFamily string
	Size        int64
	Created     int64
	OwnedBy     string
}

// FileInformation provides details for the specified model.
func (m *Models) FileInformation(modelID string) (FileInfo, error) {
	modelID, _, _ = strings.Cut(modelID, "/")

	mf, err := m.retrieveFile(modelID)
	if err != nil {
		return FileInfo{}, fmt.Errorf("retrieve-info: unable to get model file information: %w", err)
	}

	mi := FileInfo{
		ID:          modelID,
		Object:      "model",
		ModelFamily: mf.ModelFamily,
		Size:        mf.Size,
		Created:     mf.Modified.UnixMilli(),
		OwnedBy:     mf.OwnedBy,
	}

	return mi, nil
}

// =============================================================================

// Path returns file path information about a model. It is an alias for
// backend.ModelPath so cross-backend code can consume the same value
// type returned by every backend's Catalog implementation.
type Path = backend.ModelPath

// FullPath locates the physical location on disk and returns the full path.
//
// The index is keyed by the bare model name (e.g. "Qwen3-8B-Q8_0"). Callers
// may pass that bare name, an "<org>/<model>" pair (e.g. "Qwen/Qwen3-8B-Q8_0"),
// a "<model>/<variant>" pair (e.g. "Qwen3-8B-Q8_0/IMC"), or the full
// "<org>/<model>/<variant>" form. This resolver tries each segment in the
// most-specific-first order until it finds a key that exists in the index.
func (m *Models) FullPath(modelID string) (Path, error) {
	index := m.loadIndex()

	for _, key := range fullPathLookupKeys(modelID) {
		if mp, ok := index[key]; ok {
			return mp, nil
		}
	}

	return Path{}, fmt.Errorf("retrieve-path: model %q not found", modelID)
}

// LookupFile resolves a model identifier to its catalog File entry using
// the same precedence rules as FullPath. The identifier may be the bare
// model name ("Qwen3-8B-Q8_0"), an "<org>/<model>" pair, a
// "<model>/<variant>" pair, or the full "<org>/<model>/<variant>" form.
// Returns the matching File and true, or the zero value and false.
func (m *Models) LookupFile(modelID string) (File, bool) {
	files, err := m.Files()
	if err != nil {
		return File{}, false
	}

	byID := make(map[string]File, len(files))
	for _, f := range files {
		byID[f.ID] = f
	}

	for _, key := range fullPathLookupKeys(modelID) {
		if f, ok := byID[key]; ok {
			return f, true
		}
	}

	return File{}, false
}

// MustFullPath finds a model and panics if the model was not found. This
// should only be used for testing.
func (m *Models) MustFullPath(modelID string) Path {
	fi, err := m.FullPath(modelID)
	if err != nil {
		panic(err.Error())
	}

	return fi
}

// fullPathLookupKeys returns the index-key candidates to try for the given
// modelID, in order of preference. The bare-model segment is preferred over
// org/variant segments so that ambiguous two-segment input is resolved
// against actual index content.
func fullPathLookupKeys(modelID string) []string {
	parts := strings.Split(modelID, "/")

	switch len(parts) {
	case 1:
		return []string{parts[0]}
	case 2:
		// Could be "<model>/<variant>" or "<org>/<model>". Try the leading
		// segment first to preserve the original variant-strip semantics,
		// then fall back to the trailing segment for org-prefixed input.
		return []string{parts[0], parts[1]}
	default:
		// "<org>/<model>/<variant>..." — the bare model name lives in the
		// middle. Prefer that, then fall back to the leading and trailing
		// segments.
		return []string{parts[1], parts[0], parts[len(parts)-1]}
	}
}

// =============================================================================

// LoadIndex returns the catalog index.
func (m *Models) loadIndex() map[string]Path {
	m.biMutex.Lock()
	defer m.biMutex.Unlock()

	indexPath := filepath.Join(m.modelsPath, indexFile)

	data, err := os.ReadFile(indexPath)
	if err != nil {
		return make(map[string]Path)
	}

	var index map[string]Path
	if err := yaml.Unmarshal(data, &index); err != nil {
		return make(map[string]Path)
	}

	return index
}
