package h2

import "sort"

// StreamPhase is where a stream is in its lifecycle (a simplified subset
// of RFC 9113 §5.1).
type StreamPhase int

const (
	StreamIdle StreamPhase = iota
	StreamOpen
	StreamHalfClosed // END_STREAM seen in one direction
	StreamClosed
)

func (p StreamPhase) String() string {
	return [...]string{"idle", "open", "half-closed", "closed"}[p]
}

// Stream is the per-stream state a demultiplexer tracks.
type Stream struct {
	ID         uint32
	Phase      StreamPhase
	HeaderObs  int // number of HEADERS frames seen
	DataBytes  int // total DATA payload bytes
	ClientInit bool
}

// Demux reconstructs per-stream state from an interleaved frame stream —
// the heart of HTTP/2 multiplexing. Many requests share one connection;
// the StreamID demultiplexes them.
type Demux struct {
	streams map[uint32]*Stream
	dec     *Decoder
	fc      *FlowController
}

func NewDemux() *Demux {
	return &Demux{
		streams: map[uint32]*Stream{},
		dec:     NewDecoder(4096),
		fc:      NewFlowController(),
	}
}

func (d *Demux) get(id uint32) *Stream {
	s, ok := d.streams[id]
	if !ok {
		s = &Stream{ID: id, Phase: StreamIdle, ClientInit: id%2 == 1}
		d.streams[id] = s
	}
	return s
}

// Process feeds one frame into the demultiplexer, updating stream state
// and (for HEADERS) decoding the header block. Connection-level frames
// (stream 0) update flow control.
func (d *Demux) Process(f *Frame) ([]HeaderField, error) {
	switch f.Header.Type {
	case FrameHeaders:
		s := d.get(f.Header.StreamID)
		if s.Phase == StreamIdle {
			s.Phase = StreamOpen
		}
		s.HeaderObs++
		hdrs, err := d.dec.Decode(f.Payload) // may be HEADERS (request) or trailers
		if f.EndStream() {
			s.Phase = StreamHalfClosed
		}
		return hdrs, err

	case FrameData:
		s := d.get(f.Header.StreamID)
		s.DataBytes += len(f.Payload)
		// Consuming DATA means we owe the peer a WINDOW_UPDATE; model the
		// debit so a test can watch the window drain/recover.
		_ = d.fc.Sent(f.Header.StreamID, int64(len(f.Payload)))
		if f.EndStream() {
			if s.Phase == StreamHalfClosed {
				s.Phase = StreamClosed
			} else {
				s.Phase = StreamHalfClosed
			}
		}
		return nil, nil

	case FrameWindowUpdate:
		inc := int64(beUint32(f.Payload) & 0x7fffffff)
		d.fc.WindowUpdate(f.Header.StreamID, inc)
		return nil, nil

	case FrameRSTStream:
		s := d.get(f.Header.StreamID)
		s.Phase = StreamClosed
		return nil, nil
	}
	return nil, nil
}

// ActiveStreams returns the concurrently-tracked stream IDs in ascending
// order — i.e. the requests multiplexed over this one connection.
func (d *Demux) ActiveStreams() []uint32 {
	ids := make([]uint32, 0, len(d.streams))
	for id := range d.streams {
		if id != 0 {
			ids = append(ids, id)
		}
	}
	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
	return ids
}

func (d *Demux) Stream(id uint32) *Stream { return d.streams[id] }
func (d *Demux) Flow() *FlowController     { return d.fc }

func beUint32(b []byte) uint32 {
	if len(b) < 4 {
		return 0
	}
	return uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
}