package driver

import (
	"context"

	"entgo.io/ent"
	"entgo.io/ent/dialect"
	"entgo.io/ent/dialect/sql"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
)

type MultiDriver struct {
	RW               dialect.Driver
	Name, User, Host string
}

var _ dialect.Driver = (*MultiDriver)(nil)

func (d *MultiDriver) Query(ctx context.Context, query string, args, v any) error {
	e := d.RW
	span, ctx := opentracing.StartSpanFromContext(ctx,
		query,
		opentracing.Tag{Key: string(ext.Component), Value: "ent"},
		ext.SpanKindRPCClient,
	)
	defer span.Finish()
	span.SetTag(`args`, args)
	span.SetTag(`db.statement`, query)
	span.SetTag(`labels.callType`, e.Dialect())
	span.SetTag(`labels.dbName`, d.Name)
	span.SetTag(`labels.host`, d.Host)

	// Mutation statements that use the RETURNING clause.
	if ent.QueryFromContext(ctx) == nil {
		span.SetTag(`query_for_write`, false)
	} else {
		span.SetTag(`query_for_read`, true)
	}
	return e.Query(ctx, query, args, v)
}

func (d *MultiDriver) Exec(ctx context.Context, query string, args, v any) error {
	span, ctx := opentracing.StartSpanFromContext(ctx,
		query,
		opentracing.Tag{Key: string(ext.Component), Value: "ent"},
		ext.SpanKindRPCClient,
	)
	span.SetTag(`labels.callType`, d.RW.Dialect())
	span.SetTag(`labels.dbName`, d.Name)
	span.SetTag(`labels.host`, d.Host)
	defer span.Finish()
	span.SetTag(`args`, args)
	span.SetTag(`db.statement`, query)
	return d.RW.Exec(ctx, query, args, v)
}

func (d *MultiDriver) Tx(ctx context.Context) (dialect.Tx, error) {
	span, ctx := opentracing.StartSpanFromContext(ctx,
		`TX`,
		opentracing.Tag{Key: string(ext.Component), Value: "ent"},
		ext.SpanKindRPCClient,
	)
	span.SetTag(`labels.callType`, d.RW.Dialect())
	span.SetTag(`labels.dbName`, d.Name)
	span.SetTag(`labels.host`, d.Host)
	defer span.Finish()
	span.SetTag(`开启事务`, true)
	return d.RW.Tx(ctx)
}

func (d *MultiDriver) BeginTx(ctx context.Context, opts *sql.TxOptions) (dialect.Tx, error) {
	span, ctx := opentracing.StartSpanFromContext(ctx,
		`BeginTx`,
		opentracing.Tag{Key: string(ext.Component), Value: "ent"},
		ext.SpanKindRPCClient,
	)
	defer span.Finish()
	span.SetTag(`labels.callType`, d.RW.Dialect())
	span.SetTag(`labels.dbName`, d.Name)
	span.SetTag(`labels.host`, d.Host)
	span.SetTag(`开启事务`, true)
	return d.RW.(interface {
		BeginTx(context.Context, *sql.TxOptions) (dialect.Tx, error)
	}).BeginTx(ctx, opts)
}

func (d *MultiDriver) Close() error {
	err := d.RW.Close()
	if err != nil {
		return err
	}
	return nil
}

func (d *MultiDriver) Dialect() string {
	return d.RW.Dialect()
}
