package aleo

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

const (
	// 网络常量
	TcpEndpoint = "127.0.0.1:3030"
	RpcEndpoint = "http://127.0.0.1:4133"

	// 消息类型
	MsgHandshake           uint16 = 1
	MsgHandshakeResponse   uint16 = 2
	MsgPuzzleRequest       uint16 = 9
	MsgPuzzleResponse      uint16 = 10
	MsgUnconfirmedSolution uint16 = 11

	// 协议版本
	ProtocolVersion uint32 = 17

	// 挖矿参数
	MaxPuzzleInstances = 6 // 最大puzzle实例数
)

// Prover 挖矿节点结构
type Prover struct {
	conn         net.Conn
	minerAddress []byte

	// 最新状态
	latestEpochHash   []byte
	latestBlockHeader *BlockHeader
	epochHashMutex    sync.RWMutex
	headerMutex       sync.RWMutex

	// puzzle实例计数
	puzzleInstances    atomic.Int32
	maxPuzzleInstances uint32

	// 停止信号
	shutdown chan struct{}
}

// BlockHeader 区块头结构
type BlockHeader struct {
	Height           uint32
	PreviousHash     [32]byte
	Timestamp        uint64
	CoinbaseTarget   uint64
	ProofTarget      uint64
	TransactionsRoot [32]byte
	StateRoot        [32]byte
}

// Solution 挖矿解结构
type Solution struct {
	ID        [32]byte
	EpochHash []byte
	Nonce     uint64
	Address   []byte
	Proof     []byte
}

func NewProver(minerAddress []byte) (*Prover, error) {
	// 计算最大puzzle实例数
	maxInstances := uint32(runtime.NumCPU() - 2)
	if maxInstances < 1 {
		maxInstances = 1
	}
	if maxInstances > MaxPuzzleInstances {
		maxInstances = MaxPuzzleInstances
	}

	prover := &Prover{
		minerAddress:       minerAddress,
		maxPuzzleInstances: maxInstances,
		shutdown:           make(chan struct{}),
	}

	// 连接节点
	if err := prover.connect(); err != nil {
		return nil, err
	}

	return prover, nil
}

func (p *Prover) connect() error {
	conn, err := net.Dial("tcp", TcpEndpoint)
	if err != nil {
		return fmt.Errorf("连接失败: %v", err)
	}
	p.conn = conn

	// 执行握手
	if err := p.performHandshake(); err != nil {
		conn.Close()
		return fmt.Errorf("握手失败: %v", err)
	}

	return nil
}

// mining/message.go
func (p *Prover) sendMessage(msgType uint16, payload []byte) error {
	header := make([]byte, 6)
	binary.BigEndian.PutUint16(header[0:2], msgType)
	binary.BigEndian.PutUint32(header[2:6], ProtocolVersion)

	if _, err := p.conn.Write(header); err != nil {
		return err
	}

	if payload != nil {
		if _, err := p.conn.Write(payload); err != nil {
			return err
		}
	}

	return nil
}

func (p *Prover) receiveMessage() (uint16, []byte, error) {
	header := make([]byte, 6)
	if _, err := io.ReadFull(p.conn, header); err != nil {
		return 0, nil, err
	}

	msgType := binary.BigEndian.Uint16(header[0:2])
	// protocolVersion := binary.BigEndian.Uint32(header[2:6])

	var payload []byte
	if msgType == MsgPuzzleResponse {
		// 读取epoch hash
		epochHash := make([]byte, 32)
		if _, err := io.ReadFull(p.conn, epochHash); err != nil {
			return msgType, nil, err
		}

		// 读取block header
		lenBuf := make([]byte, 4)
		if _, err := io.ReadFull(p.conn, lenBuf); err != nil {
			return msgType, nil, err
		}
		headerLen := binary.BigEndian.Uint32(lenBuf)

		blockHeader := make([]byte, headerLen)
		if _, err := io.ReadFull(p.conn, blockHeader); err != nil {
			return msgType, nil, err
		}

		payload = append(epochHash, blockHeader...)
	}

	return msgType, payload, nil
}

func (p *Prover) getPuzzleTask() error {
	// 发送请求
	if err := p.sendMessage(MsgPuzzleRequest, nil); err != nil {
		return err
	}

	// 接收响应
	msgType, data, err := p.receiveMessage()
	if err != nil {
		return err
	}
	if msgType != MsgPuzzleResponse {
		return fmt.Errorf("unexpected message type: %d", msgType)
	}

	// 解析数据
	epochHash := data[0:32]
	header, err := parseBlockHeader(data[32:])
	if err != nil {
		return err
	}

	// 更新状态
	p.epochHashMutex.Lock()
	p.latestEpochHash = epochHash
	p.epochHashMutex.Unlock()

	p.headerMutex.Lock()
	p.latestBlockHeader = header
	p.headerMutex.Unlock()

	return nil
}

