package traces

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/ethclient"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Tracer extracts and stores transaction traces
type Tracer struct {
	db      *pgxpool.Pool
	client  *ethclient.Client
	chainID int
}

// NewTracer creates a new tracer
func NewTracer(db *pgxpool.Pool, client *ethclient.Client, chainID int) *Tracer {
	return &Tracer{
		db:      db,
		client:  client,
		chainID: chainID,
	}
}

// Trace represents a transaction trace
type Trace struct {
	TransactionHash common.Hash
	BlockNumber     int64
	Traces          []CallTrace
}

// CallTrace represents a single call in a trace
type CallTrace struct {
	Type         string
	From         common.Address
	To           common.Address
	Value        string
	Gas          uint64
	GasUsed      uint64
	Input        string
	Output       string
	Error        string
	Calls        []CallTrace
}

// ExtractTrace extracts trace for a transaction
func (t *Tracer) ExtractTrace(ctx context.Context, txHash common.Hash, blockNumber int64) error {
	// Call trace_block RPC method
	var result json.RawMessage
	err := t.client.Client().CallContext(ctx, &result, "debug_traceTransaction", txHash.Hex(), map[string]interface{}{
		"tracer": "callTracer",
		"tracerConfig": map[string]interface{}{
			"withLog": true,
		},
	})

	if err != nil {
		return fmt.Errorf("failed to trace transaction: %w", err)
	}

	// Parse trace
	var trace CallTrace
	if err := json.Unmarshal(result, &trace); err != nil {
		return fmt.Errorf("failed to parse trace: %w", err)
	}

	// Store trace
	return t.storeTrace(ctx, txHash, blockNumber, trace)
}

// storeTrace stores trace in database
func (t *Tracer) storeTrace(ctx context.Context, txHash common.Hash, blockNumber int64, trace CallTrace) error {
	// Create traces table if it doesn't exist
	// For now, store as JSONB
	query := `
		CREATE TABLE IF NOT EXISTS traces (
			chain_id INTEGER NOT NULL,
			transaction_hash VARCHAR(66) NOT NULL,
			block_number BIGINT NOT NULL,
			trace_data JSONB NOT NULL,
			created_at TIMESTAMP DEFAULT NOW(),
			PRIMARY KEY (chain_id, transaction_hash)
		) PARTITION BY LIST (chain_id)
	`

	_, err := t.db.Exec(ctx, query)
	if err != nil {
		// Table might already exist
	}

	// Insert trace
	insertQuery := `
		INSERT INTO traces (chain_id, transaction_hash, block_number, trace_data)
		VALUES ($1, $2, $3, $4)
		ON CONFLICT (chain_id, transaction_hash) DO UPDATE SET
			trace_data = $4,
			created_at = NOW()
	`

	traceJSON, err := json.Marshal(trace)
	if err != nil {
		return fmt.Errorf("failed to marshal trace: %w", err)
	}

	_, err = t.db.Exec(ctx, insertQuery, t.chainID, txHash.Hex(), blockNumber, traceJSON)
	return err
}