// Package models provides the models for the KermaGo application
package models

import (
	"crypto/ed25519"
	"crypto/x509"
	"errors"
	"kerma/helpers"
	"os"
	"strings"

	"github.com/docker/go/canonical/json"

	"golang.org/x/exp/maps"
	"golang.org/x/exp/slices"
)

/*
A State is an object that contains the current state of the application.
State-related fields include:
- PeerList - list of discovered peers
- Chain - chain of known valid blocks
- Transactions - hashmap of known valid transactions
- PrivateKey - ED25519 private key of the Kerma node
*/
type State struct {
	PeerList     []Peer
	Validator    helpers.Validator
	Config       helpers.Config
	Chain        []Block
	Transactions map[string]*Transaction
	PrivateKey   ed25519.PrivateKey
	Logger       helpers.Logger
	Height       uint64
}

/*
A Peer represents a peer on the Kerma network. It has a name and an indicator of whether there was a handshake
*/
type Peer struct {
	Name   string
	Active bool
}

/*
A PeerListResponse represents an object with a type and a string array, containing the addresses of all known peers
*/
type PeerListResponse struct {
	Type  string   `json:"type" binding:"required"`
	Peers []string `json:"peers"`
}

/*
Construct creates a new state object
*/
func (s *State) Construct() {
	s.PeerList = []Peer{}
	s.Validator.Construct()
	s.Config.Construct()

	privKey, err := s.parsePKIFile()
	if err != nil {
		s.Logger.Warn(err.Error())
		s.Logger.Info("Generating new key...")
		_, privKey, _ = ed25519.GenerateKey(nil)
	}
	s.PrivateKey = privKey
	s.dumpPKIFile()

	s.Logger.Info("Reading known peers from dump file")
	dumpErr := s.parsePeerListStore()

	if dumpErr != nil {
		s.Logger.Warn(dumpErr.Error() + ", reverting to defaults...")
		s.parseInitialPeerList()
	}

	s.Logger.Info("Restoring previous known state")
	s.Chain = []Block{}
	s.Transactions = map[string]*Transaction{}

	terr := s.parseTransactionStore()
	if terr != nil {
		s.Logger.Debug(terr.Error())
	}

	berr := s.parseBlockStore()
	if berr != nil {
		s.Logger.Debug(berr.Error())
	}

	if len(s.Chain) == 0 {
		s.Chain = append(s.Chain, *GetGenesisBlock())
		s.Height = 0
	}
}

/*
==============================

	HANDSHAKE OPERATIONS

==============================
*/

/*
CheckForHandshake checks whether the node has completed a handshake with the specified peer
*/
func (s *State) CheckForHandshake(peerName string) bool {
	for _, peer := range s.PeerList {
		if peer.Name == peerName {
			return peer.Active
		}
	}
	return false
}

/*
CompleteHandshake completes a handshake with the specified peer
*/
func (s *State) CompleteHandshake(peerName string) {
	peer := s.FindPeerByName(peerName)

	if peer != nil {
		peer.Active = true
	} else {
		peer = &Peer{
			Name:   peerName,
			Active: true,
		}
		s.PeerList = append(s.PeerList, *peer)
	}
}

/*
======================

	PEER OPERATIONS

======================
*/

/*
MarshalJson returns the canonical json of the PeerListResponse object
*/
func (p *PeerListResponse) MarshalJson() ([]byte, error) {
	return json.MarshalCanonical(p)
}

/*
FetchPeerListResponse returns a new PeerListResponse object from the local peer list
*/
func (s *State) FetchPeerListResponse() PeerListResponse {
	list := PeerListResponse{
		Type:  "peers",
		Peers: []string{},
	}

	for _, p := range s.PeerList {
		list.Peers = append(list.Peers, p.Name)
	}

	return list
}

/*
FindPeerByName returns the specified peer from the peer list
*/
func (s *State) FindPeerByName(peerName string) *Peer {
	for p, peer := range s.PeerList {
		if peer.Name == peerName {
			return &s.PeerList[p]
		}
	}
	return nil
}