func (p *Prover) StartMining() error {
	// 启动多个puzzle实例
	for i := uint32(0); i < p.maxPuzzleInstances; i++ {
		go p.puzzleLoop()
	}

	// 定期获取新任务
	go p.taskLoop()

	return nil
}

func (p *Prover) puzzleLoop() {
	p.puzzleInstances.Add(1)
	defer p.puzzleInstances.Add(-1)

	for {
		select {
		case <-p.shutdown:
			return
		default:
			// 获取当前状态
			p.epochHashMutex.RLock()
			epochHash := p.latestEpochHash
			p.epochHashMutex.RUnlock()

			p.headerMutex.RLock()
			header := p.latestBlockHeader
			p.headerMutex.RUnlock()

			if epochHash != nil && header != nil {
				// 尝试找到解
				solution, err := p.findSolution(epochHash, header.ProofTarget)
				if err == nil {
					// 提交解
					if err := p.submitSolution(solution); err != nil {
						log.Printf("提交解失败: %v", err)
					}
				}
			}

			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (p *Prover) taskLoop() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-p.shutdown:
			return
		case <-ticker.C:
			if err := p.getPuzzleTask(); err != nil {
				log.Printf("获取任务失败: %v", err)
			}
		}
	}
}

func (p *Prover) submitSolution(solution *Solution) error {
	// 序列化solution
	var buf bytes.Buffer

	buf.Write(solution.ID[:])
	buf.Write(solution.EpochHash)

	nonceBytes := make([]byte, 8)
	binary.LittleEndian.PutUint64(nonceBytes, solution.Nonce)
	buf.Write(nonceBytes)

	buf.Write(solution.Address)
	buf.Write(solution.Proof)

	// 发送消息
	return p.sendMessage(MsgUnconfirmedSolution, buf.Bytes())
}

func main() {
	// 创建矿工地址
	minerAddress := []byte("your_miner_address_here")

	// 创建Prover
	prover, err := NewProver(minerAddress)
	if err != nil {
		log.Fatalf("创建Prover失败: %v", err)
	}

	// 启动挖矿
	if err := prover.StartMining(); err != nil {
		log.Fatalf("启动挖矿失败: %v", err)
	}

	// 等待中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	// 优雅关闭
	close(prover.shutdown)
}

func (p *Prover) findSolution(epochHash []byte, proofTarget uint64) (*Solution, error) {
	var nonce uint64 = 0
	maxNonce := ^uint64(0) // 最大uint64值

	// 将proofTarget转换为字节数组用于比较
	targetBytes := make([]byte, 32)
	targetBytes[0] = byte(proofTarget)

	for nonce < maxNonce {
		// 计算proof
		proof, err := p.computeProof(epochHash, p.minerAddress, nonce)
		if err != nil {
			return nil, fmt.Errorf("计算proof失败: %v", err)
		}

		// 验证proof是否满足难度要求
		if p.verifyProof(proof, targetBytes) {
			// 构造解
			solution := &Solution{
				EpochHash: epochHash,
				Nonce:     nonce,
				Address:   p.minerAddress,
				Proof:     proof,
			}

			// 计算解的ID
			id, err := p.computeSolutionID(solution)
			if err != nil {
				return nil, fmt.Errorf("计算解ID失败: %v", err)
			}
			copy(solution.ID[:], id)

			log.Printf("找到有效解 - Nonce: %d, ID: %x", nonce, id)
			return solution, nil
		}

		// 每处理1000个nonce检查一次是否需要退出
		if nonce%1000 == 0 {
			select {
			case <-p.shutdown:
				return nil, fmt.Errorf("挖矿已停止")
			default:
			}
		}

		nonce++
	}

	return nil, fmt.Errorf("未找到有效解")
}

// computeProof 计算工作量证明
func (p *Prover) computeProof(epochHash []byte, address []byte, nonce uint64) ([]byte, error) {
	hasher := sha256.New()

	// 写入epoch hash
	if _, err := hasher.Write(epochHash); err != nil {
		return nil, fmt.Errorf("写入epoch hash失败: %v", err)
	}

	// 写入地址
	if _, err := hasher.Write(address); err != nil {
		return nil, fmt.Errorf("写入地址失败: %v", err)
	}

	// 写入nonce (小端字节序)
	nonceBytes := make([]byte, 8)
	binary.LittleEndian.PutUint64(nonceBytes, nonce)
	if _, err := hasher.Write(nonceBytes); err != nil {
		return nil, fmt.Errorf("写入nonce失败: %v", err)
	}

	return hasher.Sum(nil), nil
}

