// state machine binding one pion PeerConnection to the XMPP signalling
// channel (XEP-0166 + XEP-0353 JMI preamble)
package jingle

import (
	"context"
	"encoding/xml"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/pion/webrtc/v4"
	"gosrc.io/xmpp/stanza"
)

type Role uint8

const (
	RoleInitiator Role = iota // local side proposed the call
	RoleResponder             // peer proposed the call
)

// high-level call lifecycle; pion's ConnectionState owns media substates
type State uint8

const (
	StateNew       State = iota
	StateProposed        // caller: sent <propose>, awaiting proceed/reject/retract
	StateRinging         // callee: got <propose>, sent <ringing>, awaiting user accept/decline
	StateProceeded       // proceed exchanged; media setup either pending or in flight
	StateActive          // session-accept exchanged; pion drives ICE/DTLS from here
	StateTerminated
)

func (s State) String() string {
	switch s {
	case StateNew:
		return "new"
	case StateProposed:
		return "proposed"
	case StateRinging:
		return "ringing"
	case StateProceeded:
		return "proceeded"
	case StateActive:
		return "active"
	case StateTerminated:
		return "terminated"
	}
	return fmt.Sprintf("state(%d)", uint8(s))
}

// slice of gosrc.io/xmpp.Sender that Session needs; *xmpp.Client and
// *xmpp.Component both satisfy it
type Sender interface {
	Send(packet stanza.Packet) error
	SendIQ(ctx context.Context, iq *stanza.IQ) (chan stanza.IQ, error)
}

// JMI/Jingle lifecycle callbacks. Methods may fire under the session mutex,
// so observers must not call back into the session synchronously - defer
// to a goroutine. When set, Session does NOT auto-send <ringing/> on
// incoming <propose>; the observer drives that.
type SessionObserver interface {
	OnRinging(from string)
	OnProceeded(from string)
	// fires after SetRemoteDescription, before CreateAnswer on the responder,
	// so a same-stack ReplaceTrack lands in the answer SDP
	OnRemoteDescriptionApplied(sdp string)
	// reason: XEP-0166 condition local-name, or "rejected"/"retracted"
	OnTerminated(reason string)
}

// PC and Sender are required
type SessionOpts struct {
	PC        *webrtc.PeerConnection
	Sender    Sender
	LocalJID  string
	RemoteJID string // bare is fine for initial <propose>; latched to full on <proceed>
	SID       string
	Role      Role
	Media     []string // only audio supported
	Observer  SessionObserver
}

// one call. Public methods are safe from any goroutine; mu is held only
// briefly and never across pion calls or Sender I/O.
type Session struct {
	opts SessionOpts
	// latched from opts so it's readable without mu
	observer SessionObserver

	mu                sync.Mutex
	state             State
	remoteCandQueue   []remoteCandidate     // candidates from before SetRemoteDescription
	remoteDescApplied bool
	localCandQueue    []webrtc.ICECandidate // gathered before session-initiate/accept
	signaled          bool
	// captured from the local SDP after SetLocalDescription
	ufrag  string
	pwd    string
	mids   []string
	closed bool
}

type remoteCandidate struct {
	mid  string
	init webrtc.ICECandidateInit
}

// human-readable so XML dumps are easy to grep
var stanzaIDCounter atomic.Uint64

func newStanzaID() string {
	return fmt.Sprintf("jng-%d", stanzaIDCounter.Add(1))
}

// caller adds transceivers/tracks to PC before Propose/AcceptProposal
func New(opts SessionOpts) (*Session, error) {
	if opts.PC == nil {
		return nil, errors.New("jingle: PC required")
	}
	if opts.Sender == nil {
		return nil, errors.New("jingle: Sender required")
	}
	if opts.SID == "" {
		return nil, errors.New("jingle: SID required")
	}
	if opts.LocalJID == "" || opts.RemoteJID == "" {
		return nil, errors.New("jingle: LocalJID and RemoteJID required")
	}
	if len(opts.Media) == 0 {
		opts.Media = []string{"audio"}
	}
	s := &Session{opts: opts, observer: opts.Observer}
	opts.PC.OnICECandidate(s.onICECandidate)
	return s, nil
}

func (s *Session) State() State {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.state
}

func (s *Session) SID() string { return s.opts.SID }

// after <proceed> arrives this is the full JID of the device that picked up
func (s *Session) RemoteJID() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.opts.RemoteJID
}

