package host

import (
	"context"
	"errors"
	"sort"
	"sync"

	"github.com/carlosframework/platform/internal/hostbox"
	"github.com/carlosframework/platform/internal/hostguard"
)

type claimProducer struct {
	cancel context.CancelFunc
	joined chan struct{}
	err    error
}
type completeClaimProducer func(context.Context, hostguard.FenceTarget, string, hostbox.ProducerCompletion) (hostbox.ProducerCompletion, error)

type ClaimProducers struct {
	closeWriters func(context.Context, hostguard.FenceTarget) error
	completed    hostbox.ProducerCompletion
	complete     completeClaimProducer
	mu           sync.Mutex
	target       hostguard.FenceTarget
	processStart string
	fence        string
	drained      bool
	// writersStopped closes admission to every producer but the retained
	// ones (StopWriters); the full drain still cancels those.
	writersStopped bool
	workers        map[string]*claimProducer
	Record         func(context.Context, hostguard.FenceTarget, string) error
}

func NewClaimProducers(target hostguard.FenceTarget, start string, record func(context.Context, hostguard.FenceTarget, string) error, complete ...completeClaimProducer) (*ClaimProducers, error) {
	if target.Host == "" || target.ControllerID == "" || start == "" || record == nil || len(complete) > 1 || (len(complete) == 1 && complete[0] == nil) {
		return nil, ErrClaimStopUnavailable
	}
	p := &ClaimProducers{target: target, processStart: start, Record: record, workers: map[string]*claimProducer{}}
	if len(complete) == 1 {
		p.complete = complete[0]
	}
	return p, nil
}
func (p *ClaimProducers) Launch(ctx context.Context, target hostguard.FenceTarget, id string, fn func(context.Context) error) error {
	if p == nil || fn == nil || !stopComponent(id) {
		return ErrClaimStopUnavailable
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) {
		return ErrClaimStopIdentity
	}
	if p.writersStopped && !claimRetainedProducer(id) {
		return ErrClaimWritersStopped
	}
	if err := p.compact(ctx); err != nil {
		return err
	}
	if _, ok := p.workers[id]; ok {
		return ErrClaimStopIdentity
	}
	if err := p.Record(ctx, target, id); err != nil {
		return err
	}
	child, cancel := context.WithCancel(ctx)
	w := &claimProducer{cancel: cancel, joined: make(chan struct{})}
	p.workers[id] = w
	go func() { p.finish(w, fn(child)) }()
	return nil
}
func (p *ClaimProducers) DrainClaim(ctx context.Context, r StopRequest, fence string) (DrainAcknowledgement, error) {
	if p == nil || fence == "" {
		return DrainAcknowledgement{}, ErrClaimStopUnavailable
	}
	p.mu.Lock()
	if p.target.Host != r.Host || p.target.Incarnation != r.Incarnation || p.target.Stamp != r.InstallationStamp || p.target.Session != r.WriterSession || p.target.Box != r.Box || p.target.Enrollment != r.Enrollment || p.target.BootID != r.BootID || p.target.ControllerID != r.ControllerID || (p.fence != "" && p.fence != fence && !p.drained) {
		p.mu.Unlock()
		return DrainAcknowledgement{}, ErrClaimStopIdentity
	}
	p.fence, p.drained = fence, false
	if p.closeWriters != nil {
		if err := p.closeWriters(ctx, p.target); err != nil {
			p.mu.Unlock()
			return DrainAcknowledgement{}, err
		}
	}
	workers := map[string]*claimProducer{}
	for id, w := range p.workers {
		workers[id] = w
		w.cancel()
	}
	p.mu.Unlock()
	for _, w := range workers {
		select {
		case <-ctx.Done():
			return DrainAcknowledgement{}, ctx.Err()
		case <-w.joined:
		}
		if w.err != nil && !errors.Is(w.err, context.Canceled) {
			return DrainAcknowledgement{}, w.err
		}
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := p.compact(ctx); err != nil {
		return DrainAcknowledgement{}, err
	}
	ids := make([]string, 0, len(p.workers))
	for id := range p.workers {
		ids = append(ids, id)
	}
	p.drained = true
	return DrainClaimAcknowledgement(fence, r.ControllerID, p.processStart, ids, p.completed), nil
}

func (p *ClaimProducers) Register(ctx context.Context, target hostguard.FenceTarget, id string) (context.Context, func(error), error) {
	if p == nil || !stopComponent(id) {
		return nil, nil, ErrClaimStopUnavailable
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) {
		return nil, nil, ErrClaimStopIdentity
	}
	if p.writersStopped && !claimRetainedProducer(id) {
		return nil, nil, ErrClaimWritersStopped
	}
	if err := p.compact(ctx); err != nil {
		return nil, nil, err
	}
	if _, ok := p.workers[id]; ok {
		return nil, nil, ErrClaimStopIdentity
	}
	if err := p.Record(ctx, target, id); err != nil {
		return nil, nil, err
	}
	child, cancel := context.WithCancel(ctx)
	w := &claimProducer{cancel: cancel, joined: make(chan struct{})}
	p.workers[id] = w
	var once sync.Once
	finish := func(err error) { once.Do(func() { p.finish(w, err) }) }
	return child, finish, nil
}

// StopWriters is the writers-only drain (claim Plan 11, R2): it closes
// admission to every producer but the retained ones, cancels the rest (the
// role launch's producers and the schedule runner's timer deliveries) and
// joins them. The replicator, the generation publisher and the files
// publisher keep running; DrainClaim, root's full drain, still ends them.
// It returns the ids it joined, sorted. Repeating it is safe.
func (p *ClaimProducers) StopWriters(ctx context.Context, target hostguard.FenceTarget) ([]string, error) {
	if p == nil || ctx == nil {
		return nil, ErrClaimStopUnavailable
	}
	p.mu.Lock()
	if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) {
		p.mu.Unlock()
		return nil, ErrClaimStopIdentity
	}
	p.writersStopped = true
	if p.closeWriters != nil {
		if err := p.closeWriters(ctx, p.target); err != nil {
			p.mu.Unlock()
			return nil, err
		}
	}
	workers := map[string]*claimProducer{}
	for id, w := range p.workers {
		if !claimRetainedProducer(id) {
			workers[id] = w
			w.cancel()
		}
	}
	p.mu.Unlock()
	ids := make([]string, 0, len(workers))
	for id, w := range workers {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-w.joined:
		}
		if w.err != nil && !errors.Is(w.err, context.Canceled) {
			return nil, w.err
		}
		ids = append(ids, id)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := p.compact(ctx); err != nil {
		return nil, err
	}
	sort.Strings(ids)
	return ids, nil
}

