package torrent

import (
	"context"
	"sync"
)

type pieceTask struct {
	Index int
}

type assignPieceRequest struct {
	peerID     string
	have       []bool
	hasInfo    bool
	responseCh chan assignPieceResponse
}

type assignPieceResponse struct {
	task pieceTask
	ok   bool
}

type reportPieceRequest struct {
	pieceIndex int
	success    bool
	responseCh chan bool
}

// cancelPeerRequest отправляет сигнал конкретному пиру отменить кусок в endgame.
type cancelPeerRequest struct {
	pieceIndex int
	peerID     string
}

type pieceScheduler struct {
	assignCh      chan assignPieceRequest
	reportCh      chan reportPieceRequest
	releasePeerCh chan string
	progressCh    chan int
	doneCh        chan struct{}
	stopCh        chan struct{}
	stopOnce      sync.Once

	setSequentialCh chan bool

	// cancelCh рассылает сигнал отмены конкретному воркеру в endgame-режиме.
	// Ключ — peerID, значение — канал с индексом куска для отмены.
	cancelMu  sync.RWMutex
	cancelSubs map[string]chan int
}

type pieceState uint8

const (
	piecePending    pieceState = iota
	pieceInProgress            // качается одним пиром
	pieceDone
)

// endgameThreshold — количество оставшихся кусков, при котором включается endgame.
const endgameThreshold = 4

func newPieceScheduler(pieceCount int) *pieceScheduler {
	return newPieceSchedulerWithResume(pieceCount, nil)
}

func newPieceSchedulerWithResume(pieceCount int, completedIndices []int) *pieceScheduler {
	ps := &pieceScheduler{
		assignCh:        make(chan assignPieceRequest, 128),
		reportCh:        make(chan reportPieceRequest, 128),
		releasePeerCh:   make(chan string, 128),
		progressCh:      make(chan int, 1),
		doneCh:          make(chan struct{}),
		stopCh:          make(chan struct{}),
		setSequentialCh: make(chan bool),
		cancelSubs:      make(map[string]chan int),
	}

	go ps.run(pieceCount, completedIndices)
	return ps
}

// SubscribeCancel регистрирует канал для получения cancel-сигналов в endgame-режиме.
// Воркер должен вызвать это перед Acquire, и отписаться после завершения.
func (ps *pieceScheduler) SubscribeCancel(peerID string) <-chan int {
	ch := make(chan int, 8)
	ps.cancelMu.Lock()
	ps.cancelSubs[peerID] = ch
	ps.cancelMu.Unlock()
	return ch
}

// UnsubscribeCancel отписывает воркера от cancel-сигналов.
func (ps *pieceScheduler) UnsubscribeCancel(peerID string) {
	ps.cancelMu.Lock()
	delete(ps.cancelSubs, peerID)
	ps.cancelMu.Unlock()
}

// broadcastCancel рассылает сигнал отмены куска всем воркерам кроме победителя.
func (ps *pieceScheduler) broadcastCancel(pieceIndex int, winnerPeerID string) {
	ps.cancelMu.RLock()
	defer ps.cancelMu.RUnlock()
	for id, ch := range ps.cancelSubs {
		if id == winnerPeerID {
			continue
		}
		select {
		case ch <- pieceIndex:
		default:
		}
	}
}

func (ps *pieceScheduler) Acquire(ctx context.Context, peerID string, have []bool, hasInfo bool) (pieceTask, bool, error) {
	responseCh := make(chan assignPieceResponse, 1)
	req := assignPieceRequest{
		peerID:     peerID,
		have:       have,
		hasInfo:    hasInfo,
		responseCh: responseCh,
	}

	select {
	case <-ctx.Done():
		return pieceTask{}, false, ctx.Err()
	case <-ps.doneCh:
		return pieceTask{}, false, nil
	case ps.assignCh <- req:
	}

	select {
	case <-ctx.Done():
		return pieceTask{}, false, ctx.Err()
	case <-ps.doneCh:
		return pieceTask{}, false, nil
	case response := <-responseCh:
		return response.task, response.ok, nil
	}
}