// <message><propose/></message>
func (s *Session) Propose(ctx context.Context) error {
	s.mu.Lock()
	if s.state != StateNew || s.opts.Role != RoleInitiator {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("propose: bad state %s", st)
	}
	to := s.opts.RemoteJID
	from := s.opts.LocalJID
	media := append([]string(nil), s.opts.Media...)
	s.mu.Unlock()

	msg := newJMIMessage(from, to, &JMIPropose{
		ID:           s.opts.SID,
		Descriptions: mediaDescs(media),
	})
	if err := s.sendStanza(msg); err != nil {
		return fmt.Errorf("propose: send: %w", err)
	}
	s.setStateIf(StateNew, StateProposed)
	return nil
}

// <message><ringing/></message> - callee tells caller the device is alerting
func (s *Session) Ringing(ctx context.Context) error {
	s.mu.Lock()
	if s.state != StateRinging || s.opts.Role != RoleResponder {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("ringing: bad state %s", st)
	}
	to, from := s.opts.RemoteJID, s.opts.LocalJID
	s.mu.Unlock()

	msg := newJMIMessage(from, to, &JMIRinging{ID: s.opts.SID})
	return s.sendStanza(msg)
}

// <message><proceed/></message>; media setup waits for session-initiate
func (s *Session) AcceptProposal(ctx context.Context) error {
	s.mu.Lock()
	if s.state != StateRinging || s.opts.Role != RoleResponder {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("accept: bad state %s", st)
	}
	to, from := s.opts.RemoteJID, s.opts.LocalJID
	s.mu.Unlock()

	msg := newJMIMessage(from, to, &JMIProceed{ID: s.opts.SID})
	if err := s.sendStanza(msg); err != nil {
		return fmt.Errorf("accept: send: %w", err)
	}
	s.setStateIf(StateRinging, StateProceeded)
	return nil
}

// <message><reject/></message> (callee, pre-media)
func (s *Session) Decline(ctx context.Context, reason string) error {
	s.mu.Lock()
	if s.state != StateRinging || s.opts.Role != RoleResponder {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("decline: bad state %s", st)
	}
	to, from := s.opts.RemoteJID, s.opts.LocalJID
	s.mu.Unlock()

	r := &JMIReject{ID: s.opts.SID}
	if reason != "" {
		r.Reason = &Reason{Condition: &ReasonCondition{XMLName: xml.Name{Space: NSJingle, Local: reason}}}
	}
	msg := newJMIMessage(from, to, r)
	err := s.sendStanza(msg)
	s.terminate()
	return err
}

// <message><retract/></message> (caller, pre-<proceed>)
func (s *Session) Retract(ctx context.Context) error {
	s.mu.Lock()
	if s.state != StateProposed || s.opts.Role != RoleInitiator {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("retract: bad state %s", st)
	}
	to, from := s.opts.RemoteJID, s.opts.LocalJID
	s.mu.Unlock()

	msg := newJMIMessage(from, to, &JMIRetract{ID: s.opts.SID})
	err := s.sendStanza(msg)
	s.terminate()
	return err
}

// <jingle action=session-terminate>; falls through to Retract/Decline if
// invoked pre-media
func (s *Session) Terminate(ctx context.Context, reason string) error {
	s.mu.Lock()
	switch s.state {
	case StateTerminated:
		s.mu.Unlock()
		return nil
	case StateProposed, StateRinging:
		role := s.opts.Role
		s.mu.Unlock()
		if role == RoleInitiator {
			return s.Retract(ctx)
		}
		return s.Decline(ctx, reason)
	}
	to, from := s.opts.RemoteJID, s.opts.LocalJID
	role := s.opts.Role
	s.mu.Unlock()

	j := &JingleIQ{
		Action: ActionSessionTerminate,
		SID:    s.opts.SID,
	}
	if role == RoleInitiator {
		j.Initiator = from
	} else {
		j.Responder = from
	}
	if reason != "" {
		j.Reason = &Reason{Condition: &ReasonCondition{XMLName: xml.Name{Space: NSJingle, Local: reason}}}
	}
	iq, err := newJingleIQ(from, to, stanza.IQTypeSet, j)
	if err != nil {
		s.terminate()
		return err
	}
	sendErr := s.sendStanza(iq)
	s.terminate()
	return sendErr
}

