package query

import (
	"fmt"
	"regexp"
	"strings"
	"time"
)

// RewriteOptions controls server-side CTE injection for scoped datasets.
type RewriteOptions struct {
	Dataset         string
	Region          string
	LogsTableSuffix string
	Scope           Scope
	TimeRange       *TimeRange
}

var (
	fromJoinIdentRe = regexp.MustCompile(`(?i)\b(?:from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)\b`)
)

// BuildEffectiveSQL builds the Athena SQL submitted to the engine.
//
// Scoped datasets: prepend WITH <alias> AS (SELECT * FROM <physical> WHERE <scope[/optional time]>)
// and leave caller SQL unchanged (must reference the virtual alias).
//
// dataset=raw: passes SQL through; scope/time_range are not allowed.
func BuildEffectiveSQL(userSQL string, opt RewriteOptions) (string, error) {
	userSQL = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(userSQL), ";"))
	if err := ValidateSelectSQL(userSQL); err != nil {
		return "", err
	}
	if err := RejectDangerousSQL(userSQL); err != nil {
		return "", err
	}

	ds := normalizeDataset(opt.Dataset)
	switch ds {
	case DatasetRaw:
		if len(opt.Scope) > 0 {
			return "", fmt.Errorf("scope is not allowed for dataset=%q (use a scoped dataset such as %q / %q)", DatasetRaw, DatasetLogs, DatasetTopsql)
		}
		if opt.TimeRange != nil {
			return "", fmt.Errorf("time_range is not allowed for dataset=%q", DatasetRaw)
		}
		return userSQL, nil
	default:
		if !isScopedDataset(ds) {
			return "", fmt.Errorf("unsupported dataset %q (supported: %s)", opt.Dataset, supportedDatasetsHint())
		}
		return buildScopedDatasetSQL(userSQL, opt)
	}
}

func buildScopedDatasetSQL(userSQL string, opt RewriteOptions) (string, error) {
	ds := normalizeDataset(opt.Dataset)
	profile, err := datasetProfile(ds, opt.Region, opt.LogsTableSuffix)
	if err != nil {
		return "", err
	}
	if opt.Region == "" {
		return "", fmt.Errorf("region is required for dataset=%s", ds)
	}
	if len(opt.Scope) == 0 {
		return "", fmt.Errorf("scope is required for dataset=%s (required fields: %s)", ds, requiredScopeFieldNames(profile.Fields))
	}

	// Validate scope before the virtual-table check so missing/empty keys are reported first.
	if err := validateRequiredScope(opt.Scope, profile.Fields); err != nil {
		return "", err
	}

	if !virtualTableReferenced(userSQL, profile.Alias) {
		return "", virtualTableReferenceError(ds, profile.Alias, userSQL)
	}

	preds, err := buildScopePredicates(opt.Scope, profile.Fields)
	if err != nil {
		return "", err
	}
	// time_range is optional convenience: inject onto the dataset's known time column.
	timePreds, err := buildTimeRangePredicates(opt.TimeRange, profile.TimeColumn, profile.TimeFormat)
	if err != nil {
		return "", err
	}
	preds = append(preds, timePreds...)
	if len(preds) == 0 {
		return "", fmt.Errorf("internal error: empty predicates for dataset=%s", ds)
	}

	cteBody := fmt.Sprintf("  SELECT * FROM %s\n  WHERE %s", profile.PhysicalSQL, strings.Join(preds, "\n    AND "))
	return prependCTE(userSQL, profile.Alias, cteBody), nil
}

func requiredScopeFieldNames(fields []scopeFieldSpec) string {
	names := make([]string, 0, len(fields))
	for _, f := range fields {
		if f.Required {
			names = append(names, f.APIKeys[0])
		}
	}
	return strings.Join(names, ", ")
}

// validateRequiredScope collects all missing / blank required scope fields into one error.
func validateRequiredScope(scope Scope, fields []scopeFieldSpec) error {
	var missing []string
	for _, f := range fields {
		if !f.Required {
			continue
		}
		v, err := f.resolve(scope)
		if err != nil || v.isBlank() {
			missing = append(missing, f.APIKeys[0])
		}
	}
	if len(missing) == 0 {
		return nil
	}
	return fmt.Errorf("missing required scope fields: %s (required: %s)",
		strings.Join(missing, ", "), requiredScopeFieldNames(fields))
}

