package transformers

import (
	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"github.com/cloudquery/cloudquery/plugins/transformer/jsonflattener/client/recordupdater"
	"github.com/cloudquery/cloudquery/plugins/transformer/jsonflattener/client/spec"
	"github.com/cloudquery/cloudquery/plugins/transformer/jsonflattener/client/tablematcher"
	"github.com/rs/zerolog"
)

type TransformationFn = func(arrow.RecordBatch) (arrow.RecordBatch, error)
type SchemaTransformationFn = func(*arrow.Schema) (*arrow.Schema, error)

type Transformer struct {
	matcher  *tablematcher.TableMatcher
	fn       TransformationFn
	schemaFn SchemaTransformationFn
}

func NewFromSpec(logger zerolog.Logger, sp spec.Spec) (*Transformer, error) {
	tr := &Transformer{matcher: tablematcher.New(sp.Tables)}

	tr.fn = FlattenJSONFields(logger)
	tr.schemaFn = transformSchema(tr.fn)

	return tr, nil
}

func (tr *Transformer) Transform(record arrow.RecordBatch) (arrow.RecordBatch, error) {
	// Passthrough if the record's table is not a match to any of the spec's
	// tablepatterns, but error if the record doesn't have table metadata.
	isMatch, err := tr.matcher.IsSchemasTableMatch(record.Schema())
	if err != nil {
		return nil, err
	}
	if !isMatch {
		return record, nil
	}
	// Apply the specific transformation kind
	return tr.fn(record)
}

func (tr *Transformer) TransformSchema(schema *arrow.Schema) (*arrow.Schema, error) {
	// Passthrough if the record's table is not a match to any of the spec's
	// tablepatterns, but error if the record doesn't have table metadata.
	isMatch, err := tr.matcher.IsSchemasTableMatch(schema)
	if err != nil {
		return nil, err
	}
	if !isMatch {
		return schema, nil
	}

	if tr.schemaFn == nil {
		return schema, nil
	}
	return tr.schemaFn(schema)
}

func FlattenJSONFields(logger zerolog.Logger) TransformationFn {
	return func(record arrow.RecordBatch) (arrow.RecordBatch, error) {
		record, err := recordupdater.New(logger, record).FlattenJSONFields()
		if err != nil {
			return nil, err
		}
		return record, nil
	}
}

func transformSchema(tf TransformationFn) SchemaTransformationFn {
	return func(schema *arrow.Schema) (*arrow.Schema, error) {
		newRecord, err := tf(makeEmptyRecord(schema))
		if err != nil {
			return nil, err
		}
		return newRecord.Schema(), nil
	}
}

func makeEmptyRecord(s *arrow.Schema) arrow.RecordBatch {
	cols := make([]arrow.Array, 0, len(s.Fields()))
	for _, field := range s.Fields() {
		cols = append(cols, array.NewBuilder(memory.DefaultAllocator, field.Type).NewArray())
	}
	md := s.Metadata()
	return array.NewRecordBatch(arrow.NewSchema(s.Fields(), &md), cols, 0)
}