// dispatches an incoming JMI payload; Manager has matched the SID
func (s *Session) HandleJMI(from string, ext stanza.MsgExtension) {
	switch m := ext.(type) {
	case *JMIPropose:
		s.handleJMIPropose(from, *m)
	case *JMIProceed:
		s.handleJMIProceed(from)
	case *JMIRinging:
		s.handleJMIRinging(from)
	case *JMIReject:
		s.handleJMIReject(from)
	case *JMIRetract:
		s.handleJMIRetract(from)
	case *JMIFinish:
		// terminal informational marker; ignored
	}
}

func (s *Session) handleJMIPropose(from string, m JMIPropose) {
	s.mu.Lock()
	if s.state != StateNew || s.opts.Role != RoleResponder {
		s.mu.Unlock()
		return
	}
	s.opts.RemoteJID = from
	s.state = StateRinging
	s.mu.Unlock()
	// auto-<ringing/> when no observer is wired
	if s.observer == nil {
		_ = s.Ringing(context.Background())
	}
}

func (s *Session) handleJMIProceed(from string) {
	s.mu.Lock()
	if s.state != StateProposed || s.opts.Role != RoleInitiator {
		s.mu.Unlock()
		return
	}
	s.opts.RemoteJID = from // latch full JID
	s.state = StateProceeded
	s.mu.Unlock()
	if s.observer != nil {
		s.observer.OnProceeded(from)
	}
	// off the dispatcher so pion's gather doesn't block JMI routing
	go s.startOutgoingMedia()
}

func (s *Session) handleJMIRinging(from string) {
	if s.observer != nil {
		s.observer.OnRinging(from)
	}
}

