package schema

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
)

var requiredTableNames = []string{
	"organizations",
	"organization_settings",
	"organization_feishu_credentials",
	"organization_teams",
	"organization_units",
	"organization_members",
	"organization_auth_bindings",
	"organization_directory_sync_runs",
	"organization_invites",
	"organization_workspace_member_grants",
	"organization_workspace_hosted_model_authorizations",
	"organization_member_team_memberships",
	"organization_member_unit_memberships",
	"organization_workspace_team_grants",
	"organization_workspace_member_projections",
	"organization_providers",
	"organization_llm_proxy_keys",
	"organization_provider_tokens",
	"organization_provider_routes",
	"organization_provider_route_groups",
	"organization_provider_route_group_model_refs",
	"organization_provider_route_group_provider_refs",
	"organization_workspace_provider_grants",
	"organization_workspace_route_grants",
	"organization_llm_usage_buckets",
	"organization_llm_error_buckets",
	"organization_llm_daily_usage",
	"organization_llm_quota_policies",
}

var requiredColumnNames = map[string][]string{
	"workspaces": {
		"organization_id",
	},
	"organization_settings": {
		"default_workspace_id",
		"status",
		"directory_revision",
		"auth_provider",
		"directory_provider",
		"feishu_tenant_key",
		"feishu_root_department_id",
		"feishu_last_auto_join_status",
		"feishu_last_auto_join_subject_type",
		"feishu_last_auto_join_subject_hint",
		"feishu_last_auto_join_message",
		"feishu_last_auto_join_attempted_at",
		"hosted_model_picker_visibility_policy",
		"hosted_model_auto_authorization_enabled",
		"hosted_model_auto_authorization_mode",
		"hosted_model_auto_authorization_model_refs_json",
	},
	"organization_feishu_credentials": {
		"app_id_ciphertext",
		"app_secret_ciphertext",
		"updated_by_user_id",
		"last_validated_at",
	},
	"organization_teams": {
		"is_default",
		"backing_unit_id",
		"external_provider",
		"external_id",
		"synced_at",
		"sync_mode",
		"sync_source",
		"last_sync_run_id",
	},
	"organization_units": {
		"parent_unit_id",
		"manager_member_id",
		"external_provider",
		"external_id",
		"synced_at",
		"sync_mode",
		"sync_source",
		"last_sync_run_id",
	},
	"organization_members": {
		"user_id",
		"primary_team_id",
		"primary_unit_id",
		"org_role",
		"status",
		"source",
		"default_workspace_role",
		"external_provider",
		"external_id",
		"synced_at",
		"sync_mode",
		"sync_source",
		"last_sync_run_id",
		"invited_by_user_id",
		"created_by_user_id",
	},
	"organization_auth_bindings": {
		"organization_member_id",
		"auth_provider",
		"subject_key",
		"subject_value",
	},
	"organization_directory_sync_runs": {
		"directory_revision",
		"external_provider",
		"sync_source",
		"status",
	},
	"organization_invites": {
		"email",
		"org_role",
		"default_team_id",
		"default_workspace_role",
		"token_hash",
	},
	"organization_workspace_member_grants": {
		"workspace_role",
	},
	"organization_workspace_hosted_model_authorizations": {
		"workspace_id",
		"authorization_mode",
		"selected_model_refs_json",
		"status",
		"authorized_by_user_id",
	},
	"organization_member_team_memberships": {
		"organization_team_id",
		"is_primary",
	},
	"organization_member_unit_memberships": {
		"organization_unit_id",
		"is_primary",
	},
	"organization_workspace_team_grants": {
		"organization_team_id",
	},
	"organization_workspace_member_projections": {
		"member_id",
		"applied_source",
	},
	"organization_providers": {
		"driver_kind",
		"base_url",
		"endpoints_json",
		"capabilities_json",
		"status",
	},
	"organization_llm_proxy_keys": {
		"organization_id",
		"token_hash",
		"token_hint",
		"status",
	},
	"organization_provider_routes": {
		"organization_id",
		"vendor",
		"display_name",
		"models_json",
		"source_kind",
		"status",
	},
	"organization_provider_route_groups": {
		"organization_id",
		"route_id",
		"models_json",
		"provider_refs_json",
		"strategy",
		"priority",
		"managed_by_system",
		"status",
	},
	"organization_provider_route_group_model_refs": {
		"organization_id",
		"route_id",
		"route_group_id",
		"model_id",
		"ref_order",
	},
	"organization_provider_route_group_provider_refs": {
		"organization_id",
		"route_id",
		"route_group_id",
		"provider_id",
		"ref_order",
	},
	"organization_provider_tokens": {
		"organization_provider_id",
		"token_ciphertext",
		"secret_json_ciphertext",
		"token_hint",
		"config_json",
		"status",
	},
	"organization_workspace_provider_grants": {
		"workspace_id",
		"organization_provider_id",
		"status",
	},
	"organization_workspace_route_grants": {
		"workspace_id",
		"organization_route_id",
		"status",
	},
	"organization_llm_usage_buckets": {
		"organization_id",
		"workspace_id",
		"proxy_key_id",
		"source",
		"organization_provider_id",
		"model",
		"bucket_start",
		"request_count",
		"input_tokens",
		"cached_input_tokens",
		"output_tokens",
		"duration_ms",
		"queue_wait_ms",
	},
	"organization_llm_error_buckets": {
		"organization_id",
		"workspace_id",
		"proxy_key_id",
		"source",
		"organization_provider_id",
		"model",
		"error_type",
		"bucket_start",
		"error_count",
	},
	"organization_llm_daily_usage": {
		"subject_type",
		"subject_id",
		"usage_date",
		"spent_usd",
		"request_count",
		"input_tokens",
		"cached_input_tokens",
		"output_tokens",
	},
	"organization_llm_quota_policies": {
		"subject_type",
		"subject_id",
		"status",
		"daily_quota_usd",
	},
}

