package state

import (
	"context"
	"fmt"
	"math/big"
	"sync"
	"sync/atomic"
	"time"

	"github.com/ethereum/go-ethereum"
	"github.com/ethereum/go-ethereum/accounts/abi/bind"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/log"

	shastaBindings "github.com/taikoxyz/taiko-mono/packages/taiko-client/bindings/shasta"
	"github.com/taikoxyz/taiko-mono/packages/taiko-client/internal/metrics"
	"github.com/taikoxyz/taiko-mono/packages/taiko-client/pkg/rpc"
)

const (
	// l1HeadPollInterval is the cadence at which the driver polls the L1 head
	// when running against an HTTP-only L1 endpoint. Tuned well below an L1
	// slot so the driver picks up new heads within one slot.
	l1HeadPollInterval = 3 * time.Second
	// l2HeadPollInterval is the cadence at which the driver polls the L2 head
	// when running against an HTTP-only L2 endpoint. Set faster than the L2
	// block time so the driver tracks the L2 chain with sub-block latency.
	l2HeadPollInterval = 500 * time.Millisecond
)

// State contains all states which will be used by driver.
type State struct {
	// Feeds
	l1HeadsFeed event.Feed // L1 new heads notification feed

	l1Head    atomic.Value // Latest known L1 head
	l2Head    atomic.Value // Current L2 execution engine's local chain head
	l1Current atomic.Value // Current L1 block sync cursor

	// Constants
	GenesisL1Height *big.Int

	// RPC clients
	rpc *rpc.Client

	wg sync.WaitGroup
}

// New creates a new driver state instance.
func New(ctx context.Context, rpc *rpc.Client) (*State, error) {
	s := &State{rpc: rpc}

	if err := s.init(ctx); err != nil {
		return nil, fmt.Errorf("failed to initialize driver state: %w", err)
	}

	go s.eventLoop(ctx)

	return s, nil
}

// Close closes all inner subscriptions.
func (s *State) Close() {
	s.wg.Wait()
}

// init fetches the latest status and initializes the state instance.
func (s *State) init(ctx context.Context) error {
	if err := s.initGenesisHeight(ctx); err != nil {
		return fmt.Errorf("failed to initialize genesis height: %w", err)
	}

	log.Info("Genesis L1 height", "height", s.GenesisL1Height)

	// L1 head
	l1Head, err := s.rpc.L1.HeaderByNumber(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to get L1 head: %w", err)
	}
	s.setL1Head(l1Head)

	// L2 head
	l2Head, err := s.rpc.L2.HeaderByNumber(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to get L2 head: %w", err)
	}

	log.Info("L2 execution engine head", "blockID", l2Head.Number, "hash", l2Head.Hash())
	s.setL2Head(l2Head)

	// Set the L2 head's latest known L1 origin as current L1 sync cursor.
	if err := s.initL1Current(ctx, l2Head); err != nil {
		return fmt.Errorf("failed to initialize L1 current cursor: %w", err)
	}

	return nil
}

// initL1Current initializes the L1 current cursor from the L2 head's latest known L1 origin.
func (s *State) initL1Current(ctx context.Context, l2Head *types.Header) error {
	headL1Origin, err := s.rpc.L2.HeadL1Origin(ctx)
	if err != nil && err.Error() != ethereum.NotFound.Error() {
		return fmt.Errorf("failed to get head L1 origin: %w", err)
	}

	// The L2 chain is non-empty but has no head L1 origin recorded: the execution engine
	// was synced without the driver writing L1 origins (a beacon sync interrupted by a
	// restart before the first event sync insertion, or a snapshot restore without
	// L1Origin data). Derive the cursor from the head block's anchor instead of falling
	// back to the activation block, which would rescan the whole proposal history.
	// A preconfirmation-tip head is safe here: preconfirmation blocks carry the same
	// Shasta extra-data layout, proposal IDs <= 1 resolve to the activation block inside
	// ResetL1Current, and any derivation failure falls through to the previous fallback.
	if headL1Origin == nil && l2Head.Number.Sign() > 0 {
		resetErr := s.ResetL1Current(ctx, l2Head.Number)
		if resetErr == nil {
			return nil
		}
		log.Warn(
			"Failed to derive the L1 current cursor from the L2 head, use the last known L1 origin instead",
			"l2Head", l2Head.Number,
			"error", resetErr,
		)
	}

	latestL2KnownL1Header, err := s.rpc.LatestL2KnownL1Header(ctx)
	if err != nil {
		return fmt.Errorf("failed to get latest L2 known L1 header: %w", err)
	}
	s.l1Current.Store(latestL2KnownL1Header)

	return nil
}

