package query

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/glue"
	"github.com/aws/aws-sdk-go-v2/service/glue/types"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
)

type glueAPI interface {
	GetDatabase(ctx context.Context, params *glue.GetDatabaseInput, optFns ...func(*glue.Options)) (*glue.GetDatabaseOutput, error)
	CreateDatabase(ctx context.Context, params *glue.CreateDatabaseInput, optFns ...func(*glue.Options)) (*glue.CreateDatabaseOutput, error)
	GetTable(ctx context.Context, params *glue.GetTableInput, optFns ...func(*glue.Options)) (*glue.GetTableOutput, error)
	CreateTable(ctx context.Context, params *glue.CreateTableInput, optFns ...func(*glue.Options)) (*glue.CreateTableOutput, error)
	UpdateTable(ctx context.Context, params *glue.UpdateTableInput, optFns ...func(*glue.Options)) (*glue.UpdateTableOutput, error)
}

type s3ListAPI interface {
	ListObjectsV2(ctx context.Context, params *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
}

// Catalog manages built-in OLAP Glue tables for the local region.
type Catalog struct {
	cfg    Config
	tables []ResolvedTable
	glue   glueAPI
	s3     s3ListAPI
}

// TableListItem is returned by GET /olap/tables.
type TableListItem struct {
	Name             string           `json:"name"`
	FullName         string           `json:"full_name"`
	Database         string           `json:"database"`
	Dataset          string           `json:"dataset,omitempty"`
	DatasetAlias     string           `json:"dataset_alias,omitempty"`
	Description      string           `json:"description"`
	Location         string           `json:"location"`
	CatalogStatus    string           `json:"catalog_status"`
	DataStatus       string           `json:"data_status"`
	RequiresScope    bool             `json:"requires_scope"`
	ScopeFields      []ScopeFieldMeta `json:"scope_fields,omitempty"` // ordered Hive dive path
	FileColumns      []ColumnMeta     `json:"file_columns"`
	PartitionColumns []ColumnMeta     `json:"partition_columns"`
}

// TableRef identifies a database.table pair referenced in SQL.
type TableRef struct {
	Database string
	Table    string
}

func NewCatalog(cfg Config) (*Catalog, error) {
	c := &Catalog{
		cfg:    cfg,
		tables: resolveRegistryTables(cfg),
	}
	if normalizeProvider(cfg.Vendor) != "aws" {
		return c, nil
	}
	awsCfg, err := config.LoadDefaultConfig(context.Background(), config.WithRegion(cfg.Region))
	if err != nil {
		return nil, fmt.Errorf("load aws config for glue catalog: %w", err)
	}
	c.glue = glue.NewFromConfig(awsCfg)
	c.s3 = s3.NewFromConfig(awsCfg)
	return c, nil
}

func (c *Catalog) RegistryTables() []ResolvedTable {
	out := make([]ResolvedTable, len(c.tables))
	copy(out, c.tables)
	return out
}

func (c *Catalog) ListTables(ctx context.Context) ([]TableListItem, error) {
	out := make([]TableListItem, 0, len(c.tables))
	for _, t := range c.tables {
		ds := t.Definition.Dataset
		alias := datasetVirtualAlias(ds)
		item := TableListItem{
			Name:             t.Name,
			FullName:         t.FullName,
			Database:         t.Database,
			Dataset:          ds,
			DatasetAlias:     alias,
			Description:      t.Definition.Description,
			Location:         t.Location,
			RequiresScope:    t.Definition.RequiresScope,
			FileColumns:      append([]ColumnMeta(nil), t.Definition.FileColumns...),
			PartitionColumns: append([]ColumnMeta(nil), t.Definition.PartitionColumns...),
		}
		if dive, err := divePathForDataset(ds); err == nil {
			item.ScopeFields = scopeFieldMetas(dive)
		}
		if c.glue == nil {
			if normalizeProvider(c.cfg.Vendor) == "aws" {
				item.CatalogStatus = CatalogStatusMissing
			} else {
				item.CatalogStatus = CatalogStatusNotSupported
			}
			item.DataStatus = DataStatusUnknown
			out = append(out, item)
			continue
		}
		exists, err := c.tableExists(ctx, t.Database, t.Name)
		if err != nil {
			return nil, err
		}
		if exists {
			item.CatalogStatus = CatalogStatusReady
		} else {
			item.CatalogStatus = CatalogStatusMissing
		}
		item.DataStatus = c.probeDataStatus(ctx, t.LocationPrefix)
		out = append(out, item)
	}
	return out, nil
}

func (c *Catalog) EnsureAll(ctx context.Context) error {
	if c.glue == nil {
		return nil
	}
	databases := map[string]struct{}{}
	for _, t := range c.tables {
		databases[t.Database] = struct{}{}
	}
	if db := strings.TrimSpace(c.cfg.AthenaDatabase); db != "" {
		databases[db] = struct{}{}
	}
	for db := range databases {
		if err := c.ensureDatabase(ctx, db); err != nil {
			return err
		}
	}
	for _, t := range c.tables {
		if err := c.ensureTable(ctx, t); err != nil {
			return err
		}
	}
	return nil
}

func (c *Catalog) EnsureReferenced(ctx context.Context, refs []TableRef) error {
	if c.glue == nil || len(refs) == 0 {
		return nil
	}
	seen := map[string]struct{}{}
	for _, ref := range refs {
		key := ref.Database + "." + ref.Table
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		// Always ensure the Glue database exists for raw / ad-hoc tables.
		if err := c.ensureDatabase(ctx, ref.Database); err != nil {
			return err
		}
		t, ok := findResolvedTable(c.tables, ref.Database, ref.Table)
		if !ok {
			continue
		}
		if err := c.ensureTable(ctx, *t); err != nil {
			return err
		}
	}
	return nil
}

func (c *Catalog) ensureDatabase(ctx context.Context, name string) error {
	_, err := c.glue.GetDatabase(ctx, &glue.GetDatabaseInput{Name: aws.String(name)})
	if err == nil {
		return nil
	}
	if !isGlueNotFound(err) {
		return fmt.Errorf("glue GetDatabase %s: %w", name, err)
	}
	_, err = c.glue.CreateDatabase(ctx, &glue.CreateDatabaseInput{
		DatabaseInput: &types.DatabaseInput{
			Name:        aws.String(name),
			Description: aws.String("TiDB Cloud observability offline OLAP Athena database"),
		},
	})
	if err != nil && !isGlueAlreadyExists(err) {
		return fmt.Errorf("glue CreateDatabase %s: %w", name, err)
	}
	return nil
}

func (c *Catalog) ensureTable(ctx context.Context, t ResolvedTable) error {
	got, err := c.glue.GetTable(ctx, &glue.GetTableInput{
		DatabaseName: aws.String(t.Database),
		Name:         aws.String(t.Name),
	})
	if err != nil {
		if !isGlueNotFound(err) {
			return fmt.Errorf("glue GetTable %s.%s: %w", t.Database, t.Name, err)
		}
		_, err = c.glue.CreateTable(ctx, buildCreateTableInput(t))
		if err != nil && !isGlueAlreadyExists(err) {
			return fmt.Errorf("glue CreateTable %s.%s: %w", t.Database, t.Name, err)
		}
		return nil
	}
	if !glueTableNeedsSync(got.Table, t) {
		return nil
	}
	_, err = c.glue.UpdateTable(ctx, buildUpdateTableInput(t))
	if err != nil {
		return fmt.Errorf("glue UpdateTable %s.%s: %w", t.Database, t.Name, err)
	}
	return nil
}

func (c *Catalog) tableExists(ctx context.Context, database, name string) (bool, error) {
	_, err := c.glue.GetTable(ctx, &glue.GetTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(name),
	})
	if err == nil {
		return true, nil
	}
	if isGlueNotFound(err) {
		return false, nil
	}
	return false, fmt.Errorf("glue GetTable %s.%s: %w", database, name, err)
}