func (ps *pieceScheduler) Report(ctx context.Context, pieceIndex int, success bool) (bool, error) {
	responseCh := make(chan bool, 1)
	req := reportPieceRequest{
		pieceIndex: pieceIndex,
		success:    success,
		responseCh: responseCh,
	}

	select {
	case <-ctx.Done():
		return false, ctx.Err()
	case <-ps.doneCh:
		return false, nil
	case ps.reportCh <- req:
	}

	select {
	case <-ctx.Done():
		return false, ctx.Err()
	case <-ps.doneCh:
		return false, nil
	case accepted := <-responseCh:
		return accepted, nil
	}
}

func (ps *pieceScheduler) ReleasePeer(peerID string) {
	if peerID == "" {
		return
	}
	select {
	case <-ps.doneCh:
		return
	case ps.releasePeerCh <- peerID:
	default:
	}
}

func (ps *pieceScheduler) Progress() <-chan int {
	return ps.progressCh
}

func (ps *pieceScheduler) Done() <-chan struct{} {
	return ps.doneCh
}

func (ps *pieceScheduler) Stop() {
	ps.stopOnce.Do(func() {
		close(ps.stopCh)
	})
}

func (ps *pieceScheduler) SetSequential(mode bool) {
	select {
	case ps.setSequentialCh <- mode:
	case <-ps.stopCh:
	}
}

func (ps *pieceScheduler) run(pieceCount int, completedIndices []int) {
	states := make([]pieceState, pieceCount)
	availability := make([]int, pieceCount)
	peerAvailability := make(map[string][]bool, 256)
	endgamePeers := make(map[int][]string)
	completed := 0
	sequentialMode := false

	for _, idx := range completedIndices {
		if idx >= 0 && idx < pieceCount {
			states[idx] = pieceDone
			completed++
		}
	}

	finish := func() {
		close(ps.doneCh)
		close(ps.progressCh)
	}

	if pieceCount == 0 {
		finish()
		return
	}

	for {
		if completed >= pieceCount {
			finish()
			return
		}

		remaining := pieceCount - completed
		endgame := remaining <= endgameThreshold

		select {
		case <-ps.stopCh:
			finish()
			return
		case peerID := <-ps.releasePeerCh:
			releasePeerAvailability(peerAvailability, availability, peerID)
		case req := <-ps.assignCh:
			if req.peerID != "" && req.hasInfo {
				updatePeerAvailability(peerAvailability, availability, req.peerID, req.have, pieceCount)
			}

			pieceIndex := -1
			if endgame {
				pieceIndex = selectPieceEndgame(states, availability, req.have, req.hasInfo, endgamePeers, req.peerID)
			} else {
				if sequentialMode {
					pieceIndex = selectPendingPieceSequential(states, req.have, req.hasInfo)
				} else {
					pieceIndex = selectPendingPieceRarest(states, availability, req.have, req.hasInfo)
				}
			}

			if pieceIndex >= 0 {
				if !endgame {
					states[pieceIndex] = pieceInProgress
				}
				endgamePeers[pieceIndex] = append(endgamePeers[pieceIndex], req.peerID)
				req.responseCh <- assignPieceResponse{
					task: pieceTask{Index: pieceIndex},
					ok:   true,
				}
				continue
			}
			req.responseCh <- assignPieceResponse{ok: false}

		case req := <-ps.reportCh:
			if req.pieceIndex < 0 || req.pieceIndex >= pieceCount {
				req.responseCh <- false
				continue
			}

			if req.success {
				if states[req.pieceIndex] == pieceDone {
					req.responseCh <- false
					continue
				}
				states[req.pieceIndex] = pieceDone
				completed++

				// Endgame: уведомить других пиров отменить этот кусок
				if peers, ok := endgamePeers[req.pieceIndex]; ok && len(peers) > 1 {
					go ps.broadcastCancel(req.pieceIndex, "")
				}
				delete(endgamePeers, req.pieceIndex)

				select {
				case ps.progressCh <- req.pieceIndex:
				default:
				}
			} else {
				if states[req.pieceIndex] != pieceDone {
					states[req.pieceIndex] = piecePending
				}
			}
			req.responseCh <- true

		case mode := <-ps.setSequentialCh:
			sequentialMode = mode
		}
	}
}

