package client

import (
	"crypto/sha1"

	"storrent/bt"
	"storrent/metainfo"
)

type pieceScheduler struct {
	pieces     []piece
	incomplete []int
	verified   int
	meta       *metainfo.File
	dir        string
}

func newPieceScheduler(m *metainfo.File, dir string) *pieceScheduler {
	n := len(m.Info.Pieces)
	pieces := make([]piece, n)
	incomplete := make([]int, 0, n)
	verified := 0

	for i := range n {
		size := m.Info.PieceSize
		if i == n-1 {
			size = m.Size - int64(n-1)*m.Info.PieceSize
		}
		nblocks := (size + bt.BlockSize - 1) / bt.BlockSize
		pieces[i] = piece{
			hash:    m.Info.Pieces[i],
			size:    size,
			data:    make([]byte, size),
			have:    make([]bool, nblocks),
			reqs:    make([]bool, nblocks),
			reqPeer: make([]*bt.Peer, nblocks),
		}

		data := m.Info.Read(dir, i, pieces[i].size)
		if data != nil && sha1.Sum(data) == pieces[i].hash {
			pieces[i].data = data
			pieces[i].done = true
			verified++
		} else {
			incomplete = append(incomplete, i)
		}
	}

	return &pieceScheduler{
		pieces:     pieces,
		incomplete: incomplete,
		verified:   verified,
		meta:       m,
		dir:        dir,
	}
}

func (ps *pieceScheduler) nextRequest(p *bt.Peer) *bt.Msg {
	for _, idx := range ps.incomplete {
		pc := &ps.pieces[idx]
		if pc.done {
			continue
		}
		if !p.HasPiece(idx) {
			continue
		}
		if begin, length, ok := pc.next(p); ok {
			return &bt.Msg{
				Kind:   bt.Request,
				Index:  uint32(idx),
				Begin:  uint32(begin),
				Length: uint32(length),
			}
		}
	}
	return nil
}

func (ps *pieceScheduler) put(index int, begin int, block []byte) bool {
	pc := &ps.pieces[index]
	if pc.put(begin, block) {
		ps.meta.Info.Write(ps.dir, index, pc.data)
		ps.verified++
		ps.removeIncomplete(index)
		return true
	}
	return false
}

func (ps *pieceScheduler) removeIncomplete(index int) {
	for i, idx := range ps.incomplete {
		if idx == index {
			ps.incomplete = append(ps.incomplete[:i], ps.incomplete[i+1:]...)
			return
		}
	}
}

func (ps *pieceScheduler) peerDisconnected(p *bt.Peer) {
	for i := range ps.pieces {
		ps.pieces[i].resetPeer(p)
	}
}

func (ps *pieceScheduler) isComplete() bool {
	return ps.verified == len(ps.pieces)
}

func (ps *pieceScheduler) progress() (verified, total int) {
	return ps.verified, len(ps.pieces)
}

func (ps *pieceScheduler) pieceData(index int) []byte {
	pc := &ps.pieces[index]
	if !pc.done {
		return nil
	}
	return pc.data
}

func (ps *pieceScheduler) count() int {
	return len(ps.pieces)
}

func (ps *pieceScheduler) bitfield(all bool) []byte {
	n := len(ps.pieces)
	bf := make([]byte, (n+7)/8)
	for i := range ps.pieces {
		if all || ps.pieces[i].done {
			bt.SetBit(bf, i)
		}
	}
	return bf
}