package store

import (
	"context"
	"database/sql/driver"
	"time"

	"modernc.org/sqlite"

	"g.fc.run/ai/byre/internal/trace"
)

type sqliteConn interface {
	driver.Conn
	driver.ConnBeginTx
	driver.ConnPrepareContext
	driver.ExecerContext
	driver.QueryerContext
	driver.Pinger
	driver.SessionResetter
	driver.Validator
}

type tracedConnector struct {
	driver.Connector
}

func newTracedConnector(dsn string) (driver.Connector, error) {
	c, err := sqlite.NewConnector(dsn)
	if err != nil {
		return nil, err
	}
	return tracedConnector{c}, nil
}

func (c tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.Connector.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return tracedConn{conn.(sqliteConn)}, nil
}

type tracedConn struct {
	sqliteConn
}

func (c tracedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	defer trace.Start(ctx, "sql begin")()
	return c.sqliteConn.BeginTx(ctx, opts)
}

func (c tracedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	defer trace.Start(ctx, "sql exec")()
	return c.sqliteConn.ExecContext(ctx, query, args)
}

func (c tracedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	start := time.Now()
	rows, err := c.sqliteConn.QueryContext(ctx, query, args)
	if err != nil {
		trace.Add(ctx, "sql query", time.Since(start))
		return nil, err
	}
	return &tracedRows{sqliteRows: rows.(sqliteRows), ctx: ctx, took: time.Since(start)}, nil
}

type sqliteRows interface {
	driver.Rows
	driver.RowsColumnTypeDatabaseTypeName
	driver.RowsColumnTypeLength
	driver.RowsColumnTypeNullable
	driver.RowsColumnTypePrecisionScale
	driver.RowsColumnTypeScanType
}

type tracedRows struct {
	sqliteRows
	ctx  context.Context
	took time.Duration
}

func (r *tracedRows) Next(dest []driver.Value) error {
	start := time.Now()
	err := r.sqliteRows.Next(dest)
	r.took += time.Since(start)
	return err
}

func (r *tracedRows) Close() error {
	start := time.Now()
	err := r.sqliteRows.Close()
	trace.Add(r.ctx, "sql query", r.took+time.Since(start))
	return err
}