package storage

import (
	"context"
	"fmt"

	"archivdms/internal/classifier"
)

// MLClassifierKinds is the fixed set of taxonomy kinds the Naive-Bayes
// classifier is trained/predicted for, in a stable order (used by the retrain
// CLI so runs are deterministic).
var MLClassifierKinds = []string{
	classifier.KindDocumentTypes,
	classifier.KindCorrespondents,
	classifier.KindTags,
}

// TrainClassifier rebuilds the Naive-Bayes model for one tenant and one kind and
// returns the number of training documents used. Thin wrapper around
// classifier.Train that keeps the Store's *pgxpool.Pool encapsulated (the
// classifier package must not import storage). See classifier.Train for
// semantics (full rebuild, MinDocsPerClass threshold, no error on empty data).
func (s *Store) TrainClassifier(ctx context.Context, tenantID int64, kind string) (int, error) {
	return classifier.New(s.db).Train(ctx, tenantID, kind)
}

// StartMLRun inserts a fresh ml_classifier_runs row in status 'running' for a
// tenant and returns its id. The retrain CLI creates one run per tenant that
// spans all kinds, then finalises it with FinishMLRun.
func (s *Store) StartMLRun(ctx context.Context, tenantID int64) (int64, error) {
	var id int64
	if err := s.db.QueryRow(ctx,
		`INSERT INTO ml_classifier_runs (tenant_id, status) VALUES ($1, 'running') RETURNING id`,
		tenantID).Scan(&id); err != nil {
		return 0, fmt.Errorf("storage: start ml_classifier run: %w", err)
	}
	return id, nil
}

// FinishMLRun finalises an ml_classifier_runs row: sets completed_at=now(),
// doc_count, the terminal status ('completed', 'failed' or
// 'skipped_insufficient_data') and an optional error message. errMsg "" stores
// SQL NULL.
func (s *Store) FinishMLRun(ctx context.Context, runID int64, docCount int, status, errMsg string) error {
	var errArg any
	if errMsg != "" {
		errArg = errMsg
	}
	_, err := s.db.Exec(ctx,
		`UPDATE ml_classifier_runs
		 SET completed_at = now(), doc_count = $2, status = $3, error = $4
		 WHERE id = $1`,
		runID, docCount, status, errArg)
	if err != nil {
		return fmt.Errorf("storage: finish ml_classifier run: %w", err)
	}
	return nil
}