// verifyProof 验证工作量证明是否满足难度要求
func (p *Prover) verifyProof(proof []byte, target []byte) bool {
	// 比较proof和target
	// proof必须小于target才是有效的
	for i := 0; i < len(proof) && i < len(target); i++ {
		if proof[i] < target[i] {
			return true
		}
		if proof[i] > target[i] {
			return false
		}
	}
	return false
}

// computeSolutionID 计算解的唯一标识符
func (p *Prover) computeSolutionID(solution *Solution) ([]byte, error) {
	hasher := sha256.New()

	// 写入所有字段
	if _, err := hasher.Write(solution.EpochHash); err != nil {
		return nil, fmt.Errorf("写入epoch hash失败: %v", err)
	}

	if _, err := hasher.Write(solution.Address); err != nil {
		return nil, fmt.Errorf("写入地址失败: %v", err)
	}

	nonceBytes := make([]byte, 8)
	binary.LittleEndian.PutUint64(nonceBytes, solution.Nonce)
	if _, err := hasher.Write(nonceBytes); err != nil {
		return nil, fmt.Errorf("写入nonce失败: %v", err)
	}

	if _, err := hasher.Write(solution.Proof); err != nil {
		return nil, fmt.Errorf("写入proof失败: %v", err)
	}

	return hasher.Sum(nil), nil
}

type PuzzleInstances struct {
	count uint32
}

// Sub 原子减少实例计数
func (p *PuzzleInstances) Sub(delta uint32) {
	atomic.AddUint32(&p.count, ^(delta - 1))
}

// Add 原子增加实例计数
func (p *PuzzleInstances) Add(delta uint32) {
	atomic.AddUint32(&p.count, delta)
}

// Get 获取当前实例数
func (p *PuzzleInstances) Get() uint32 {
	return atomic.LoadUint32(&p.count)
}

func parseBlockHeader(data []byte) (*BlockHeader, error) {
	if len(data) < 96 { // 最小区块头长度
		return nil, fmt.Errorf("区块头数据太短")
	}

	reader := bytes.NewReader(data)
	header := &BlockHeader{}

	// 读取高度 (4字节)
	if err := binary.Read(reader, binary.LittleEndian, &header.Height); err != nil {
		return nil, fmt.Errorf("读取区块高度失败: %v", err)
	}

	// 读取前一个区块哈希 (32字节)
	if _, err := io.ReadFull(reader, header.PreviousHash[:]); err != nil {
		return nil, fmt.Errorf("读取前一个区块哈希失败: %v", err)
	}

	// 读取时间戳 (8字节)
	if err := binary.Read(reader, binary.LittleEndian, &header.Timestamp); err != nil {
		return nil, fmt.Errorf("读取时间戳失败: %v", err)
	}

	// 读取coinbase目标值 (8字节)
	if err := binary.Read(reader, binary.LittleEndian, &header.CoinbaseTarget); err != nil {
		return nil, fmt.Errorf("读取coinbase目标值失败: %v", err)
	}

	// 读取proof目标值 (8字节)
	if err := binary.Read(reader, binary.LittleEndian, &header.ProofTarget); err != nil {
		return nil, fmt.Errorf("读取proof目标值失败: %v", err)
	}

	// 读取交易根 (32字节)
	if _, err := io.ReadFull(reader, header.TransactionsRoot[:]); err != nil {
		return nil, fmt.Errorf("读取交易根失败: %v", err)
	}

	// 读取状态根 (32字节)
	if _, err := io.ReadFull(reader, header.StateRoot[:]); err != nil {
		return nil, fmt.Errorf("读取状态根失败: %v", err)
	}

	return header, nil
}

func (p *Prover) performHandshake() error {
	// 发送握手消息
	if err := p.sendMessage(MsgHandshake, nil); err != nil {
		return fmt.Errorf("发送握手消息失败: %v", err)
	}

	// 设置读取超时
	if err := p.conn.SetReadDeadline(time.Now().Add(10 * time.Second)); err != nil {
		return fmt.Errorf("设置读取超时失败: %v", err)
	}
	defer p.conn.SetReadDeadline(time.Time{}) // 清除超时

	// 接收握手响应
	msgType, payload, err := p.receiveMessage()
	if err != nil {
		return fmt.Errorf("接收握手响应失败: %v", err)
	}

	// 验证消息类型
	if msgType != MsgHandshakeResponse {
		return fmt.Errorf("收到意外的消息类型: %d, 期望类型: %d", msgType, MsgHandshakeResponse)
	}

	// 验证协议版本
	if len(payload) >= 4 {
		version := binary.BigEndian.Uint32(payload[0:4])
		if version != ProtocolVersion {
			return fmt.Errorf("协议版本不匹配: 收到 %d, 期望 %d", version, ProtocolVersion)
		}
	}

	log.Printf("握手成功 - 协议版本: %d", ProtocolVersion)
	return nil
}