// eventLoop initializes and starts all subscriptions and callbacks in the given state instance.
func (s *State) eventLoop(ctx context.Context) {
	s.wg.Add(1)
	defer s.wg.Done()

	var (
		// Channels for subscriptions.
		l1HeadCh = make(chan *types.Header, 10)
		l2HeadCh = make(chan *types.Header, 10)
		provedCh = make(chan *shastaBindings.ShastaInboxClientProved, 10)

		// Subscriptions. L2 head uses a tighter polling cadence than L1 when
		// running against HTTP-only endpoints; L1-derived events use the L1 cadence.
		l1HeadSub         = rpc.SubscribeChainHead(s.rpc.L1, l1HeadCh, l1HeadPollInterval)
		l2HeadSub         = rpc.SubscribeChainHead(s.rpc.L2, l2HeadCh, l2HeadPollInterval)
		l2ProvedShastaSub = rpc.SubscribeProved(s.rpc.L1, s.rpc.ShastaClients.Inbox, provedCh)
	)

	defer func() {
		l1HeadSub.Unsubscribe()
		l2HeadSub.Unsubscribe()
		l2ProvedShastaSub.Unsubscribe()
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case e := <-provedCh:
			coreState, err := s.rpc.GetCoreState(&bind.CallOpts{Context: ctx})
			if err != nil {
				log.Error("Failed to get core state", "err", err)
				continue
			}
			header, err := s.rpc.L2.HeaderByHash(ctx, coreState.LastFinalizedBlockHash)
			if err != nil {
				log.Error("Failed to get finalized block header", "err", err)
				continue
			}
			log.Info(
				"📈 Proposals proven and verified",
				"firstProposalID", e.FirstNewProposalId,
				"lastProposalID", e.LastProposalId,
				"checkpointNumber", header.Number,
				"checkpointHash", common.Hash(coreState.LastFinalizedBlockHash),
			)
		case newHead := <-l1HeadCh:
			s.setL1Head(newHead)
			s.l1HeadsFeed.Send(newHead)
		case newHead := <-l2HeadCh:
			s.setL2Head(newHead)
		}
	}
}

// setL1Head sets the L1 head concurrent safely.
func (s *State) setL1Head(l1Head *types.Header) {
	if l1Head == nil {
		log.Warn("Empty new L1 head")
		return
	}

	log.Debug("New L1 head", "height", l1Head.Number, "hash", l1Head.Hash(), "timestamp", l1Head.Time)
	metrics.DriverL1HeadHeightGauge.Set(float64(l1Head.Number.Int64()))

	s.l1Head.Store(l1Head)
}

// GetL1Head reads the L1 head concurrent safely.
func (s *State) GetL1Head() *types.Header {
	return s.l1Head.Load().(*types.Header)
}

// SetL2Head exposes setL2Head to callers outside this package.
// Used by the chain syncer to refresh the cached L2 head from RPC when the
// `newHeads` subscription has not observed an engine-API-driven backfill.
func (s *State) SetL2Head(l2Head *types.Header) {
	s.setL2Head(l2Head)
}

// setL2Head sets the L2 head concurrent safely.
func (s *State) setL2Head(l2Head *types.Header) {
	if l2Head == nil {
		log.Warn("Empty new L2 head")
		return
	}

	log.Trace("New L2 head", "blockID", l2Head.Number, "hash", l2Head.Hash(), "timestamp", l2Head.Time)

	s.l2Head.Store(l2Head)
}

// GetL2Head reads the L2 head concurrent safely.
func (s *State) GetL2Head() *types.Header {
	return s.l2Head.Load().(*types.Header)
}

// SubL1HeadsFeed registers a subscription of new L1 heads.
func (s *State) SubL1HeadsFeed(ch chan *types.Header) event.Subscription {
	return s.l1HeadsFeed.Subscribe(ch)
}

// initGenesisHeight fetches the L1 activation height for the current inbox.
func (s *State) initGenesisHeight(ctx context.Context) error {
	// GetActivationBlockNumber returns the L1 block number whose timestamp matches the
	// inbox activation timestamp.
	genesisHeight, err := s.rpc.GetActivationBlockNumber(ctx)
	if err != nil {
		return fmt.Errorf("failed to get activation block number: %w", err)
	}

	s.GenesisL1Height = genesisHeight
	return nil
}
