package schema

import (
	"context"
	"fmt"
	"strconv"

	"github.com/mem9-ai/drive9/pkg/embedding"
	"github.com/mem9-ai/drive9/pkg/logger"
	"go.uber.org/zap"
)

// Keep this statement list aligned with any external workflow that manages the
// app-embedding TiDB schema directly. If you change columns, indexes,
// generated expressions, or constraints here, update that external workflow at
// the same time.
//
// Export path: `drive9-server schema dump-init-sql --provider local`.
// The dump includes optional FTS/VECTOR ALTER statements unconditionally;
// runtime local init applies those best-effort (AllowUnsupportedOptionalIndexes).
func tidbAppEmbeddingSchemaStatements() []string {
	stmts := tidbAppEmbeddingBaseSchemaStatements()
	stmts = append(stmts, tidbAppEmbeddingOptionalSchemaStatements()...)
	return stmts
}

// TiDBAppEmbeddingTenantSchemaStatements returns the required app-managed
// tenant DDL. Optional FTS/VECTOR indexes are omitted so tests and local
// bootstrap can run on TiDB builds that do not support those features.
func TiDBAppEmbeddingTenantSchemaStatements() []string {
	return CloneStatements(tidbAppEmbeddingBaseSchemaStatements())
}

// TiDBAppEmbeddingLegacyFilesStatements returns the legacy `files` table for
// the app-managed schema. Fresh tenant init does not create this table; tests
// that exercise the dual-write path can opt in.
func TiDBAppEmbeddingLegacyFilesStatements() []string {
	dim := strconv.Itoa(embedding.Dimensions)
	return []string{
		`CREATE TABLE IF NOT EXISTS files (
			file_id                    VARCHAR(64) PRIMARY KEY,
			storage_type               VARCHAR(32) NOT NULL,
			content_layout             VARCHAR(32) NOT NULL DEFAULT 'single',
			storage_ref                TEXT NOT NULL,
			storage_ref_hash           VARCHAR(64) NOT NULL DEFAULT '',
			storage_encryption_mode    VARCHAR(16) NOT NULL DEFAULT 'legacy',
			storage_encryption_key_id  VARCHAR(256) NOT NULL DEFAULT '',
			content_blob               LONGBLOB,
			content_type               VARCHAR(255),
			size_bytes                 BIGINT NOT NULL DEFAULT 0,
			checksum_sha256            VARCHAR(128),
			revision                   BIGINT NOT NULL DEFAULT 1,
			status                     VARCHAR(32) NOT NULL DEFAULT 'PENDING',
			source_id                  VARCHAR(255),
			content_text               LONGTEXT,
			description                LONGTEXT,
			embedding                  VECTOR(` + dim + `),
			embedding_revision         BIGINT,
			description_embedding      VECTOR(` + dim + `),
			description_embedding_revision BIGINT,
			created_at                 DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			confirmed_at               DATETIME(3),
			expires_at                 DATETIME(3)
		)`,
		`CREATE INDEX idx_status ON files(status, created_at)`,
		`CREATE INDEX idx_files_storage_ref_hash ON files(storage_ref_hash)`,
	}
}