/*
RemovePeerByName removes the specified peer from the peer list
*/
func (s *State) RemovePeerByName(peerName string) {
	for i, peer := range s.PeerList {
		if peer.Name == peerName {
			slices.Delete(s.PeerList, i, i)
		}
	}
	s.DumpPeerListStore()
}

/*
ParsePeerListResponse parses the given PeerListResponse object to the local peer list
*/
func (s *State) ParsePeerListResponse(p *PeerListResponse) {
	for _, peer := range p.Peers {
		if s.Validator.IsValidPeerName(peer) && s.FindPeerByName(peer) == nil {
			p := Peer{
				Name:   peer,
				Active: false,
			}
			s.PeerList = append(s.PeerList, p)
		}
	}
	s.DumpPeerListStore()
}

/*
DumpPeerListStore stores the local peer list to a JSON file
*/
func (s *State) DumpPeerListStore() {
	list := []string{}
	for _, p := range s.PeerList {
		if s.Validator.IsValidPeerName(p.Name) {
			list = append(list, p.Name)
		}
	}
	file, _ := json.MarshalCanonical(list)
	_ = os.WriteFile(s.Config.PeerListStore, file, 0644)
}

func (s *State) parsePeerListStore() error {
	if s.Validator.CheckIfFileExists(s.Config.PeerListStore) {
		file, err := os.ReadFile(s.Config.PeerListStore)
		if err != nil {
			return errors.New("cannot read peer list store " + s.Config.PeerListStore)
		}
		peerList := []string{}
		err = json.Unmarshal([]byte(file), &peerList)
		if err != nil {
			return errors.New("cannot parse peer list store at " + s.Config.PeerListStore)
		}
		for _, peer := range peerList {
			if s.Validator.IsValidPeerName(peer) {
				p := Peer{
					Name:   peer,
					Active: false,
				}
				s.PeerList = append(s.PeerList, p)
			}
		}
		return nil
	}
	return errors.New("cannot find peer list store at " + s.Config.PeerListStore)
}

func (s *State) parseInitialPeerList() {
	s.PeerList = []Peer{}

	if strings.Contains(s.Config.InitialPeerList, ",") {
		peers := strings.Split(s.Config.InitialPeerList, ",")
		for _, peerName := range peers {
			if s.Validator.IsValidPeerName(peerName) {
				p := Peer{
					Name:   peerName,
					Active: false,
				}
				s.PeerList = append(s.PeerList, p)
			} else {
				s.Logger.Warn("Failed parsing the initial peer list. Please check your configuration")
			}
		}
	} else {
		if s.Validator.IsValidPeerName(s.Config.InitialPeerList) {
			p := Peer{
				Name:   s.Config.InitialPeerList,
				Active: false,
			}
			s.PeerList = append(s.PeerList, p)
		} else {
			s.Logger.Warn("Failed parsing the initial peer list. Please check your configuration")
		}
	}
}

/*
======================

	PKI OPERATIONS

======================
*/

func (s *State) parsePKIFile() (ed25519.PrivateKey, error) {
	privKeyBytes, err := os.ReadFile(s.Config.PrivateKeyFile)
	if err != nil {
		return nil, errors.New("cannot read private key file at " + s.Config.PrivateKeyFile)
	}

	res, err := x509.ParsePKCS8PrivateKey(privKeyBytes)
	if err != nil {
		return nil, errors.New("cannot parse private key from " + s.Config.PrivateKeyFile)
	}

	privKey := res.(ed25519.PrivateKey)

	return privKey, nil
}

func (s *State) dumpPKIFile() {
	privKeyBytes, _ := x509.MarshalPKCS8PrivateKey(s.PrivateKey)

	_ = os.WriteFile(s.Config.PrivateKeyFile, privKeyBytes, 0600)
}