type schemaRows interface {
	Close()
	Err() error
	Next() bool
	Scan(...any) error
}

type schemaRow interface {
	Scan(...any) error
}

type schemaQueryer interface {
	Query(context.Context, string, ...any) (schemaRows, error)
	QueryRow(context.Context, string, ...any) schemaRow
}

type sqlDBQueryer struct {
	db *sql.DB
}

type sqlRowsAdapter struct {
	rows *sql.Rows
}

func (adapter sqlRowsAdapter) Close() {
	_ = adapter.rows.Close()
}

func (adapter sqlRowsAdapter) Err() error {
	return adapter.rows.Err()
}

func (adapter sqlRowsAdapter) Next() bool {
	return adapter.rows.Next()
}

func (adapter sqlRowsAdapter) Scan(dest ...any) error {
	return adapter.rows.Scan(dest...)
}

func (queryer sqlDBQueryer) Query(ctx context.Context, query string, args ...any) (schemaRows, error) {
	rows, err := queryer.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	return sqlRowsAdapter{rows: rows}, nil
}

func (queryer sqlDBQueryer) QueryRow(ctx context.Context, query string, args ...any) schemaRow {
	return queryer.db.QueryRowContext(ctx, query, args...)
}

// VerifyTiDB checks the org provider schema on TiDB.
func VerifyTiDB(ctx context.Context, db *sql.DB) error {
	if db == nil {
		return fmt.Errorf("tidb db is nil")
	}
	return verify(ctx, sqlDBQueryer{db: db}, `SELECT DATABASE()`, `
SELECT table_name
FROM information_schema.tables
WHERE table_schema = ?
  AND table_type = 'BASE TABLE'
`, `
SELECT table_name, column_name
FROM information_schema.columns
WHERE table_schema = ?
`)
}

func verify(ctx context.Context, queryer schemaQueryer, schemaNameQuery string, tableQuery string, columnQuery string) error {
	var schemaName string
	if err := queryer.QueryRow(ctx, schemaNameQuery).Scan(&schemaName); err != nil {
		return fmt.Errorf("load org schema name: %w", err)
	}
	schemaName = strings.TrimSpace(schemaName)
	if schemaName == "" {
		return fmt.Errorf("organization schema name is empty")
	}

	actualTables, err := loadActualTables(ctx, queryer, tableQuery, schemaName)
	if err != nil {
		return err
	}
	actualColumns, err := loadActualColumns(ctx, queryer, columnQuery, schemaName)
	if err != nil {
		return err
	}

	missingTables := make([]string, 0)
	for _, tableName := range requiredTableNames {
		if _, ok := actualTables[tableName]; !ok {
			missingTables = append(missingTables, tableName)
		}
	}

	missingColumns := make([]string, 0)
	for tableName, columnNames := range requiredColumnNames {
		available := actualColumns[tableName]
		for _, columnName := range columnNames {
			if _, ok := available[columnName]; !ok {
				missingColumns = append(missingColumns, fmt.Sprintf("%s.%s", tableName, columnName))
			}
		}
	}
	if len(missingTables) == 0 && len(missingColumns) == 0 {
		return nil
	}

	parts := make([]string, 0, 2)
	if len(missingTables) > 0 {
		parts = append(parts, "missing org tables: "+strings.Join(missingTables, ", "))
	}
	if len(missingColumns) > 0 {
		parts = append(parts, "missing org columns: "+strings.Join(missingColumns, ", "))
	}
	return fmt.Errorf("agentteam-org schema mismatch (%s)", strings.Join(parts, "; "))
}

func loadActualTables(ctx context.Context, queryer schemaQueryer, query string, schemaName string) (map[string]struct{}, error) {
	rows, err := queryer.Query(ctx, query, schemaName)
	if err != nil {
		return nil, fmt.Errorf("check org tables: %w", err)
	}
	defer rows.Close()

	actual := make(map[string]struct{}, len(requiredTableNames))
	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err != nil {
			return nil, fmt.Errorf("scan org table: %w", err)
		}
		actual[strings.TrimSpace(tableName)] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate org tables: %w", err)
	}
	return actual, nil
}

func loadActualColumns(ctx context.Context, queryer schemaQueryer, query string, schemaName string) (map[string]map[string]struct{}, error) {
	rows, err := queryer.Query(ctx, query, schemaName)
	if err != nil {
		return nil, fmt.Errorf("check org columns: %w", err)
	}
	defer rows.Close()

	actual := make(map[string]map[string]struct{}, len(requiredColumnNames))
	for rows.Next() {
		var tableName string
		var columnName string
		if err := rows.Scan(&tableName, &columnName); err != nil {
			return nil, fmt.Errorf("scan org column: %w", err)
		}
		tableName = strings.TrimSpace(tableName)
		columnName = strings.TrimSpace(columnName)
		if actual[tableName] == nil {
			actual[tableName] = make(map[string]struct{})
		}
		actual[tableName][columnName] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate org columns: %w", err)
	}
	return actual, nil
}