func tidbAppEmbeddingBaseSchemaStatements() []string {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS file_nodes (
			node_id      VARCHAR(64) PRIMARY KEY,
			path         TEXT NOT NULL,
			path_hash    VARCHAR(64) NOT NULL DEFAULT '',
			parent_path  TEXT NOT NULL,
			parent_path_hash VARCHAR(64) NOT NULL DEFAULT '',
			name         VARCHAR(255) NOT NULL,
			is_directory BOOLEAN NOT NULL DEFAULT FALSE,
			file_id      VARCHAR(64),
			inode_id     VARCHAR(64),
			created_at   DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3)
		)`,
		`CREATE UNIQUE INDEX idx_path ON file_nodes(path_hash)`,
		`CREATE INDEX idx_parent ON file_nodes(parent_path_hash, name)`,
		`CREATE INDEX idx_file_id ON file_nodes(file_id)`,
		`CREATE INDEX idx_inode_id ON file_nodes(inode_id)`,

		`CREATE TABLE IF NOT EXISTS inodes (
			inode_id     VARCHAR(64) PRIMARY KEY,
			size_bytes   BIGINT NOT NULL DEFAULT 0,
			revision     BIGINT NOT NULL DEFAULT 1,
			mode         INT UNSIGNED NOT NULL DEFAULT 420,
			status       VARCHAR(32) NOT NULL DEFAULT 'PENDING',
			created_at   DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			mtime        DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			confirmed_at DATETIME(3),
			expires_at   DATETIME(3)
		)`,
		`CREATE INDEX idx_inodes_status ON inodes(status, created_at)`,
		`CREATE TABLE IF NOT EXISTS contents (
			inode_id                   VARCHAR(64) PRIMARY KEY,
			storage_type               VARCHAR(32) NOT NULL,
			content_layout             VARCHAR(32) NOT NULL DEFAULT 'single',
			storage_ref                TEXT NOT NULL,
			storage_ref_hash           VARCHAR(64) NOT NULL DEFAULT '',
			storage_encryption_mode    VARCHAR(16) NOT NULL DEFAULT 'legacy',
			storage_encryption_key_id  VARCHAR(256) NOT NULL DEFAULT '',
			content_blob               LONGBLOB,
			content_type               VARCHAR(255),
			checksum_sha256            VARCHAR(128),
			source_id                  VARCHAR(255)
		)`,
		`CREATE INDEX idx_contents_storage_ref_hash ON contents(storage_ref_hash)`,
		`CREATE TABLE IF NOT EXISTS semantic (
			inode_id                           VARCHAR(64) PRIMARY KEY,
			content_text                       LONGTEXT,
			description                        LONGTEXT,
			embedding                          VECTOR(` + strconv.Itoa(embedding.Dimensions) + `),
			embedding_revision                 BIGINT,
			embedding_generation               BIGINT UNSIGNED,
			description_embedding              VECTOR(` + strconv.Itoa(embedding.Dimensions) + `),
			description_embedding_revision     BIGINT,
			description_embedding_generation   BIGINT UNSIGNED
		)`,
		`CREATE TABLE IF NOT EXISTS file_tags (
			file_id   VARCHAR(64) NOT NULL,
			inode_id  VARCHAR(64),
			tag_key   VARCHAR(255) NOT NULL,
			tag_value VARCHAR(255) NOT NULL DEFAULT '',
			source    VARCHAR(32) NOT NULL DEFAULT 'user',
			PRIMARY KEY (file_id, tag_key)
		)`,
		`CREATE INDEX idx_kv ON file_tags(tag_key, tag_value)`,
		`CREATE TABLE IF NOT EXISTS uploads (
			upload_id          VARCHAR(64) PRIMARY KEY,
			file_id            VARCHAR(64) NOT NULL,
			inode_id           VARCHAR(64),
			target_kind        VARCHAR(16) NOT NULL DEFAULT 'file',
			layer_id           VARCHAR(64) NOT NULL DEFAULT '',
			target_path        TEXT NOT NULL,
			target_path_hash   VARCHAR(64) NOT NULL DEFAULT '',
			base_inode_id      VARCHAR(64) NOT NULL DEFAULT '',
			mode               INT UNSIGNED NOT NULL DEFAULT 420,
			s3_upload_id       VARCHAR(255) NOT NULL,
			s3_key             VARCHAR(2048) NOT NULL,
			storage_encryption_mode VARCHAR(16) NOT NULL DEFAULT 'none',
			storage_encryption_key_id VARCHAR(256) NOT NULL DEFAULT '',
			total_size         BIGINT NOT NULL,
			part_size          BIGINT NOT NULL,
			parts_total        INT NOT NULL,
			expected_revision  BIGINT NULL,
			status             VARCHAR(32) NOT NULL DEFAULT 'UPLOADING',
			fingerprint_sha256 VARCHAR(128),
			idempotency_key    VARCHAR(255),
			description        LONGTEXT,
			created_at         DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			updated_at         DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
			expires_at         DATETIME(3) NOT NULL,
			active_target_path_hash VARCHAR(64) AS (CASE WHEN status = 'UPLOADING' THEN target_path_hash ELSE NULL END) VIRTUAL
		)`,
		`ALTER TABLE uploads ADD COLUMN expected_revision BIGINT NULL`,
		`ALTER TABLE uploads ADD COLUMN target_kind VARCHAR(16) NOT NULL DEFAULT 'file'`,
		`ALTER TABLE uploads ADD COLUMN layer_id VARCHAR(64) NOT NULL DEFAULT ''`,
		`ALTER TABLE uploads ADD COLUMN base_inode_id VARCHAR(64) NOT NULL DEFAULT ''`,
		`ALTER TABLE uploads ADD COLUMN mode INT UNSIGNED NOT NULL DEFAULT 420`,
		`CREATE INDEX idx_upload_path ON uploads(target_path_hash, status)`,
		`CREATE UNIQUE INDEX idx_idempotency ON uploads(idempotency_key)`,
		`CREATE UNIQUE INDEX idx_uploads_active ON uploads(active_target_path_hash)`,
		`CREATE TABLE IF NOT EXISTS semantic_tasks (
			task_id           VARCHAR(64) PRIMARY KEY,
			task_type         VARCHAR(32) NOT NULL,
			resource_id       VARCHAR(64) NOT NULL,
			resource_version  BIGINT NOT NULL,
			status            VARCHAR(20) NOT NULL,
			attempt_count     INT NOT NULL DEFAULT 0,
			max_attempts      INT NOT NULL DEFAULT 5,
			receipt           VARCHAR(128) NULL,
			leased_at         DATETIME(3) NULL,
			lease_until       DATETIME(3) NULL,
			available_at      DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			payload_json      JSON NULL,
			last_error        TEXT NULL,
			created_at        DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			updated_at        DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			completed_at      DATETIME(3) NULL
		)`,
		`CREATE UNIQUE INDEX uk_task_resource_version ON semantic_tasks(task_type, resource_id, resource_version)`,
		`CREATE INDEX idx_task_claim ON semantic_tasks(status, available_at, lease_until, created_at)`,
		`CREATE INDEX idx_task_claim_type ON semantic_tasks(status, task_type, available_at, created_at, task_id)`,
		`CREATE TABLE IF NOT EXISTS file_gc_tasks (
			task_id       VARCHAR(64) PRIMARY KEY,
			file_id       VARCHAR(64) NOT NULL,
			inode_id      VARCHAR(64),
			storage_type  VARCHAR(32) NOT NULL,
			content_layout VARCHAR(32) NOT NULL DEFAULT 'single',
			storage_ref   TEXT NOT NULL,
			size_bytes    BIGINT NOT NULL DEFAULT 0,
			content_type  VARCHAR(255) NULL,
			status        VARCHAR(20) NOT NULL,
			attempt_count INT NOT NULL DEFAULT 0,
			max_attempts  INT NOT NULL DEFAULT 0,
			receipt       VARCHAR(128) NULL,
			leased_at     DATETIME(3) NULL,
			lease_until   DATETIME(3) NULL,
			available_at  DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			last_error    TEXT NULL,
			created_at    DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			updated_at    DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			completed_at  DATETIME(3) NULL
		)`,
		`CREATE UNIQUE INDEX uk_file_gc_file_id ON file_gc_tasks(file_id)`,
		`CREATE UNIQUE INDEX uk_file_gc_inode_id ON file_gc_tasks(inode_id)`,
		`CREATE INDEX idx_file_gc_claim ON file_gc_tasks(status, available_at, lease_until, created_at)`,
		`CREATE TABLE IF NOT EXISTS llm_usage (
			id              BIGINT AUTO_INCREMENT PRIMARY KEY,
			task_type       VARCHAR(32) NOT NULL,
			task_id         VARCHAR(64) NOT NULL,
			cost_millicents BIGINT NOT NULL DEFAULT 0,
			raw_units       BIGINT NOT NULL DEFAULT 0,
			raw_unit_type   VARCHAR(16) NOT NULL,
			created_at      DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3)
		)`,
		`CREATE INDEX idx_llm_usage_created ON llm_usage(created_at)`,
		`CREATE TABLE IF NOT EXISTS fs_events (
			seq        BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			path       TEXT NOT NULL,
			op         VARCHAR(64) NOT NULL,
			actor      VARCHAR(255),
			ts         BIGINT NOT NULL,
			created_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3)
		)`,
		`CREATE INDEX idx_fs_events_created ON fs_events(created_at)`,
	}
	stmts = append(stmts, GitWorkspaceTiDBSchemaStatements()...)
	stmts = append(stmts, FSLayerTiDBSchemaStatements()...)
	stmts = append(stmts, JournalTiDBSchemaStatements()...)
	stmts = append(stmts, VaultTiDBSchemaStatements()...)
	return stmts
}

func tidbAppEmbeddingOptionalSchemaStatements() []string {
	return []string{

		`ALTER TABLE semantic
			ADD FULLTEXT INDEX idx_semantic_fts_content(content_text)
			WITH PARSER MULTILINGUAL
			ADD_COLUMNAR_REPLICA_ON_DEMAND`,
		`ALTER TABLE semantic
			ADD FULLTEXT INDEX idx_semantic_fts_description(description)
			WITH PARSER MULTILINGUAL
			ADD_COLUMNAR_REPLICA_ON_DEMAND`,
		`ALTER TABLE semantic
			ADD VECTOR INDEX idx_semantic_cosine((VEC_COSINE_DISTANCE(embedding)))
			ADD_COLUMNAR_REPLICA_ON_DEMAND`,
		`ALTER TABLE semantic
			ADD VECTOR INDEX idx_semantic_desc_cosine((VEC_COSINE_DISTANCE(description_embedding)))
			ADD_COLUMNAR_REPLICA_ON_DEMAND`,
	}
}

func initTiDBAppEmbeddingSchema(ctx context.Context, dsn string, opts InitTiDBTenantSchemaOptions) error {
	db, err := OpenTiDBSchemaDB(ctx, dsn)
	if err != nil {
		return err
	}
	defer func() { _ = closeTiDBSchemaDB(db) }()
	if !IsTiDBCluster(ctx, db) {
		return fmt.Errorf("provider requires TiDB capabilities (FTS/VECTOR)")
	}
	if err := repairMySQLPathHashSchema(ctx, db); err != nil {
		return err
	}
	if err := ExecSchemaStatementsParallelByTableContext(ctx, db, tidbAppEmbeddingBaseSchemaStatements()); err != nil {
		return err
	}
	if opts.AllowUnsupportedOptionalIndexes {
		// Local-only compatibility mode can continue without these indexes; other
		// callers still execute the same statements strictly below.
		skipped, err := ExecOptionalSchemaStatements(ctx, db, tidbAppEmbeddingOptionalSchemaStatements())
		if err != nil {
			return err
		}
		if skipped > 0 {
			logger.Warn(ctx, "tidb_app_optional_indexes_skipped",
				zap.Int("skipped_count", skipped),
				zap.String("reason", "allow_unsupported_optional_indexes"))
		}
	} else if err := ExecSchemaStatementsContext(ctx, db, tidbAppEmbeddingOptionalSchemaStatements()); err != nil {
		return err
	}
	return EnsureTiDBSchemaForMode(ctx, db, TiDBEmbeddingModeApp)
}