func glueTableNeedsSync(existing *types.Table, want ResolvedTable) bool {
	if existing == nil {
		return true
	}
	wantParams := cloneMap(want.Definition.PartitionProjection)
	wantTpl := strings.TrimSpace(wantParams["storage.location.template"])
	haveTpl := ""
	if existing.Parameters != nil {
		haveTpl = strings.TrimSpace(existing.Parameters["storage.location.template"])
	}
	if wantTpl != "" && wantTpl != haveTpl {
		return true
	}
	wantLoc := strings.TrimRight(want.Location, "/")
	haveLoc := ""
	if existing.StorageDescriptor != nil && existing.StorageDescriptor.Location != nil {
		haveLoc = strings.TrimRight(aws.ToString(existing.StorageDescriptor.Location), "/")
	}
	if wantLoc != "" && wantLoc != haveLoc {
		return true
	}
	// Partition key set changed (e.g. slowlogs dropped date path partition).
	haveParts := map[string]struct{}{}
	for _, p := range existing.PartitionKeys {
		haveParts[aws.ToString(p.Name)] = struct{}{}
	}
	if len(haveParts) != len(want.Definition.PartitionColumns) {
		return true
	}
	for _, p := range want.Definition.PartitionColumns {
		if _, ok := haveParts[p.Name]; !ok {
			return true
		}
	}
	return false
}

