package chain

import (
	"blockchain/config"
	"bytes"
	"sync"
)

// Snapshot is a read-only view of the chain state for API queries.

type Blockchain struct {
	blocks []*Block
	coins  CoinStore
	mx     sync.RWMutex
}

func NewBlockchain() *Blockchain {
	genesis := NewBlock(0, nil, "")
	genesis.Hash = genesis.CalcHash()
	return &Blockchain{
		blocks: []*Block{genesis},
		coins:  make(CoinStore),
	}
}

type Snapshot struct {
	Blocks []*Block
	Coins  CoinStore
}

// Snapshot returns a consistent read-only view of the chain state.
func (cs *Blockchain) Snapshot() Snapshot {
	cs.mx.RLock()
	defer cs.mx.RUnlock()

	blocks := make([]*Block, len(cs.blocks))
	copy(blocks, cs.blocks)

	return Snapshot{
		Blocks: blocks,
		Coins:  cs.coins,
	}
}

func (cs *Blockchain) LastBlock() *Block {
	cs.mx.RLock()
	defer cs.mx.RUnlock()

	return cs.blocks[len(cs.blocks)-1]
}

func (cs *Blockchain) AddBlock(b *Block) bool {
	cs.mx.Lock()
	defer cs.mx.Unlock()

	return cs.addBlock(b)
}

// addBlock validates and appends a block. Caller must hold write lock.
func (cs *Blockchain) addBlock(b *Block) bool {
	last := cs.blocks[len(cs.blocks)-1]
	if b.PrevHash != last.Hash || b.Height != last.Height+1 {
		return false
	}
	if !b.CheckProofOfWork() {
		return false
	}
	if !cs.validateTransactions(b.Transactions) {
		return false
	}
	cs.applyBlock(b)
	cs.blocks = append(cs.blocks, b)
	return true
}

func (cs *Blockchain) validateTransactions(txs []*Transaction) bool {
	working := cs.coins.Clone()
	coinbaseCount := 0

	for _, tx := range txs {
		if tx.IsCoinbase() {
			coinbaseCount++
			if coinbaseCount > 1 {
				return false
			}
			if len(tx.Coins) != 1 || tx.Coins[0].Amount != config.Default.Chain.MiningReward {
				return false
			}
			continue
		}

		if !cs.validateTx(tx, working) {
			return false
		}

		for _, s := range tx.Spends {
			working.Spend(s.CoinID)
		}
		txID := tx.TxID()
		for i, c := range tx.Coins {
			working.Add(txID, i, c)
		}
	}
	return true
}

func (cs *Blockchain) validateTx(tx *Transaction, coins CoinStore) bool {
	if !tx.Verify() {
		return false
	}

	inputSum := 0
	for _, s := range tx.Spends {
		coin := coins.Get(s.CoinID)
		if coin == nil {
			return false
		}
		if !bytes.Equal(coin.Owner, s.Owner) {
			return false
		}
		inputSum += coin.Amount
	}

	outputSum := 0
	for _, c := range tx.Coins {
		if c.Amount <= 0 {
			return false
		}
		outputSum += c.Amount
	}

	return outputSum <= inputSum
}

func (cs *Blockchain) applyBlock(b *Block) {
	for _, tx := range b.Transactions {
		for _, s := range tx.Spends {
			cs.coins.Spend(s.CoinID)
		}
		txID := tx.TxID()
		for i, c := range tx.Coins {
			cs.coins.Add(txID, i, c)
		}
	}
}

func (cs *Blockchain) rebuildCoins() {
	cs.coins = make(CoinStore)
	for _, b := range cs.blocks {
		cs.applyBlock(b)
	}
}

func (cs *Blockchain) ReplaceBlocks(blocks []*Block) bool {
	cs.mx.Lock()
	defer cs.mx.Unlock()

	if len(blocks) == 0 {
		return false
	}

	first := blocks[0]
	attachIdx := -1
	for i, b := range cs.blocks {
		if b.Hash == first.PrevHash {
			attachIdx = i
			break
		}
	}
	if attachIdx == -1 {
		return false
	}
	if attachIdx+1+len(blocks) <= len(cs.blocks) {
		return false
	}

	oldBlocks := make([]*Block, len(cs.blocks))
	copy(oldBlocks, cs.blocks)
	oldCoins := cs.coins

	cs.blocks = cs.blocks[:attachIdx+1]
	cs.rebuildCoins()

	for _, b := range blocks {
		if !cs.addBlock(b) {
			cs.blocks = oldBlocks
			cs.coins = oldCoins
			return false
		}
	}
	return true
}

// BlocksAfterHeight returns all blocks after the given height.
func (cs *Blockchain) BlocksAfterHeight(height int) []*Block {
	cs.mx.RLock()
	defer cs.mx.RUnlock()
	if height+1 >= len(cs.blocks) {
		return nil
	}
	result := make([]*Block, len(cs.blocks)-height-1)
	copy(result, cs.blocks[height+1:])
	return result
}