package telemetry

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-sql-driver/mysql"
	pkgerrors "github.com/nexuscrm/kernel/pkg/errors"
	"github.com/nexuscrm/kernel/pkg/perftiming"
	"github.com/nexuscrm/shared/pkg/constants"
	"github.com/rs/zerolog/log"
)

func init() {
	sql.Register("mysql-tracing", &TracingDriver{Parent: mysql.MySQLDriver{}})
}

// TracingDriver wraps the MySQL driver to inject trace context into SQL comments
type TracingDriver struct { //nolint:revive
	Parent mysql.MySQLDriver
}

// SQLStatementCounter is an opt-in, request-scoped observation of statements
// that reach this driver through QueryContext or ExecContext. The counter
// records only aggregate counts and durations. An independently enabled
// diagnostic event may include the SQL template, but never its bound arguments.
// The counter is deliberately carried by context so one Runtime request cannot
// observe another request's work.
type SQLStatementCounter struct {
	count      atomic.Uint64
	totalNanos atomic.Int64
	maxNanos   atomic.Int64
}

type sqlStatementCounterKey struct{}

// WithSQLStatementCounter binds a new counter to ctx. Callers should read it
// only after the observed operation returns.
func WithSQLStatementCounter(ctx context.Context) (context.Context, *SQLStatementCounter) {
	counter := &SQLStatementCounter{}
	counters, _ := ctx.Value(sqlStatementCounterKey{}).([]*SQLStatementCounter)
	bound := make([]*SQLStatementCounter, 0, len(counters)+1)
	bound = append(bound, counters...)
	bound = append(bound, counter)
	return context.WithValue(ctx, sqlStatementCounterKey{}, bound), counter
}

func (c *SQLStatementCounter) Count() uint64 {
	if c == nil {
		return 0
	}
	return c.count.Load()
}

func (c *SQLStatementCounter) TotalDuration() time.Duration {
	if c == nil {
		return 0
	}
	return time.Duration(c.totalNanos.Load())
}

func (c *SQLStatementCounter) MaxDuration() time.Duration {
	if c == nil {
		return 0
	}
	return time.Duration(c.maxNanos.Load())
}

type sqlStatementObservation struct {
	start     time.Time
	counters  []*SQLStatementCounter
	ctx       context.Context
	query     string
	operation string
	started   atomic.Bool
	once      sync.Once
}

func beginSQLStatement(ctx context.Context, operation, query string) *sqlStatementObservation {
	counters, _ := ctx.Value(sqlStatementCounterKey{}).([]*SQLStatementCounter)
	if len(counters) == 0 && !perftiming.Enabled(ctx) {
		return nil
	}
	return &sqlStatementObservation{
		start: time.Now(), counters: counters, ctx: ctx,
		operation: operation, query: query,
	}
}

func (o *sqlStatementObservation) markStarted() {
	if o == nil || !o.started.CompareAndSwap(false, true) {
		return
	}
	for _, counter := range o.counters {
		if counter != nil {
			counter.count.Add(1)
		}
	}
}

func (o *sqlStatementObservation) finish(err error) {
	if o == nil {
		return
	}
	o.once.Do(func() {
		o.markStarted()
		elapsed := time.Since(o.start)
		for _, counter := range o.counters {
			if counter == nil {
				continue
			}
			counter.totalNanos.Add(elapsed.Nanoseconds())
			for observed := counter.maxNanos.Load(); elapsed.Nanoseconds() > observed; observed = counter.maxNanos.Load() {
				if counter.maxNanos.CompareAndSwap(observed, elapsed.Nanoseconds()) {
					break
				}
			}
		}
		if perftiming.Enabled(o.ctx) {
			event := log.Info()
			status := "success"
			if err != nil && !errors.Is(err, io.EOF) {
				event = log.Error().Err(err)
				status = "error"
			}
			event.
				Str("event", "sql_statement_timing").
				Str("request_id", perftiming.RequestID(o.ctx)).
				Str("status", status).
				Str("operation", o.operation).
				Int64("elapsed_ms", elapsed.Milliseconds()).
				Str("sql", diagnosticSQL(o.query)).
				Msg("SQL statement timing")
		}
	})
}

func diagnosticSQL(query string) string {
	query = strings.Join(strings.Fields(query), " ")
	const maxLength = 4000
	if len(query) <= maxLength {
		return query
	}
	return query[:maxLength] + "...<truncated>"
}

// TracingStmt preserves database/sql's prepared-statement execution path.
// go-sql-driver/mysql deliberately returns driver.ErrSkip from connection
// QueryContext for parameterized statements, so the real network/result-set
// work happens here rather than in TracingConn.QueryContext.
type TracingStmt struct {
	driver.Stmt
	query string
}