func (c *Catalog) probeDataStatus(ctx context.Context, locationPrefix string) string {
	if c.s3 == nil {
		return DataStatusUnknown
	}
	bucket, prefix, err := parseS3URI(strings.TrimSuffix(locationPrefix, "/") + "/")
	if err != nil {
		return DataStatusUnknown
	}
	out, err := c.s3.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
		Bucket:  aws.String(bucket),
		Prefix:  aws.String(prefix),
		MaxKeys: aws.Int32(1),
	})
	if err != nil {
		return DataStatusUnknown
	}
	if len(out.Contents) > 0 || (out.KeyCount != nil && *out.KeyCount > 0) {
		return DataStatusAvailable
	}
	return DataStatusEmpty
}

func buildCreateTableInput(t ResolvedTable) *glue.CreateTableInput {
	return &glue.CreateTableInput{
		DatabaseName: aws.String(t.Database),
		TableInput:   buildTableInput(t),
	}
}

func buildUpdateTableInput(t ResolvedTable) *glue.UpdateTableInput {
	return &glue.UpdateTableInput{
		DatabaseName: aws.String(t.Database),
		TableInput:   buildTableInput(t),
	}
}

func buildTableInput(t ResolvedTable) *types.TableInput {
	columns := make([]types.Column, 0, len(t.Definition.FileColumns))
	for _, col := range t.Definition.FileColumns {
		columns = append(columns, types.Column{
			Name: aws.String(col.Name),
			Type: aws.String(col.Type),
		})
	}
	partitions := make([]types.Column, 0, len(t.Definition.PartitionColumns))
	for _, col := range t.Definition.PartitionColumns {
		partitions = append(partitions, types.Column{
			Name: aws.String(col.Name),
			Type: aws.String(col.Type),
		})
	}
	params := cloneMap(t.Definition.PartitionProjection)
	return &types.TableInput{
		Name:          aws.String(t.Name),
		TableType:     aws.String("EXTERNAL_TABLE"),
		Parameters:    params,
		PartitionKeys: partitions,
		StorageDescriptor: &types.StorageDescriptor{
			Columns:      columns,
			Location:     aws.String(t.Location),
			InputFormat:  aws.String("org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat"),
			OutputFormat: aws.String("org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat"),
			SerdeInfo: &types.SerDeInfo{
				SerializationLibrary: aws.String("org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe"),
			},
		},
	}
}

func isGlueNotFound(err error) bool {
	var ae smithy.APIError
	if errors.As(err, &ae) {
		return ae.ErrorCode() == "EntityNotFoundException"
	}
	return false
}

func isGlueAlreadyExists(err error) bool {
	var ae smithy.APIError
	if errors.As(err, &ae) {
		return ae.ErrorCode() == "AlreadyExistsException"
	}
	return false
}