// selectPieceEndgame выбирает кусок в endgame режиме:
// разрешает назначать куски в состоянии InProgress другим пирам.
func selectPieceEndgame(states []pieceState, availability []int, have []bool, hasInfo bool, endgamePeers map[int][]string, peerID string) int {
	// Сначала пробуем найти pending кусок (обычный путь)
	if idx := selectPendingPieceRarest(states, availability, have, hasInfo); idx >= 0 {
		return idx
	}

	// Endgame: ищем InProgress кусок, который у нас ещё не назначен этому пиру
	bestPiece := -1
	bestAvailability := 0

	for pieceIndex, state := range states {
		if state != pieceInProgress {
			continue
		}
		// Проверяем, не назначен ли уже этому пиру
		alreadyAssigned := false
		for _, pid := range endgamePeers[pieceIndex] {
			if pid == peerID {
				alreadyAssigned = true
				break
			}
		}
		if alreadyAssigned {
			continue
		}
		// Проверяем наличие у пира
		if hasInfo && (pieceIndex >= len(have) || !have[pieceIndex]) {
			continue
		}

		count := availability[pieceIndex]
		if bestPiece == -1 || count < bestAvailability {
			bestPiece = pieceIndex
			bestAvailability = count
		}
	}

	return bestPiece
}

func updatePeerAvailability(peerAvailability map[string][]bool, availability []int, peerID string, have []bool, pieceCount int) {
	normalized := make([]bool, pieceCount)
	copy(normalized, have)

	if prev, ok := peerAvailability[peerID]; ok {
		same := true
		for i := 0; i < pieceCount; i++ {
			if prev[i] != normalized[i] {
				same = false
				break
			}
		}
		if same {
			return
		}

		for i := 0; i < pieceCount; i++ {
			if prev[i] && availability[i] > 0 {
				availability[i]--
			}
			if normalized[i] {
				availability[i]++
			}
		}
		peerAvailability[peerID] = normalized
		return
	}

	for i := 0; i < pieceCount; i++ {
		if normalized[i] {
			availability[i]++
		}
	}
	peerAvailability[peerID] = normalized
}

func releasePeerAvailability(peerAvailability map[string][]bool, availability []int, peerID string) {
	prev, ok := peerAvailability[peerID]
	if !ok {
		return
	}

	for i := range prev {
		if prev[i] && availability[i] > 0 {
			availability[i]--
		}
	}
	delete(peerAvailability, peerID)
}

func selectPendingPieceRarest(states []pieceState, availability []int, have []bool, hasInfo bool) int {
	firstPending := firstPendingPiece(states)
	if firstPending < 0 {
		return -1
	}

	// If peer has not sent bitfield/have yet, optimistically probe.
	if !hasInfo {
		return firstPending
	}

	bestPiece := -1
	bestAvailability := 0

	for pieceIndex, state := range states {
		if state != piecePending {
			continue
		}
		if pieceIndex >= len(have) || !have[pieceIndex] {
			continue
		}

		count := availability[pieceIndex]
		if bestPiece == -1 || count < bestAvailability || (count == bestAvailability && pieceIndex < bestPiece) {
			bestPiece = pieceIndex
			bestAvailability = count
		}
	}

	if bestPiece >= 0 {
		return bestPiece
	}

	// Some peers send truncated availability info; allow fallback probing.
	if len(have) == 0 || len(have) < len(states) {
		return firstPending
	}

	return -1
}

func selectPendingPieceSequential(states []pieceState, have []bool, hasInfo bool) int {
	firstPending := firstPendingPiece(states)
	if firstPending < 0 {
		return -1
	}

	// If peer has not sent bitfield/have yet, optimistically probe.
	if !hasInfo {
		return firstPending
	}

	for pieceIndex, state := range states {
		if state != piecePending {
			continue
		}
		if pieceIndex >= len(have) || !have[pieceIndex] {
			continue
		}
		return pieceIndex // В последовательном режиме сразу берём первый доступный
	}

	// Some peers send truncated availability info; allow fallback probing.
	if len(have) == 0 || len(have) < len(states) {
		return firstPending
	}

	return -1
}

func firstPendingPiece(states []pieceState) int {
	for pieceIndex, state := range states {
		if state == piecePending {
			return pieceIndex
		}
	}
	return -1
}