// Package modelprojection renders disposable runtime adapters from the installation model catalog.
package modelprojection

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"

	"github.com/aritmolab/thothii/tools/tht/internal/config"
)

const (
	CatalogFile    = "catalog.json"
	PiModelsFile   = "pi/models.json"
	PiSettingsFile = "pi/settings.json"
	ComposeFile    = "compose.models.yaml"
)

var renameProjectionDirectory = os.Rename

type runtimeCatalog struct {
	SchemaVersion      int                          `json:"schemaVersion"`
	DefaultInteraction string                       `json:"defaultInteraction"`
	Embedding          config.ModelCatalogEmbedding `json:"embedding"`
	Models             []runtimeCatalogModel        `json:"models"`
}

type runtimeCatalogModel struct {
	ID                 string                          `json:"id"`
	Provider           string                          `json:"provider"`
	Model              string                          `json:"model"`
	Label              string                          `json:"label"`
	UpstreamModel      string                          `json:"upstreamModel"`
	Endpoint           *config.ModelEndpoint           `json:"endpoint,omitempty"`
	Authentication     config.ModelAuthentication      `json:"authentication"`
	SessionAdapter     *config.ModelSessionAdapter     `json:"sessionAdapter,omitempty"`
	MetadataAdapter    *config.ModelMetadataAdapter    `json:"metadataAdapter,omitempty"`
	Session            *config.SessionModel            `json:"session,omitempty"`
	MetadataGeneration *config.MetadataGenerationModel `json:"metadataGeneration,omitempty"`
}

type piModels struct {
	Providers map[string]piProvider `json:"providers"`
}

type piProvider struct {
	BaseURL string    `json:"baseUrl"`
	API     string    `json:"api"`
	APIKey  string    `json:"apiKey"`
	Models  []piModel `json:"models"`
}

type piModel struct {
	ID            string                     `json:"id"`
	Name          string                     `json:"name"`
	Reasoning     bool                       `json:"reasoning"`
	Input         []string                   `json:"input,omitempty"`
	Cost          *config.ModelCost          `json:"cost,omitempty"`
	ContextWindow int                        `json:"contextWindow"`
	MaxTokens     int                        `json:"maxTokens"`
	Compat        *config.ModelCompatibility `json:"compat,omitempty"`
}

type piSettings struct {
	DefaultProjectTrust string   `json:"defaultProjectTrust"`
	EnabledModels       []string `json:"enabledModels"`
}