func (s *Session) handleJMIReject(from string) {
	s.mu.Lock()
	if s.state == StateTerminated {
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	// terminate before observer; observers re-entering via the bridge then
	// see StateTerminated and short-circuit their own teardown
	s.terminate()
	if s.observer != nil {
		s.observer.OnTerminated("rejected")
	}
}

func (s *Session) handleJMIRetract(from string) {
	// caller cancelled before responder proceeded
	s.mu.Lock()
	if s.state == StateTerminated {
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	s.terminate()
	if s.observer != nil {
		s.observer.OnTerminated("retracted")
	}
}

// processes an inbound Jingle IQ; Manager has matched the SID. Returns the
// result/error condition for the IQ reply.
func (s *Session) HandleJingleIQ(iq *stanza.IQ, j *JingleIQ) (resultPayload stanza.IQPayload, errCondition string) {
	switch j.Action {
	case ActionSessionInitiate:
		if err := s.handleSessionInitiate(iq, j); err != nil {
			return nil, "bad-request"
		}
	case ActionSessionAccept:
		if err := s.handleSessionAccept(iq, j); err != nil {
			return nil, "bad-request"
		}
	case ActionSessionTerminate:
		s.handleSessionTerminate(iq, j)
	case ActionTransportInfo:
		s.handleTransportInfo(iq, j)
	default:
		// unsupported action: ack but no-op
	}
	return nil, ""
}

func (s *Session) handleSessionInitiate(iq *stanza.IQ, j *JingleIQ) error {
	s.mu.Lock()
	if s.state != StateProceeded || s.opts.Role != RoleResponder {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("out-of-order in state %s", st)
	}
	s.mu.Unlock()

	sdp, err := JingleToSDP(j, ConvertOpts{Initiator: false, LocalCreator: "initiator", SessionID: s.opts.SID})
	if err != nil {
		return fmt.Errorf("convert offer: %w", err)
	}
	if err := s.opts.PC.SetRemoteDescription(webrtc.SessionDescription{
		Type: webrtc.SDPTypeOffer,
		SDP:  sdp,
	}); err != nil {
		return fmt.Errorf("SetRemoteDescription(offer): %w", err)
	}
	s.drainRemoteCandidates()
	s.notifyRemoteDescriptionApplied(sdp)

	answer, err := s.opts.PC.CreateAnswer(nil)
	if err != nil {
		return fmt.Errorf("CreateAnswer: %w", err)
	}
	if err := s.opts.PC.SetLocalDescription(answer); err != nil {
		return fmt.Errorf("SetLocalDescription(answer): %w", err)
	}
	s.captureLocalMeta()

	ja, err := SDPToJingle(answer.SDP, ConvertOpts{Initiator: false, LocalCreator: "initiator", SessionID: s.opts.SID})
	if err != nil {
		return fmt.Errorf("convert answer: %w", err)
	}
	ja.Action = ActionSessionAccept
	ja.Responder = s.opts.LocalJID

	out, err := newJingleIQ(s.opts.LocalJID, s.opts.RemoteJID, stanza.IQTypeSet, ja)
	if err != nil {
		return err
	}
	if err := s.sendStanza(out); err != nil {
		return fmt.Errorf("send session-accept: %w", err)
	}
	s.mu.Lock()
	s.state = StateActive
	s.mu.Unlock()
	s.drainLocalCandidates()
	return nil
}

func (s *Session) handleSessionAccept(iq *stanza.IQ, j *JingleIQ) error {
	s.mu.Lock()
	if s.state != StateProceeded || s.opts.Role != RoleInitiator {
		st := s.state
		s.mu.Unlock()
		return fmt.Errorf("out-of-order in state %s", st)
	}
	s.mu.Unlock()

	sdp, err := JingleToSDP(j, ConvertOpts{Initiator: true, LocalCreator: "initiator", SessionID: s.opts.SID})
	if err != nil {
		return fmt.Errorf("convert answer: %w", err)
	}
	if err := s.opts.PC.SetRemoteDescription(webrtc.SessionDescription{
		Type: webrtc.SDPTypeAnswer,
		SDP:  sdp,
	}); err != nil {
		return fmt.Errorf("SetRemoteDescription(answer): %w", err)
	}
	s.drainRemoteCandidates()
	s.notifyRemoteDescriptionApplied(sdp)
	s.mu.Lock()
	s.state = StateActive
	s.mu.Unlock()
	return nil
}

// dispatched inline so an answerer-side ReplaceTrack from the observer
// lands before CreateAnswer runs
func (s *Session) notifyRemoteDescriptionApplied(sdp string) {
	if s.observer != nil {
		s.observer.OnRemoteDescriptionApplied(sdp)
	}
}

func (s *Session) handleTransportInfo(iq *stanza.IQ, j *JingleIQ) {
	for _, c := range j.Contents {
		if c.Transport == nil {
			continue
		}
		for _, cand := range c.Transport.Candidates {
			candCopy := cand
			init, err := candidateToPionInit(&candCopy, c.Name)
			if err != nil {
				continue
			}
			s.applyRemoteCandidate(c.Name, init)
		}
	}
}

func (s *Session) handleSessionTerminate(iq *stanza.IQ, j *JingleIQ) {
	reason := ""
	if j != nil && j.Reason != nil && j.Reason.Condition != nil {
		reason = j.Reason.Condition.XMLName.Local
	}
	// terminate before observer; observers re-entering via the bridge then
	// see StateTerminated and short-circuit their own teardown
	s.terminate()
	if s.observer != nil {
		s.observer.OnTerminated(reason)
	}
}

func (s *Session) startOutgoingMedia() {
	offer, err := s.opts.PC.CreateOffer(nil)
	if err != nil {
		s.terminate()
		return
	}
	if err := s.opts.PC.SetLocalDescription(offer); err != nil {
		s.terminate()
		return
	}
	s.captureLocalMeta()

	j, err := SDPToJingle(offer.SDP, ConvertOpts{Initiator: true, LocalCreator: "initiator", SessionID: s.opts.SID})
	if err != nil {
		s.terminate()
		return
	}
	j.Action = ActionSessionInitiate
	j.Initiator = s.opts.LocalJID

	iq, err := newJingleIQ(s.opts.LocalJID, s.opts.RemoteJID, stanza.IQTypeSet, j)
	if err != nil {
		s.terminate()
		return
	}
	if err := s.sendStanza(iq); err != nil {
		s.terminate()
		return
	}
	s.drainLocalCandidates()
}

// ICE candidates are buffered until SetLocalDescription/SetRemoteDescription
// has run; otherwise pion drops them

func (s *Session) onICECandidate(c *webrtc.ICECandidate) {
	if c == nil {
		return // end-of-gathering; no wire marker
	}
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	if !s.signaled {
		s.localCandQueue = append(s.localCandQueue, *c)
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	_ = s.sendTransportInfo(c)
}

func (s *Session) drainLocalCandidates() {
	s.mu.Lock()
	s.signaled = true
	queue := s.localCandQueue
	s.localCandQueue = nil
	s.mu.Unlock()
	for i := range queue {
		_ = s.sendTransportInfo(&queue[i])
	}
}

func (s *Session) applyRemoteCandidate(mid string, init webrtc.ICECandidateInit) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	if !s.remoteDescApplied {
		s.remoteCandQueue = append(s.remoteCandQueue, remoteCandidate{mid: mid, init: init})
		s.mu.Unlock()
		return
	}
	s.mu.Unlock()
	_ = s.opts.PC.AddICECandidate(init)
}

func (s *Session) drainRemoteCandidates() {
	s.mu.Lock()
	s.remoteDescApplied = true
	queue := s.remoteCandQueue
	s.remoteCandQueue = nil
	s.mu.Unlock()
	for _, rc := range queue {
		_ = s.opts.PC.AddICECandidate(rc.init)
	}
}

func (s *Session) sendTransportInfo(c *webrtc.ICECandidate) error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil
	}
	ufrag, pwd := s.ufrag, s.pwd
	mids := s.mids
	role := s.opts.Role
	localJID, remoteJID := s.opts.LocalJID, s.opts.RemoteJID
	s.mu.Unlock()

	cand, err := pionToCandidate(c)
	if err != nil {
		return err
	}
	mid := c.SDPMid
	if mid == "" && len(mids) > 0 {
		mid = mids[0]
	}
	j := buildTransportInfo(s.opts.SID, mid, ufrag, pwd, "initiator", cand)
	if role == RoleInitiator {
		j.Initiator = localJID
	} else {
		j.Responder = localJID
	}
	iq, err := newJingleIQ(localJID, remoteJID, stanza.IQTypeSet, j)
	if err != nil {
		return err
	}
	return s.sendStanza(iq)
}