/*
======================

	BLOCK OPERATIONS

======================
*/

func (s *State) validateBlock(b *Block) error {
	err := b.Validate()
	if err != nil {
		return err
	}

	if b.IsGenesisBlock() {
		return nil
	}

	bid := b.GetID()

	prevBlock := s.Chain[len(s.Chain)-1]
	prevHash, err := prevBlock.Hash()

	if err != nil {
		return err
	}

	if *b.Previd != prevHash || prevBlock.Created > b.Created {
		return errors.New("invalid chain extension, block invalid: " + bid)
	}

	cbid := "invalid"
	for i, txid := range b.Txids {
		tx := s.Transactions[txid]
		if tx != nil {
			err = tx.Validate()
			if err != nil {
				return errors.New("detected invalid transaction, block invalid: " + bid)
			}

			if tx.IsCoinbaseTx() {
				if i != 0 {
					return errors.New("invalid coinbase transaction, block invalid: " + bid)
				}

				if tx.Outputs[0].Value > s.getTransactionFeesSum(b)+GetBlockReward() {
					return errors.New("law of conservation for coinbase transaction broken, block invalid: " + bid)
				}

				txBlock := s.findBlockForTransaction(txid)
				if txBlock != nil && *tx.Height != txBlock.Height {
					return errors.New("incorrect height on coinbase transaction, block invalid: " + bid)
				}

				cbid, _ = tx.Hash()

			} else {
				if cbid != "invalid" && tx.IsTxInput(cbid) {
					return errors.New("coinbase transaction cannot be spent in same block, block invalid: " + bid)
				}
			}
		}
	}

	return nil
}

/*
AppendToChain adds the specified block to the chain
*/
func (s *State) AppendToChain(b *Block) error {
	if len(s.Chain) == 0 {
		s.Chain = append(s.Chain, *GetGenesisBlock())
		s.Height = 0
	}

	bid, err := b.Hash()
	if err != nil {
		s.Logger.Debug(err.Error())
		return err
	}

	if s.GetBlock(bid) != nil {
		s.Logger.Debug("Block already in chain: " + bid)
		return nil
	}

	s.Logger.Debug("Validating block: " + bid)
	err = s.validateBlock(b)
	if err != nil {
		return err
	}

	if b.Previd != nil {
		prevBlock := s.GetBlock(*b.Previd)
		b.Height = prevBlock.Height + 1
	} else {
		if b.IsGenesisBlock() {
			b.Height = GetGenesisBlock().Height
		} else {
			return errors.New("invalid genesis block")
		}
	}

	if b.Height > s.Height {
		s.Logger.Info("New longest chain extension found: " + bid)
		s.Height = b.Height
		err = s.computeUTXOSet(b)

		if err == nil {
			s.Chain = append(s.Chain, *b)
			for _, txid := range b.Txids {
				s.confirmTransaction(txid)
			}
			s.DumpTransactionStore()
		} else {
			s.Logger.Debug(err.Error())
			return err
		}
	}
	return nil
}

func (s *State) computeUTXOSet(b *Block) error {
	if b.IsGenesisBlock() {
		b.UTXOSet = GetGenesisBlock().UTXOSet
		return nil
	}

	s.Logger.Debug("Computing UTXO set for block")
	prevBlock := s.GetBlock(*b.Previd)

	b.UTXOSet = make(map[string]uint64)
	maps.Copy(prevBlock.UTXOSet, b.UTXOSet)

	for _, txid := range b.Txids {
		s.Logger.Debug("Computing UTXO for transaction " + txid)
		transaction := s.GetTransaction(txid)
		err := transaction.Validate()

		if err != nil {
			return errors.New("invalid utxo set for block " + b.GetID() + ", transaction " + txid)
		}

		if transaction.IsCoinbaseTx() {
			output := transaction.Outputs[0]
			b.UTXOSet[output.Pubkey] = output.Value
		} else {
			for _, input := range transaction.Inputs {
				output := s.getOutputFromTransactionInput(&input)
				var balance uint64
				balance = output.Value
				if val, ok := b.UTXOSet[output.Pubkey]; ok {
					balance = val + output.Value
				}

				for _, o := range transaction.Outputs {
					balance -= o.Value
				}

				if balance > 0 {
					b.UTXOSet[output.Pubkey] = balance
				} else {
					delete(b.UTXOSet, output.Pubkey)
				}
			}
		}
	}

	return nil
}