// Render builds every projection in memory so no partial candidate can be published.
func Render(installation config.Installation) (map[string][]byte, error) {
	if err := installation.ModelCatalog.NormalizeDefaults(); err != nil {
		return nil, err
	}
	models := installation.ModelCatalog.RuntimeModels()
	hasMetadata := false
	for _, model := range models {
		hasMetadata = hasMetadata || model.MetadataGeneration != nil
	}
	runtimeModels := make([]runtimeCatalogModel, 0, len(models))
	customProviders := make(map[string]piProvider)
	enabled := make([]string, 0)
	for _, model := range models {
		sessionAdapter, metadataAdapter := model.SessionAdapter, model.MetadataAdapter
		if model.Session == nil {
			sessionAdapter = nil
		}
		if model.MetadataGeneration == nil {
			metadataAdapter = nil
		}
		runtimeModels = append(runtimeModels, runtimeCatalogModel{
			ID: model.ID, Provider: model.Provider, Model: model.Model, Label: model.Label,
			UpstreamModel: model.UpstreamModel, Endpoint: model.Endpoint,
			Authentication: model.Authentication, SessionAdapter: sessionAdapter,
			MetadataAdapter: metadataAdapter, Session: model.Session,
			MetadataGeneration: model.MetadataGeneration,
		})
		if model.Session == nil || (hasMetadata && model.MetadataGeneration == nil) {
			continue
		}
		if model.SessionAdapter == nil || model.SessionAdapter.Mode != "openai_compatible" {
			enabled = append(enabled, model.ID)
			continue
		}
		enabled = append(enabled, model.Provider+"/"+model.UpstreamModel)
		provider := customProviders[model.Provider]
		provider.BaseURL = model.Endpoint.BaseURL
		provider.API = "openai-completions"
		if model.Authentication.Mode == "secret_env" {
			provider.APIKey = "$" + model.Authentication.APIKeyEnv
		} else {
			provider.APIKey = "local"
		}
		provider.Models = append(provider.Models, piModel{
			ID: model.UpstreamModel, Name: model.Label, Reasoning: model.Session.Reasoning,
			Input: model.Session.Input, Cost: model.Session.Cost,
			ContextWindow: model.Session.ContextWindow, MaxTokens: model.Session.MaxTokens,
			Compat: model.Session.Compatibility,
		})
		customProviders[model.Provider] = provider
	}
	for id, provider := range customProviders {
		sort.Slice(provider.Models, func(left, right int) bool { return provider.Models[left].ID < provider.Models[right].ID })
		customProviders[id] = provider
	}
	sort.Strings(enabled)

	catalogBytes, err := marshalJSON(runtimeCatalog{
		SchemaVersion: 2, DefaultInteraction: installation.ModelCatalog.Defaults.Interaction,
		Embedding: installation.ModelCatalog.Embedding, Models: runtimeModels,
	})
	if err != nil {
		return nil, fmt.Errorf("render runtime model catalog: %w", err)
	}
	piModelsBytes, err := marshalJSON(piModels{Providers: customProviders})
	if err != nil {
		return nil, fmt.Errorf("render Pi model projection: %w", err)
	}
	piSettingsBytes, err := marshalJSON(piSettings{DefaultProjectTrust: "always", EnabledModels: enabled})
	if err != nil {
		return nil, fmt.Errorf("render Pi settings projection: %w", err)
	}
	fingerprint := sha256.Sum256(bytes.Join([][]byte{catalogBytes, piModelsBytes, piSettingsBytes}, nil))
	composeBytes := renderCompose(installation, fmt.Sprintf("sha256:%x", fingerprint))
	return map[string][]byte{
		CatalogFile: catalogBytes, PiModelsFile: piModelsBytes,
		PiSettingsFile: piSettingsBytes, ComposeFile: composeBytes,
	}, nil
}

// Generate publishes all adapters as one directory generation. An unchanged generation stays in
// place so active file bind mounts keep their source identity. A failed replacement restores the
// previous directory, so callers never observe a successfully returned mixed generation.
func Generate(installation config.Installation) error {
	artifacts, err := Render(installation)
	if err != nil {
		return err
	}
	target := installation.GeneratedDirectory()
	parent := filepath.Dir(target)
	if err := os.MkdirAll(parent, 0o755); err != nil {
		return fmt.Errorf("create model projection parent: %w", err)
	}
	info, statErr := os.Lstat(target)
	if statErr == nil {
		if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			return fmt.Errorf("current model projection path is not a regular directory")
		}
		if projectionMatches(target, artifacts) {
			return nil
		}
	} else if !os.IsNotExist(statErr) {
		return fmt.Errorf("inspect current model projection generation: %w", statErr)
	}
	candidate, err := os.MkdirTemp(parent, ".model-projections-candidate-*")
	if err != nil {
		return fmt.Errorf("create model projection candidate: %w", err)
	}
	defer func() { _ = os.RemoveAll(candidate) }()
	if err := writeProjectionCandidate(candidate, artifacts); err != nil {
		return err
	}

	if os.IsNotExist(statErr) {
		if err := renameProjectionDirectory(candidate, target); err != nil {
			return fmt.Errorf("publish model projection generation: %w", err)
		}
		return nil
	}
	previous, err := absentTemporaryPath(parent)
	if err != nil {
		return fmt.Errorf("reserve previous model projection generation: %w", err)
	}
	if err := renameProjectionDirectory(target, previous); err != nil {
		return fmt.Errorf("prepare model projection generation replacement: %w", err)
	}
	if err := renameProjectionDirectory(candidate, target); err != nil {
		if restoreErr := renameProjectionDirectory(previous, target); restoreErr != nil {
			return fmt.Errorf("publish model projection generation: %v; restore previous generation: %w", err, restoreErr)
		}
		return fmt.Errorf("publish model projection generation: %w", err)
	}
	_ = os.RemoveAll(previous)
	return nil
}

func projectionMatches(directory string, artifacts map[string][]byte) bool {
	for _, relative := range sortedArtifactPaths(artifacts) {
		path := filepath.Join(directory, filepath.FromSlash(relative))
		info, err := os.Lstat(path)
		if err != nil || !info.Mode().IsRegular() {
			return false
		}
		actual, err := os.ReadFile(path)
		if err != nil || !bytes.Equal(actual, artifacts[relative]) {
			return false
		}
	}
	return true
}