func virtualTableReferenceError(ds, alias, userSQL string) error {
	msg := fmt.Sprintf("dataset=%s requires querying virtual table %q (e.g. FROM %s / JOIN %s)", ds, alias, alias, alias)
	if ds != alias {
		lower := strings.ToLower(userSQL)
		if strings.Contains(lower, strings.ToLower(ds)) {
			msg += fmt.Sprintf("; note: %q is the dataset id — in SQL use %q (hyphen is not a valid unquoted identifier)", ds, alias)
		} else {
			msg += fmt.Sprintf("; note: dataset id is %q, SQL virtual table is %q", ds, alias)
		}
	}
	return fmt.Errorf("%s", msg)
}

type datasetSQLProfile struct {
	Alias       string
	PhysicalSQL string
	Fields      []scopeFieldSpec
	// TimeColumn / TimeFormat describe optional time_range injection. Empty = not supported.
	TimeColumn string
	TimeFormat string
}

func datasetProfile(ds, region, suffix string) (datasetSQLProfile, error) {
	switch ds {
	case DatasetLogs:
		return datasetSQLProfile{
			Alias:       DatasetLogs,
			PhysicalSQL: logsPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"env"}, SQLColumn: "env", Required: true},
				{APIKeys: []string{"deploy_type", "deploytype"}, SQLColumn: "deploytype", Required: true, Normalize: normalizeDeployTypeValue},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
				{APIKeys: []string{"components", "component"}, SQLColumn: "component", Required: true},
			},
			TimeColumn: "timestamp",
			TimeFormat: TimeFormatUnixSeconds,
		}, nil
	case DatasetLogStats:
		return datasetSQLProfile{
			Alias:       DatasetLogStatsAlias,
			PhysicalSQL: logStatsPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogStatsDate},
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
			},
			// date is the partition (via scope); optional time_range maps onto file column hour (0-23).
			TimeColumn: "hour",
			TimeFormat: TimeFormatHourOfDay,
		}, nil
	case DatasetTopsql:
		return datasetSQLProfile{
			Alias:       DatasetTopsql,
			PhysicalSQL: topsqlPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				{APIKeys: []string{"component", "components"}, SQLColumn: "component", Required: true},
				{APIKeys: []string{"instance"}, SQLColumn: "instance", Required: true},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
			},
			TimeColumn: "timestamps",
			TimeFormat: TimeFormatUnixSeconds,
		}, nil
	case DatasetTopsqlSQLMeta:
		return datasetSQLProfile{
			Alias:       DatasetTopsqlSQLMetaAlias,
			PhysicalSQL: topsqlSQLMetaPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: false, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
			},
		}, nil
	case DatasetTopsqlPlanMeta:
		return datasetSQLProfile{
			Alias:       DatasetTopsqlPlanMetaAlias,
			PhysicalSQL: topsqlPlanMetaPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: false, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
			},
		}, nil
	case DatasetTopru:
		return datasetSQLProfile{
			Alias:       DatasetTopru,
			PhysicalSQL: topruPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
			},
			TimeColumn: "timestamps",
			TimeFormat: TimeFormatUnixSeconds,
		}, nil
	case DatasetSQLStatement:
		return datasetSQLProfile{
			Alias:       DatasetSQLStatement,
			PhysicalSQL: sqlstatementPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogsDate},
			},
			TimeColumn: "begin,end",
			TimeFormat: TimeFormatWindowOverlap,
		}, nil
	case DatasetSlowlogs:
		return datasetSQLProfile{
			Alias:       DatasetSlowlogs,
			PhysicalSQL: slowlogsPhysicalTableSQLName(region, suffix),
			Fields: []scopeFieldSpec{
				{APIKeys: []string{"org", "tenant", "tenant_id"}, SQLColumn: "org", Required: true},
				{APIKeys: []string{"cluster"}, SQLColumn: "cluster", Required: true},
				// Path partition date=yyyyMMdd (required for Athena partition projection pruning).
				{APIKeys: []string{"date"}, SQLColumn: "date", Required: true, AllowLegacyDatePair: true, Normalize: normalizeLogStatsDate},
			},
			TimeColumn: "time",
			TimeFormat: TimeFormatUnixFloat,
		}, nil
	case DatasetMetricsStatusSources:
		return datasetSQLProfile{
			Alias:       DatasetMetricsStatusSourcesAlias,
			PhysicalSQL: metricsStatusSourcesPhysicalTableSQLName(region, suffix),
			Fields:      metricsStatusScopeFieldSpecs(),
		}, nil
	case DatasetMetricsStatusByMetric:
		return datasetSQLProfile{
			Alias:       DatasetMetricsStatusByMetricAlias,
			PhysicalSQL: metricsStatusByMetricPhysicalTableSQLName(region, suffix),
			Fields:      metricsStatusScopeFieldSpecs(),
		}, nil
	case DatasetMetricsStatusByLabel:
		return datasetSQLProfile{
			Alias:       DatasetMetricsStatusByLabelAlias,
			PhysicalSQL: metricsStatusByLabelPhysicalTableSQLName(region, suffix),
			Fields:      metricsStatusScopeFieldSpecs(),
		}, nil
	case DatasetMetricsStatusByLabelValue:
		return datasetSQLProfile{
			Alias:       DatasetMetricsStatusByLabelValueAlias,
			PhysicalSQL: metricsStatusByLabelValuePhysicalTableSQLName(region, suffix),
			Fields:      metricsStatusScopeFieldSpecs(),
		}, nil
	default:
		return datasetSQLProfile{}, fmt.Errorf("no scope profile for dataset %q", ds)
	}
}