// Running reports the producers that have not finished, sorted.
func (p *ClaimProducers) Running() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	ids := make([]string, 0, len(p.workers))
	for id, w := range p.workers {
		select {
		case <-w.joined:
		default:
			ids = append(ids, id)
		}
	}
	sort.Strings(ids)
	return ids
}

type ClaimProducerSet struct {
	CloseWriters func(context.Context, hostguard.FenceTarget) error
	mu           sync.Mutex
	items        []*ClaimProducers
}

func (s *ClaimProducerSet) Add(p *ClaimProducers) error {
	if p == nil {
		return ErrClaimStopUnavailable
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, old := range s.items {
		if hostbox.SameLaunchOwner(old.target, p.target) {
			return ErrClaimStopIdentity
		}
	}
	s.items = append(s.items, p)
	return nil
}
func (s *ClaimProducerSet) For(t hostguard.FenceTarget) (*ClaimProducers, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, p := range s.items {
		if hostbox.SameLaunchOwner(p.target, t) {
			return p, nil
		}
	}
	return nil, ErrClaimStopUnavailable
}
func (s *ClaimProducerSet) DrainClaim(ctx context.Context, r StopRequest, fence string) (DrainAcknowledgement, error) {
	t := hostguard.FenceTarget{Host: r.Host, Incarnation: r.Incarnation, Stamp: r.InstallationStamp, Session: r.WriterSession, Box: r.Box, Enrollment: r.Enrollment, BootID: r.BootID, ControllerID: r.ControllerID}
	p, err := s.For(t)
	if err != nil {
		return DrainAcknowledgement{}, err
	}
	return p.DrainClaim(ctx, r, fence)
}

func (s *ClaimProducerSet) Ensure(target hostguard.FenceTarget, start string, record func(context.Context, hostguard.FenceTarget, string) error, complete ...completeClaimProducer) (*ClaimProducers, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, old := range s.items {
		if hostbox.SameLaunchOwner(old.target, target) {
			if old.processStart != start {
				return nil, ErrClaimStopIdentity
			}
			return old, nil
		}
	}
	p, err := NewClaimProducers(target, start, record, complete...)
	if err != nil {
		return nil, err
	}
	p.closeWriters = s.CloseWriters
	s.items = append(s.items, p)
	return p, nil
}

func (p *ClaimProducers) finish(w *claimProducer, err error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	w.err = err
	w.cancel()
	close(w.joined)
	ctx, cancel := context.WithTimeout(context.Background(), hostbox.RequestLimit)
	defer cancel()
	// A failed durable write retains the joined worker. Admission and drain
	// retry it before accepting more work or acknowledging completion.
	_ = p.compact(ctx)
}

// compact runs under mu so the journal and the memory summary advance together.
func (p *ClaimProducers) compact(ctx context.Context) error {
	if p.complete == nil {
		return nil
	}
	for id, w := range p.workers {
		select {
		case <-w.joined:
		default:
			continue
		}
		if w.err != nil && !errors.Is(w.err, context.Canceled) {
			return w.err
		}
		next, err := p.complete(ctx, p.target, id, p.completed)
		if err != nil {
			return err
		}
		if !next.Valid() || next.Count != p.completed.Count+1 {
			return ErrClaimStopIdentity
		}
		p.completed = next
		delete(p.workers, id)
	}
	return nil
}