/*
GetBlock returns the block specified by the supplied ID or nil
*/
func (s *State) GetBlock(bid string) *Block {
	for _, b := range s.Chain {
		if bid == b.GetID() {
			return &b
		}
	}
	return nil
}

/*
GetChainTip returns the tip of our local chain as an object
*/
func (s *State) GetChainTip() *Block {
	if len(s.Chain) == 0 {
		return GetGenesisBlock()
	}

	return &s.Chain[len(s.Chain)-1]
}

func (s *State) getTransactionFeesSum(b *Block) uint64 {
	var fees uint64
	fees = 0

	for _, txid := range b.Txids {
		fees += s.getTransactionFees(txid)
	}

	return fees
}

/*
GetMissingTransactionsInBlock returns the transaction ids from the block that are not present in the mempool
*/
func (s *State) GetMissingTransactionsInBlock(b *Block) ([]string, error) {
	err := b.Validate()
	if err != nil {
		return []string{}, err
	}

	res := []string{}

	for _, txid := range b.Txids {
		tx := s.Transactions[txid]
		if tx == nil {
			res = append(res, txid)
		}
	}

	return res, nil
}

/*
DumpBlockStore stores the local block chain to a JSON file
*/
func (s *State) DumpBlockStore() {
	file, _ := json.MarshalCanonical(s.Chain)
	_ = os.WriteFile(s.Config.BlockStore, file, 0644)
}

func (s *State) parseBlockStore() error {
	if s.Validator.CheckIfFileExists(s.Config.BlockStore) {
		file, err := os.ReadFile(s.Config.BlockStore)
		if err != nil {
			return errors.New("cannot read block store " + s.Config.BlockStore)
		}
		chain := []Block{}
		err = json.Unmarshal([]byte(file), &chain)
		if err != nil {
			return errors.New("cannot parse block store at " + s.Config.BlockStore)
		}

		for i, block := range chain {
			if block.IsGenesisBlock() {
				continue
			}
			block.Height = uint64(i)
			s.AppendToChain(&block)
		}
		return nil
	}
	return errors.New("cannot find block store at " + s.Config.BlockStore)
}

/*
=============================

	TRANSACTION OPERATIONS

=============================
*/

func (s *State) validateTransaction(t *Transaction) error {
	txid, err := t.Hash()

	if err != nil {
		s.Logger.Debug(err.Error())
		return err
	}

	if len(t.Outputs) == 0 {
		return errors.New("outputs cannot be empty, transaction invalid: " + txid)
	}

	if t.Height != nil {
		s.Logger.Debug("Processing a coinbase transaction: " + txid)
	} else {
		inputValues := 0
		outputValues := 0

		if len(t.Inputs) == 0 {
			return errors.New("inputs cannot be empty in a non-coinbase transaction, transaction invalid: " + txid)
		}

		for _, input := range t.Inputs {
			prevTransaction := s.Transactions[input.Outpoint.Txid]
			if prevTransaction == nil {
				return errors.New("could not find output transaction in pool, transaction invalid: " + txid)
			}

			if len(prevTransaction.Outputs)-1 < int(input.Outpoint.Index) {
				return errors.New("could not find output specified by index in pool, transaction invalid: " + txid)
			}

			outpoint := prevTransaction.Outputs[int(input.Outpoint.Index)]

			if !t.VerifySign(input.Sig, outpoint.Pubkey) {
				return errors.New("could not validate input signatures, transaction invalid: " + txid)
			}

			inputValues += int(outpoint.Value)
		}

		for _, output := range t.Outputs {
			outputValues += int(output.Value)
		}

		if inputValues != 0 && outputValues > inputValues {
			return errors.New("overspending, transaction invalid: " + txid)
		}
	}
	return nil
}