func writeProjectionCandidate(directory string, artifacts map[string][]byte) error {
	if err := os.Chmod(directory, 0o755); err != nil {
		return fmt.Errorf("protect model projection candidate: %w", err)
	}
	for _, relative := range sortedArtifactPaths(artifacts) {
		destination := filepath.Join(directory, filepath.FromSlash(relative))
		if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil {
			return fmt.Errorf("create model projection directory: %w", err)
		}
		if err := os.WriteFile(destination, artifacts[relative], 0o644); err != nil {
			return fmt.Errorf("write model projection candidate: %w", err)
		}
	}
	return nil
}

func absentTemporaryPath(parent string) (string, error) {
	path, err := os.MkdirTemp(parent, ".model-projections-previous-*")
	if err != nil {
		return "", err
	}
	if err := os.Remove(path); err != nil {
		return "", err
	}
	return path, nil
}

// Check returns relative artifact names whose current bytes differ from the catalog projection.
func Check(installation config.Installation) ([]string, error) {
	artifacts, err := Render(installation)
	if err != nil {
		return nil, err
	}
	drift := make([]string, 0)
	for _, relative := range sortedArtifactPaths(artifacts) {
		actual, readErr := os.ReadFile(filepath.Join(installation.GeneratedDirectory(), filepath.FromSlash(relative)))
		if readErr != nil || !bytes.Equal(actual, artifacts[relative]) {
			drift = append(drift, relative)
		}
	}
	return drift, nil
}

func marshalJSON(value any) ([]byte, error) {
	contents, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(contents, '\n'), nil
}

func renderCompose(installation config.Installation, fingerprint string) []byte {
	embedding := installation.ModelCatalog.Embedding
	return []byte(fmt.Sprintf(`services:
  core:
    environment:
      THT_HOST_PLATFORM: %s
      THT_MODEL_CATALOG_FILE: /run/thothii-model-catalog/catalog.json
      THT_MODEL_CATALOG_REVISION: %s
      THT_DEFAULT_SESSION_MODEL: %s
      THT_INTERNAL_EMBEDDING_ID: %s
      THT_INTERNAL_EMBEDDING_MODEL: %s
      THT_INTERNAL_EMBEDDING_DIMENSIONS: %s
    volumes:
      - type: bind
        source: %s
        target: /run/thothii-model-catalog/catalog.json
        read_only: true
      - type: bind
        source: %s
        target: /home/thoth/.pi/agent/models.json
        read_only: true
      - type: bind
        source: %s
        target: /home/thoth/.pi/agent/settings.json
        read_only: true
  workspace-maintenance:
    environment:
      THT_INTERNAL_EMBEDDING_ID: %s
      THT_INTERNAL_EMBEDDING_MODEL: %s
      THT_INTERNAL_EMBEDDING_DIMENSIONS: %s
  embedding-model-init:
    environment:
      OLLAMA_MODEL: %s
`, strconv.Quote(runtime.GOOS), strconv.Quote(fingerprint), strconv.Quote(installation.ModelCatalog.Defaults.Interaction),
		strconv.Quote(embedding.ID), strconv.Quote(embeddingModelName(embedding.ID)), strconv.Quote(strconv.Itoa(embedding.Dimensions)),
		strconv.Quote(installation.GeneratedModelCatalogPath()), strconv.Quote(installation.GeneratedPiModelsPath()),
		strconv.Quote(installation.GeneratedPiSettingsPath()), strconv.Quote(embedding.ID),
		strconv.Quote(embeddingModelName(embedding.ID)), strconv.Quote(strconv.Itoa(embedding.Dimensions)),
		strconv.Quote(embeddingModelName(embedding.ID))))
}

func embeddingModelName(canonical string) string {
	for index := 0; index < len(canonical); index++ {
		if canonical[index] == '/' {
			return canonical[index+1:]
		}
	}
	return canonical
}

func sortedArtifactPaths(artifacts map[string][]byte) []string {
	paths := make([]string, 0, len(artifacts))
	for path := range artifacts {
		paths = append(paths, path)
	}
	sort.Strings(paths)
	return paths
}