func (s *Session) captureLocalMeta() {
	local := s.opts.PC.LocalDescription()
	if local == nil {
		return
	}
	ufrag, pwd, mids := extractSDPMeta(local.SDP)
	s.mu.Lock()
	s.ufrag, s.pwd, s.mids = ufrag, pwd, mids
	s.mu.Unlock()
}

// with BUNDLE, ufrag/pwd are identical across m= blocks; take the first
func extractSDPMeta(s string) (ufrag, pwd string, mids []string) {
	doc, err := parseSDP(s)
	if err != nil {
		return
	}
	for _, m := range doc.media {
		if ufrag == "" {
			ufrag = mmFirst(m.attrs, "ice-ufrag", "")
		}
		if pwd == "" {
			pwd = mmFirst(m.attrs, "ice-pwd", "")
		}
		if mid := mmFirst(m.attrs, "mid", ""); mid != "" {
			mids = append(mids, mid)
		}
	}
	return
}

func (s *Session) setStateIf(want, next State) {
	s.mu.Lock()
	if s.state == want {
		s.state = next
	}
	s.mu.Unlock()
}

func (s *Session) terminate() {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	s.state = StateTerminated
	s.remoteCandQueue = nil
	s.localCandQueue = nil
	pc := s.opts.PC
	s.mu.Unlock()
	if pc != nil {
		_ = pc.Close()
	}
}

// release pion + buffers without sending anything on the wire
func (s *Session) Close() error {
	s.terminate()
	return nil
}

func mediaDescs(media []string) []JMIDescription {
	out := make([]JMIDescription, 0, len(media))
	for _, m := range media {
		out = append(out, JMIDescription{Media: m})
	}
	return out
}

func newJMIMessage(from, to string, ext stanza.MsgExtension) *stanza.Message {
	msg := stanza.NewMessage(stanza.Attrs{
		Type: stanza.MessageTypeChat,
		From: from,
		To:   to,
		Id:   newStanzaID(),
	})
	msg.Extensions = []stanza.MsgExtension{ext}
	return &msg
}

func (s *Session) sendStanza(p stanza.Packet) error {
	return s.opts.Sender.Send(p)
}

func newJingleIQ(from, to string, t stanza.StanzaType, j *JingleIQ) (*stanza.IQ, error) {
	iq, err := stanza.NewIQ(stanza.Attrs{
		Type: t,
		From: from,
		To:   to,
		Id:   newStanzaID(),
	})
	if err != nil {
		return nil, err
	}
	iq.Payload = j
	return iq, nil
}