/*
AppendTransaction adds the specified transaction to the store
*/
func (s *State) AppendTransaction(t *Transaction) error {
	txid, err := t.Hash()
	if err != nil {
		s.Logger.Debug(err.Error())
		return err
	}

	s.Logger.Debug("Validating transaction: " + txid)

	err = s.validateTransaction(t)
	if err != nil {
		s.Logger.Debug(err.Error())
		return err
	}

	s.Logger.Debug("Appending transaction to store: " + txid)

	if t.IsCoinbaseTx() {
		t.Confirmed = true
	} else {
		t.Confirmed = false
	}

	s.Transactions[txid] = t
	return nil
}

func (s *State) findBlockForTransaction(txid string) *Block {
	for _, block := range s.Chain {
		if block.ContainsTransaction(txid) {
			return &block
		}
	}
	return nil
}

/*
GetTransaction returns the specified transaction from the store
*/
func (s *State) GetTransaction(txid string) *Transaction {
	if v, ok := s.Transactions[txid]; ok {
		return v
	}
	return nil
}

/*
GetMempoolTransactionIDs returns the list of transaction IDs in the mempool
*/
func (s *State) GetMempoolTransactionIDs() []string {
	res := []string{}

	for txid, tx := range s.Transactions {
		if !tx.Confirmed {
			res = append(res, txid)
		}
	}

	return res
}

func (s *State) getConfirmedTransactions() map[string]*Transaction {
	res := map[string]*Transaction{}

	for txid, tx := range s.Transactions {
		if tx.Confirmed {
			res[txid] = tx
		}
	}

	return res
}

func (s *State) confirmTransaction(txid string) {
	tx := s.GetTransaction(txid)
	if tx != nil {
		tx.Confirmed = true
	}
}

func (s *State) getTransactionFees(txid string) uint64 {
	var fees uint64
	fees = 0

	transaction := s.GetTransaction(txid)

	if transaction.IsCoinbaseTx() {
		return 0
	}

	for _, input := range transaction.Inputs {
		prevOutput := s.getOutputFromTransactionInput(&input)
		fees += prevOutput.Value
	}

	for _, output := range transaction.Outputs {
		fees -= output.Value
	}

	return fees
}

func (s *State) getOutputFromTransactionInput(i *Input) *Output {
	txid := i.Outpoint.Txid
	index := i.Outpoint.Index

	outputTx := s.GetTransaction(txid)
	for i, tx := range outputTx.Outputs {
		if int64(i) == index {
			return &tx
		}
	}

	return nil
}

/*
DumpTransactionStore stores the local store to a JSON file
*/
func (s *State) DumpTransactionStore() {
	file, _ := json.MarshalCanonical(s.getConfirmedTransactions())
	_ = os.WriteFile(s.Config.TransactionStore, file, 0644)
}

func (s *State) parseTransactionStore() error {
	if s.Validator.CheckIfFileExists(s.Config.TransactionStore) {
		file, err := os.ReadFile(s.Config.TransactionStore)
		if err != nil {
			return errors.New("cannot read transaction store " + s.Config.TransactionStore)
		}
		var pool map[string]*Transaction

		err = json.Unmarshal([]byte(file), &pool)
		if err != nil {
			return errors.New("cannot parse transaction store at " + s.Config.TransactionStore)
		}
		for id, transaction := range pool {
			s.Logger.Debug("Appending local transaction to store: " + id + ", tx: " + transaction.String())
			transaction.Confirmed = true
			s.Transactions[id] = transaction
		}
		return nil
	}
	return errors.New("cannot find transaction store at " + s.Config.TransactionStore)
}