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

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"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"
)

type runtimeCatalog struct {
	SchemaVersion             int                          `json:"schemaVersion"`
	DefaultSession            string                       `json:"defaultSession"`
	DefaultMetadataGeneration string                       `json:"defaultMetadataGeneration,omitempty"`
	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) {
	models := installation.ModelCatalog.RuntimeModels()
	runtimeModels := make([]runtimeCatalogModel, 0, len(models))
	customProviders := make(map[string]piProvider)
	enabled := make([]string, 0)
	for _, model := range models {
		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: model.SessionAdapter,
			MetadataAdapter: model.MetadataAdapter, Session: model.Session,
			MetadataGeneration: model.MetadataGeneration,
		})
		if model.Session == nil {
			continue
		}
		enabled = append(enabled, model.ID)
		if model.SessionAdapter == nil || model.SessionAdapter.Mode != "openai_compatible" {
			continue
		}
		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: 1, DefaultSession: installation.ModelCatalog.Defaults.Session,
		DefaultMetadataGeneration: installation.ModelCatalog.Defaults.MetadataGeneration,
		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 atomically replaces each generated adapter after the complete candidate has rendered.
func Generate(installation config.Installation) error {
	artifacts, err := Render(installation)
	if err != nil {
		return err
	}
	paths := sortedArtifactPaths(artifacts)
	for _, relative := range paths {
		destination := filepath.Join(installation.GeneratedDirectory(), filepath.FromSlash(relative))
		// The generated catalog contains references to secret environment variable names, never
		// secret values. Core runs as an unprivileged container user and must be able to traverse
		// the bind-mounted host directories and read the projections.
		if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil {
			return fmt.Errorf("create model projection directory: %w", err)
		}
		if err := os.Chmod(filepath.Dir(destination), 0o755); err != nil {
			return fmt.Errorf("protect model projection directory: %w", err)
		}
		temporary, err := os.CreateTemp(filepath.Dir(destination), ".projection-*")
		if err != nil {
			return fmt.Errorf("create model projection candidate: %w", err)
		}
		temporaryName := temporary.Name()
		published := false
		defer func() {
			if !published {
				_ = os.Remove(temporaryName)
			}
		}()
		if err := temporary.Chmod(0o644); err == nil {
			_, err = temporary.Write(artifacts[relative])
		}
		if closeErr := temporary.Close(); err == nil {
			err = closeErr
		}
		if err != nil {
			return fmt.Errorf("write model projection candidate: %w", err)
		}
		if err := os.Rename(temporaryName, destination); err != nil {
			return fmt.Errorf("publish model projection: %w", err)
		}
		published = true
	}
	return 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_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(fingerprint), strconv.Quote(installation.ModelCatalog.Defaults.Session),
		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
}