package pipeline

import (
	"context"
	"fmt"
	"strings"

	mErrors "github.com/caraml-dev/merlin/pkg/errors"
	"github.com/caraml-dev/merlin/pkg/transformer/spec"
	"github.com/caraml-dev/merlin/pkg/transformer/symbol"
	"github.com/caraml-dev/merlin/pkg/transformer/types"
	"github.com/caraml-dev/merlin/pkg/transformer/types/table"
	"go.opentelemetry.io/otel/attribute"
)

type TableJoinOp struct {
	tableJoinSpec *spec.TableJoin
	*OperationTracing
}

func NewTableJoinOp(tableJoinSpec *spec.TableJoin, tracingEnabled bool) Op {
	tableJoinOp := &TableJoinOp{
		tableJoinSpec: tableJoinSpec,
	}

	if tracingEnabled {
		tableJoinOp.OperationTracing = NewOperationTracing(tableJoinSpec, types.TableJoinOpType)
	}
	return tableJoinOp
}

func (t TableJoinOp) Execute(ctx context.Context, environment *Environment) error {
	_, span := tracer.Start(ctx, "pipeline.TableJoin")
	defer span.End()

	span.SetAttributes(attribute.String("table.left", t.tableJoinSpec.LeftTable))
	span.SetAttributes(attribute.String("table.right", t.tableJoinSpec.RightTable))
	span.SetAttributes(attribute.String("table.output", t.tableJoinSpec.OutputTable))
	span.SetAttributes(attribute.String("table.how", t.tableJoinSpec.How.String()))

	leftTable, err := getTable(environment, t.tableJoinSpec.LeftTable)
	if err != nil {
		return err
	}

	rightTable, err := getTable(environment, t.tableJoinSpec.RightTable)
	if err != nil {
		return err
	}

	var resultTable *table.Table

	joinColumns := []string{t.tableJoinSpec.OnColumn}
	if len(t.tableJoinSpec.OnColumns) > 0 {
		joinColumns = t.tableJoinSpec.OnColumns
	}

	switch t.tableJoinSpec.How {
	case spec.JoinMethod_LEFT:
		err := validateJoinColumns(leftTable, rightTable, joinColumns)
		if err != nil {
			return err
		}

		resultTable, err = leftTable.LeftJoin(rightTable, joinColumns)
		if err != nil {
			return err
		}
	case spec.JoinMethod_RIGHT:
		err := validateJoinColumns(leftTable, rightTable, joinColumns)
		if err != nil {
			return err
		}

		resultTable, err = leftTable.RightJoin(rightTable, joinColumns)
		if err != nil {
			return err
		}
	case spec.JoinMethod_INNER:
		err := validateJoinColumns(leftTable, rightTable, joinColumns)
		if err != nil {
			return err
		}

		resultTable, err = leftTable.InnerJoin(rightTable, joinColumns)
		if err != nil {
			return err
		}
	case spec.JoinMethod_OUTER:
		err := validateJoinColumns(leftTable, rightTable, joinColumns)
		if err != nil {
			return err
		}

		resultTable, err = leftTable.OuterJoin(rightTable, joinColumns)
		if err != nil {
			return err
		}
	case spec.JoinMethod_CROSS:
		resultTable, err = leftTable.CrossJoin(rightTable)
		if err != nil {
			return err
		}
	case spec.JoinMethod_CONCAT:
		resultTable, err = leftTable.Concat(rightTable)
		if err != nil {
			return err
		}
	default:
		return mErrors.NewInvalidInputErrorf("unknown join method: %s", t.tableJoinSpec.How)
	}

	environment.SetSymbol(t.tableJoinSpec.OutputTable, resultTable)
	if t.OperationTracing != nil {
		if err := t.AddInputOutput(
			map[string]interface{}{t.tableJoinSpec.LeftTable: leftTable, t.tableJoinSpec.RightTable: rightTable},
			map[string]interface{}{t.tableJoinSpec.OutputTable: resultTable},
		); err != nil {
			return nil
		}
	}
	environment.LogOperation("table_join", t.tableJoinSpec.OutputTable)
	return nil
}

func getTable(env *Environment, tableName string) (*table.Table, error) {
	return getTableFromSymbolRegistry(env.symbolRegistry, tableName)
}

func getTableFromSymbolRegistry(sr symbol.Registry, tableName string) (*table.Table, error) {
	tableRaw := sr[tableName]
	if tableRaw == nil {
		return nil, mErrors.NewInvalidInputErrorf("table '%s' is not declared", tableName)
	}

	tableOut, ok := tableRaw.(*table.Table)
	if !ok {
		return nil, mErrors.NewInvalidInputErrorf("variable '%s' is not a table", tableName)
	}
	return tableOut, nil
}

func validateJoinColumns(leftTable *table.Table, rightTable *table.Table, joinColumns []string) error {
	var leftTableColumns, rightTableColumns []string

	for _, joinColumn := range joinColumns {
		if _, err := leftTable.GetColumn(joinColumn); err != nil {
			leftTableColumns = append(leftTableColumns, joinColumn)
		}

		if _, err := rightTable.GetColumn(joinColumn); err != nil {
			rightTableColumns = append(rightTableColumns, joinColumn)
		}
	}

	if len(leftTableColumns) > 0 && len(rightTableColumns) > 0 {
		if strings.Join(leftTableColumns, ", ") == strings.Join(rightTableColumns, ", ") {
			return mErrors.NewInvalidInputErrorf("invalid join column: column %s does not exist in %s", strings.Join(leftTableColumns, ", "), "left table and right table")
		}
	}

	var errMessages []string
	if len(leftTableColumns) > 0 {
		missingColumns := strings.Join(leftTableColumns, ", ")
		errMessages = append(errMessages, fmt.Sprintf("column %s does not exist in %s", missingColumns, "left table"))
	}

	if len(rightTableColumns) > 0 {
		missingColumns := strings.Join(rightTableColumns, ", ")
		errMessages = append(errMessages, fmt.Sprintf("column %s does not exist in %s", missingColumns, "right table"))
	}

	if len(errMessages) > 0 {
		return mErrors.NewInvalidInputErrorf("invalid join column: %s", strings.Join(errMessages, " and "))
	}

	return nil
}
