package storage

import (
	"context"
	"fmt"
)

// initMLClassifierSchema creates the Naive-Bayes retraining/classification
// tables (ml_classifier_tokens/ml_classifier_classes/ml_classifier_runs) plus
// the assigned_via provenance columns on document_tags/documents that let the
// classifier tell apart manually-set, rule-engine-set and ML-accepted
// assignments. Idempotent, called from (*Store).initSchema AFTER
// initTaxonomySchema (depends on document_types/correspondents/tags/
// document_tags/documents existing). See migrations/020_ml_classifier.sql.
func (s *Store) initMLClassifierSchema(ctx context.Context) error {
	_, err := s.db.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS ml_classifier_tokens (
			id           BIGSERIAL PRIMARY KEY,
			tenant_id    BIGINT NOT NULL,
			kind         TEXT NOT NULL CHECK (kind IN ('document_types','correspondents','tags')),
			entity_id    BIGINT NOT NULL,
			token        TEXT NOT NULL,
			count        BIGINT NOT NULL DEFAULT 0,
			UNIQUE (tenant_id, kind, entity_id, token)
		);
		CREATE INDEX IF NOT EXISTS idx_ml_tokens_lookup ON ml_classifier_tokens(tenant_id, kind, token);

		CREATE TABLE IF NOT EXISTS ml_classifier_classes (
			tenant_id      BIGINT NOT NULL,
			kind           TEXT NOT NULL,
			entity_id      BIGINT NOT NULL,
			doc_count      BIGINT NOT NULL DEFAULT 0,
			total_tokens   BIGINT NOT NULL DEFAULT 0,
			PRIMARY KEY (tenant_id, kind, entity_id)
		);

		CREATE TABLE IF NOT EXISTS ml_classifier_runs (
			id               BIGSERIAL PRIMARY KEY,
			tenant_id        BIGINT NOT NULL,
			started_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
			completed_at     TIMESTAMPTZ,
			doc_count        BIGINT NOT NULL DEFAULT 0,
			status           TEXT NOT NULL DEFAULT 'running' CHECK (status IN ('running','completed','failed','skipped_insufficient_data')),
			error            TEXT
		);
		CREATE INDEX IF NOT EXISTS idx_ml_classifier_runs_tenant ON ml_classifier_runs(tenant_id);
	`)
	if err != nil {
		return fmt.Errorf("storage: create ml_classifier tables: %w", err)
	}

	// Provenance columns: default 'manual' so existing rows are NOT retroactively
	// (mis)classified as rule-assigned — conservative default keeps them out of
	// automatic ML-accepted reclassification too.
	_, err = s.db.Exec(ctx, `
		ALTER TABLE document_tags ADD COLUMN IF NOT EXISTS assigned_via TEXT NOT NULL DEFAULT 'manual' CHECK (assigned_via IN ('manual','rule','ml_accepted'));
		ALTER TABLE documents ADD COLUMN IF NOT EXISTS doc_type_assigned_via TEXT NOT NULL DEFAULT 'manual' CHECK (doc_type_assigned_via IN ('manual','rule','ml_accepted'));
		ALTER TABLE documents ADD COLUMN IF NOT EXISTS correspondent_assigned_via TEXT NOT NULL DEFAULT 'manual' CHECK (correspondent_assigned_via IN ('manual','rule','ml_accepted'));
	`)
	if err != nil {
		return fmt.Errorf("storage: alter documents/document_tags for ml_classifier provenance: %w", err)
	}
	return nil
}