func (s *TracingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := s.Stmt.(driver.StmtExecContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	observation := beginSQLStatement(ctx, "exec", s.query)
	result, err := execer.ExecContext(ctx, args)
	if errors.Is(err, driver.ErrSkip) {
		return nil, driver.ErrSkip
	}
	observation.markStarted()
	observation.finish(err)
	return result, err
}

func (s *TracingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := s.Stmt.(driver.StmtQueryContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	observation := beginSQLStatement(ctx, "query", s.query)
	rows, err := queryer.QueryContext(ctx, args)
	if errors.Is(err, driver.ErrSkip) {
		return nil, driver.ErrSkip
	}
	if err != nil || rows == nil {
		observation.finish(err)
		return rows, err
	}
	observation.markStarted()
	if observation == nil {
		return rows, nil
	}
	return &observedRows{Rows: rows, observation: observation}, nil
}

func (s *TracingStmt) CheckNamedValue(value *driver.NamedValue) error {
	if checker, ok := s.Stmt.(driver.NamedValueChecker); ok {
		return checker.CheckNamedValue(value)
	}
	return driver.ErrSkip
}

func (s *TracingStmt) ColumnConverter(index int) driver.ValueConverter {
	if converter, ok := s.Stmt.(driver.ColumnConverter); ok {
		return converter.ColumnConverter(index)
	}
	return driver.DefaultParameterConverter
}

// observedRows keeps a query observation open until its result set has really
// been consumed. QueryContext returning a driver.Rows value does not mean the
// database/network work is complete: most row data is read by Next. Measuring
// only QueryContext therefore makes remote result-set time look like zero.
type observedRows struct {
	driver.Rows
	observation *sqlStatementObservation
}

func (r *observedRows) Close() error {
	err := r.Rows.Close()
	r.observation.finish(err)
	return err
}

func (r *observedRows) Next(dest []driver.Value) error {
	err := r.Rows.Next(dest)
	if err != nil {
		r.observation.finish(err)
	}
	return err
}

func (r *observedRows) HasNextResultSet() bool {
	rows, ok := r.Rows.(driver.RowsNextResultSet)
	return ok && rows.HasNextResultSet()
}

func (r *observedRows) NextResultSet() error {
	rows, ok := r.Rows.(driver.RowsNextResultSet)
	if !ok {
		r.observation.finish(io.EOF)
		return io.EOF
	}
	err := rows.NextResultSet()
	if err != nil {
		r.observation.finish(err)
	}
	return err
}

func (r *observedRows) ColumnTypeScanType(index int) reflect.Type {
	if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok {
		return rows.ColumnTypeScanType(index)
	}
	return reflect.TypeOf(new(any)).Elem()
}

func (r *observedRows) ColumnTypeDatabaseTypeName(index int) string {
	if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok {
		return rows.ColumnTypeDatabaseTypeName(index)
	}
	return ""
}

func (r *observedRows) ColumnTypeNullable(index int) (bool, bool) {
	if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok {
		return rows.ColumnTypeNullable(index)
	}
	return false, false
}

func (r *observedRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) {
	if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok {
		return rows.ColumnTypePrecisionScale(index)
	}
	return 0, 0, false
}

// Open returns a wrapped connection
func (d *TracingDriver) Open(name string) (driver.Conn, error) {
	conn, err := d.Parent.Open(name)
	if err != nil {
		return nil, err
	}
	return &TracingConn{Conn: conn}, nil
}

// OpenConnector preserves database/sql's context-aware connection creation.
// Without this optional driver contract, database/sql falls back to Open and a
// canceled request cannot stop an in-flight network dial or handshake.
func (d *TracingDriver) OpenConnector(name string) (driver.Connector, error) {
	parent, err := d.Parent.OpenConnector(name)
	if err != nil {
		return nil, err
	}
	return &tracingConnector{parent: parent, driver: d}, nil
}

type tracingConnector struct {
	parent driver.Connector
	driver *TracingDriver
}

func (c *tracingConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.parent.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return &TracingConn{Conn: conn}, nil
}

func (c *tracingConnector) Driver() driver.Driver {
	return c.driver
}

// TracingConn wraps a driver.Conn
type TracingConn struct {
	driver.Conn
}

type tracingTx struct {
	driver.Tx
	ctx     context.Context
	started time.Time
	id      uint64
	once    sync.Once
}

var tracingTransactionSequence atomic.Uint64

func (t *tracingTx) Commit() error {
	started := time.Now()
	err := t.Tx.Commit()
	t.finish("commit", started, err)
	return err
}

func (t *tracingTx) Rollback() error {
	started := time.Now()
	err := t.Tx.Rollback()
	t.finish("rollback", started, err)
	return err
}

func (t *tracingTx) finish(stage string, stageStarted time.Time, err error) {
	if t == nil {
		return
	}
	t.once.Do(func() {
		perftiming.Event(t.ctx, "sql_transaction_timing", stageStarted, err).
			Uint64("driver_transaction_id", t.id).
			Str("stage", stage).
			Int64("transaction_elapsed_ms", time.Since(t.started).Milliseconds()).
			Msg("SQL transaction timing")
	})
}

// Prepare preserves callers that use the legacy driver.Conn contract. The
// context-aware database/sql path uses PrepareContext below, but both entry
// points must return the same observed statement wrapper.
func (c *TracingConn) Prepare(query string) (driver.Stmt, error) {
	stmt, err := c.Conn.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &TracingStmt{Stmt: stmt, query: query}, nil
}

// ResetSession preserves the optional driver.SessionResetter contract exposed
// by the wrapped connection. database/sql calls this before reusing a pooled
// connection, allowing the MySQL driver to reject a connection it has closed.
func (c *TracingConn) ResetSession(ctx context.Context) error {
	if resetter, ok := c.Conn.(driver.SessionResetter); ok {
		return resetter.ResetSession(ctx)
	}
	return nil
}

// IsValid preserves the optional driver.Validator contract exposed by the
// wrapped connection. database/sql calls this before returning a connection to
// the pool and discards connections the underlying driver has marked closed.
func (c *TracingConn) IsValid() bool {
	if validator, ok := c.Conn.(driver.Validator); ok {
		return validator.IsValid()
	}
	return true
}

// CheckNamedValue preserves driver-specific value conversion. ErrSkip asks
// database/sql to use its default converter when the wrapped driver does not
// implement the optional interface.
func (c *TracingConn) CheckNamedValue(value *driver.NamedValue) error {
	if checker, ok := c.Conn.(driver.NamedValueChecker); ok {
		return checker.CheckNamedValue(value)
	}
	return driver.ErrSkip
}

// Ping implements driver.Pinger
func (c *TracingConn) Ping(ctx context.Context) error {
	if pinger, ok := c.Conn.(driver.Pinger); ok {
		return pinger.Ping(ctx)
	}
	return nil
}

// BeginTx implements driver.ConnBeginTx
func (c *TracingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	started := time.Now()
	var tx driver.Tx
	var err error
	if connBeginTx, ok := c.Conn.(driver.ConnBeginTx); ok {
		tx, err = connBeginTx.BeginTx(ctx, opts)
	} else {
		tx, err = c.Conn.Begin()
	}
	if !perftiming.Enabled(ctx) {
		return tx, err
	}
	transactionID := tracingTransactionSequence.Add(1)
	perftiming.Event(ctx, "sql_transaction_timing", started, err).
		Uint64("driver_transaction_id", transactionID).
		Str("stage", "begin").
		Msg("SQL transaction timing")
	if err != nil || tx == nil {
		return tx, err
	}
	return &tracingTx{Tx: tx, ctx: ctx, started: started, id: transactionID}, nil
}

// PrepareContext implements driver.ConnPrepareContext
func (c *TracingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	diagnosticQuery := query
	query = injectTraceComment(ctx, query)
	if connPrepareCtx, ok := c.Conn.(driver.ConnPrepareContext); ok {
		stmt, err := connPrepareCtx.PrepareContext(ctx, query)
		if err != nil {
			return nil, err
		}
		return &TracingStmt{Stmt: stmt, query: diagnosticQuery}, nil
	}
	stmt, err := c.Conn.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &TracingStmt{Stmt: stmt, query: diagnosticQuery}, nil
}

// QueryContext implements driver.QueryerContext
func (c *TracingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	observation := beginSQLStatement(ctx, "query", query)
	query = injectTraceComment(ctx, query)
	if queryer, ok := c.Conn.(driver.QueryerContext); ok {
		rows, err := queryer.QueryContext(ctx, query, args)
		if errors.Is(err, driver.ErrSkip) {
			return nil, driver.ErrSkip
		}
		if err != nil || rows == nil {
			observation.finish(err)
			return rows, err
		}
		observation.markStarted()
		if observation == nil {
			return rows, nil
		}
		return &observedRows{Rows: rows, observation: observation}, nil
	}
	err := pkgerrors.NewInternalError("underlying connection does not support QueryContext", nil)
	observation.finish(err)
	return nil, err
}

// ExecContext implements driver.ExecerContext
func (c *TracingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	observation := beginSQLStatement(ctx, "exec", query)
	query = injectTraceComment(ctx, query)
	if execer, ok := c.Conn.(driver.ExecerContext); ok {
		result, err := execer.ExecContext(ctx, query, args)
		if errors.Is(err, driver.ErrSkip) {
			return nil, driver.ErrSkip
		}
		observation.markStarted()
		observation.finish(err)
		return result, err
	}
	err := pkgerrors.NewInternalError("underlying connection does not support ExecContext", nil)
	observation.finish(err)
	return nil, err
}

// injectTraceComment appends the W3C traceparent to the SQL query
func injectTraceComment(ctx context.Context, query string) string {
	// Use W3C Traceparent only
	if val := ctx.Value(constants.ContextKeyTraceparent); val != nil {
		if str, ok := val.(string); ok && str != "" {
			// Prevent SQL injection via header (basic sanitization)
			clean := strings.ReplaceAll(str, "*/", "")
			return fmt.Sprintf("%s /* traceparent='%s' */", query, clean)
		}
	}
	return query
}