func buildScopePredicates(scope Scope, fields []scopeFieldSpec) ([]string, error) {
	var preds []string
	for _, f := range fields {
		v, err := f.resolve(scope)
		if err != nil {
			return nil, err
		}
		if v.isZero() {
			continue
		}
		ps, err := scopeValueToPredicates(f.SQLColumn, v, f.Normalize)
		if err != nil {
			return nil, fmt.Errorf("scope.%s: %w", f.APIKeys[0], err)
		}
		preds = append(preds, ps...)
	}
	return preds, nil
}

func virtualTableReferenced(userSQL, alias string) bool {
	for _, m := range fromJoinIdentRe.FindAllStringSubmatch(userSQL, -1) {
		if len(m) >= 2 && strings.EqualFold(m[1], alias) {
			return true
		}
	}
	return false
}

// prependCTE inserts `name AS (body)` at the front of the WITH list, or wraps with a new WITH.
func prependCTE(userSQL, name, body string) string {
	trimmed := strings.TrimSpace(userSQL)
	cte := fmt.Sprintf("%s AS (\n%s\n)", name, body)
	if len(trimmed) >= 4 && strings.EqualFold(trimmed[:4], "WITH") {
		rest := strings.TrimSpace(trimmed[4:])
		return "WITH " + cte + ",\n" + rest
	}
	return "WITH " + cte + "\n" + trimmed
}

func normalizeDeployTypeValue(v string) (string, error) {
	v = strings.ToLower(strings.TrimSpace(v))
	switch v {
	case "":
		return "", fmt.Errorf("deploy_type must not be empty")
	case "shared":
		return "serverless", nil
	case "nextgen", "premium":
		return "premium", nil
	default:
		return v, nil
	}
}

func normalizeLogsDate(v string) (string, error) {
	v = strings.TrimSpace(v)
	if _, err := time.Parse("2006-01-02", v); err != nil {
		return "", fmt.Errorf("invalid date %q (want yyyy-MM-dd): %w", v, err)
	}
	return v, nil
}

// normalizeLogStatsDate accepts yyyyMMdd or yyyy-MM-dd and returns yyyyMMdd.
func normalizeLogStatsDate(v string) (string, error) {
	v = strings.TrimSpace(v)
	if len(v) == 8 {
		if _, err := time.Parse("20060102", v); err != nil {
			return "", fmt.Errorf("invalid date %q (want yyyyMMdd or yyyy-MM-dd): %w", v, err)
		}
		return v, nil
	}
	t, err := time.Parse("2006-01-02", v)
	if err != nil {
		return "", fmt.Errorf("invalid date %q (want yyyyMMdd or yyyy-MM-dd): %w", v, err)
	}
	return t.Format("20060102"), nil
}

func quoteLiteral(v string) string {
	return "'" + strings.ReplaceAll(v, "'", "''") + "'"
}

func quoteIdent(v string) string {
	return `"` + strings.ReplaceAll(v, `"`, `""`) + `"`
}