// Package kobject interprets uevent messages from a NETLINK_KOBJECT_UEVENT socket.
package kobject

import (
	"context"
	"fmt"
	"maps"
	"slices"
	"strconv"
	"sync"

	"hakurei.app/internal/uevent"
)

const (
	// StateColdboot denotes an [Object] populated by a coldboot event. It is
	// eligible for all event actions.
	StateColdboot = iota
	// StateNew denotes an [Object] previously populated by a [uevent.KOBJ_ADD]
	// event, but has not yet been targeted by a [uevent.KOBJ_BIND] event, or
	// has been targeted by a [uevent.KOBJ_UNBIND] event.
	StateNew
	// StateBound denotes an [Object] that has been targeted by a
	// [uevent.KOBJ_BIND] and has not been targeted by a [uevent.KOBJ_UNBIND]
	// after that.
	StateBound
)

// Object represents a kernel object.
type Object struct {
	// Origin of the object.
	State int `json:"-"`
	// Set by [uevent.KOBJ_OFFLINE] and [uevent.KOBJ_ONLINE].
	Offline bool

	// alloc_uevent_skb: devpath
	DevPath string `json:"devpath"`
	// registered per-driver (optional)
	ModAlias string `json:"modalias,omitempty"`
	// dev_driver_uevent: drv->name (optional)
	Driver string `json:"driver,omitempty"`

	// SUBSYSTEM value set by the kernel.
	Subsystem string `json:"subsystem"`

	// Uninterpreted environment variable pairs. An entry missing a separator
	// gains the value "\x00".
	Env map[string]string `json:"env"`
}

// Clone returns a copy of o.
func (o *Object) Clone() Object {
	v := *o
	v.Env = maps.Clone(o.Env)
	return v
}

// GoString returns compound literal for the underlying value.
func (o *Object) GoString() string {
	return fmt.Sprintf("&%#v", *o)
}

// merge merges uninterpreted environment variable pairs from an [Event].
func (o *Object) merge(env map[string]string) {
	for k, v := range env {
		if v == "\x00" {
			continue
		}

		switch k {
		case "MODALIAS":
			o.ModAlias = v
			continue

		case "DRIVER":
			o.Driver = v
			continue

		default:
			if o.Env == nil {
				o.Env = make(map[string]string)
			}
			o.Env[k] = v
		}
	}
}

// update updates o with pairs from env, optionally stripping visited pairs.
func (o *Object) update(env map[string]string, strip bool) {
	for k := range o.Env {
		if v, ok := env[k]; ok {
			if strip {
				delete(env, k)
			}
			o.Env[k] = v
		}
	}
}

// A pendingIterator is a callback currently iterating through objects targeted
// by ongoing events.
type pendingIterator struct {
	f    func(o *Object) bool
	done chan<- struct{}
}

// State processes a stream of [Event] populated from [uevent.Message] received
// from a NETLINK_KOBJECT_UEVENT socket and presents an efficient representation
// of kernel state.
type State struct {
	// Next expected SEQNUM.
	seq uint64
	// DevPath to environment variables.
	uevent map[string]*Object
	// Synchronises access to uevent and its objects.
	ueventMu sync.RWMutex
	// Alive iterators.
	iter []*pendingIterator
	// Synchronises access to iter.
	iterMu sync.Mutex
	// UUID for synthetic [uevent.Coldboot] events.
	coldboot uevent.UUID
	// Called on [uevent.KOBJ_CHANGE] with stripped environment variables.
	handleChange func(o *Object, env map[string]string)
	// Reports errors populating [Event] from [uevent.Message]. A user-supplied
	// nil value is replaced with a noop.
	reportErr func(error)
}

// New returns the address of a new [State].
func New(
	coldboot uevent.UUID,
	handleChange func(o *Object, env map[string]string),
	reportErr func(error),
) *State {
	return &State{
		uevent:       make(map[string]*Object),
		coldboot:     coldboot,
		handleChange: handleChange,
		reportErr:    reportErr,
	}
}

// deleteIter removes an iterator from s. Must be called after acquiring iterMu.
func (s *State) deleteIter(p *pendingIterator) {
	s.iter = slices.DeleteFunc(s.iter, func(v *pendingIterator) bool {
		return p == v
	})
}

// dispatchIter broadcasts an [Object] to all alive iterators.
func (s *State) dispatchIter(o *Object) {
	s.iterMu.Lock()
	defer s.iterMu.Unlock()

	for _, p := range s.iter {
		if !p.f(o) {
			s.deleteIter(p)
			close(p.done)
		}
	}
}

// Range calls f on all current and upcoming [Object] values tracked by s until
// f returns false or the context is cancelled. f must not retain o or modify
// the value it points to.
func (s *State) Range(ctx context.Context, f func(o *Object) bool) {
	done := make(chan struct{})
	p := pendingIterator{f, done}

	s.iterMu.Lock()
	s.ueventMu.RLock()
	for _, o := range s.uevent {
		if !f(o) {
			s.ueventMu.RUnlock()
			s.iterMu.Unlock()
			return
		}
	}
	s.ueventMu.RUnlock()
	s.iter = append(s.iter, &p)
	s.iterMu.Unlock()

	select {
	case <-ctx.Done():
		s.iterMu.Lock()
		s.deleteIter(&p)
		s.iterMu.Unlock()
		return

	case <-done:
		// deregistered by dispatchIter
		return
	}
}

// UnexpectedColdbootError is reported by [State.Consume] for a coldboot event
// with action other than the expected [uevent.KOBJ_ADD].
type UnexpectedColdbootError Event

func (e UnexpectedColdbootError) Error() string {
	return "unexpected " + e.Action.String() + " coldboot event"
}

// DuplicateAddError is reported by [State.Consume] for a [uevent.KOBJ_ADD]
// event on a still-existing entry that was not the result of a coldboot.
type DuplicateAddError Event

func (e DuplicateAddError) Error() string {
	return "duplicate add event on devpath " + strconv.Quote(e.DevPath)
}

// TargetError is reported by [State.Consume] for an event on a nonexistent
// entry. This is generally only possible before coldboot completes.
type TargetError Event

func (e TargetError) Error() string {
	return "unexpected " + e.Action.String() +
		" event on devpath " + strconv.Quote(e.DevPath)
}

// RemoveStateError is reported by [State.Consume] for a [uevent.KOBJ_REMOVE]
// event targeting an entry in a state other than [StateColdboot] and [StateNew].
type RemoveStateError Object

func (e RemoveStateError) Error() string {
	return "remove event targeting devpath " + strconv.Quote(e.DevPath) +
		" in state " + strconv.Itoa(e.State)
}

// BindStateError is reported by [State.Consume] for a [uevent.KOBJ_BIND] event
// targeting an entry in a state other than [StateColdboot] and [StateNew].
type BindStateError Object

func (e BindStateError) Error() string {
	return "bind event targeting devpath " + strconv.Quote(e.DevPath) +
		" in state " + strconv.Itoa(e.State)
}

// UnbindStateError is reported by [State.Consume] for a [uevent.KOBJ_UNBIND]
// event targeting an entry in a state other than [StateBound].
type UnbindStateError Object

func (e UnbindStateError) Error() string {
	return "unbind event targeting devpath " + strconv.Quote(e.DevPath) +
		" in state " + strconv.Itoa(e.State)
}

// MalformedMoveError is reported by [State.Consume] for a [uevent.KOBJ_MOVE]
// event missing the DEVPATH_OLD environment variable.
type MalformedMoveError Event

func (e MalformedMoveError) Error() string {
	return "move event targeting devpath " + strconv.Quote(e.DevPath) +
		" missing DEVPATH_OLD"
}

// UnexpectedOfflineError is reported by [State.Consume] for a
// [uevent.KOBJ_OFFLINE] or [uevent.KOBJ_ONLINE] event targeting an already
// offline or online object.
type UnexpectedOfflineError Object

func (e UnexpectedOfflineError) Error() string {
	if e.Offline {
		return "offline event targeting devpath " + strconv.Quote(e.DevPath)
	}
	return "online event targeting devpath " + strconv.Quote(e.DevPath)
}

// processEvent merges an event into s.
func (s *State) processEvent(e *Event) {
	s.ueventMu.Lock()
	defer s.ueventMu.Unlock()

	coldboot := e.Synth != nil
	if e.Action != uevent.KOBJ_ADD && coldboot {
		s.reportErr(UnexpectedColdbootError(e.Clone()))
		return
	}

	switch e.Action {
	case uevent.KOBJ_ADD:
		if e.Synth == nil {
			if o, ok := s.uevent[e.DevPath]; ok {
				s.reportErr(DuplicateAddError(e.Clone()))
				o.merge(e.Env)
				s.dispatchIter(o)
				return
			}
		}
		o := e.makeColdboot()
		if !coldboot {
			o.State = StateNew
		}
		o.merge(e.Env)
		s.uevent[e.DevPath] = o
		s.dispatchIter(o)
		return

	case uevent.KOBJ_REMOVE:
		if o, ok := s.uevent[e.DevPath]; !ok {
			s.reportErr(TargetError(e.Clone()))
			return
		} else if o.State != StateColdboot && o.State != StateNew {
			s.reportErr(RemoveStateError(o.Clone()))
		}
		delete(s.uevent, e.DevPath)
		return

	case uevent.KOBJ_CHANGE:
		o, ok := s.uevent[e.DevPath]
		if !ok {
			s.reportErr(TargetError(e.Clone()))
			// this suffers from the coldboot race window similar to KOBJ_MOVE,
			// however this action combines driver-specific and change-specific
			// environment variables and combines them with environment
			// variables meant to convey state of the kobject, and it is not
			// possible to reliably separate them, so this fallback avoids the
			// race at the cost of including some garbage in tracked state
			o = e.makeColdboot()
			o.merge(e.Env)
			s.uevent[e.DevPath] = o
			s.dispatchIter(o)
			return
		}
		o.update(e.Env, true)
		if s.handleChange != nil {
			s.handleChange(o, e.Env)
		}
		s.dispatchIter(o)
		return

	case uevent.KOBJ_MOVE:
		var o *Object
		if old, ok := e.Env["DEVPATH_OLD"]; !ok {
			s.reportErr(MalformedMoveError(e.Clone()))
			// not reached
			o = e.makeColdboot()
		} else if o, ok = s.uevent[old]; !ok {
			s.reportErr(TargetError(e.Clone()))
			// this generally happens during coldboot, dropping the event here
			// may cause inconsistent state if the coldboot event for this
			// object was generated before the bind event
			delete(e.Env, "DEVPATH_OLD")
			o = e.makeColdboot()
		} else {
			delete(s.uevent, old)
			delete(e.Env, "DEVPATH_OLD")
		}
		o.merge(e.Env)
		s.uevent[e.DevPath] = o
		o.DevPath = e.DevPath
		s.dispatchIter(o)
		return

	case uevent.KOBJ_ONLINE:
		o, ok := s.uevent[e.DevPath]
		if !ok {
			s.reportErr(TargetError(e.Clone()))
			// coldboot race window similar to an unexpected KOBJ_MOVE
			o = e.makeColdboot()
			s.uevent[e.DevPath] = o
			o.merge(e.Env)
		}
		if !o.Offline {
			s.reportErr(UnexpectedOfflineError(o.Clone()))
		}
		o.Offline = false
		s.dispatchIter(o)
		return

	case uevent.KOBJ_OFFLINE:
		o, ok := s.uevent[e.DevPath]
		if !ok {
			s.reportErr(TargetError(e.Clone()))
			// coldboot race window similar to an unexpected KOBJ_MOVE
			o = e.makeColdboot()
			s.uevent[e.DevPath] = o
			o.merge(e.Env)
		}
		if o.Offline {
			s.reportErr(UnexpectedOfflineError(o.Clone()))
		}
		o.Offline = true
		s.dispatchIter(o)
		return

	case uevent.KOBJ_BIND:
		o, ok := s.uevent[e.DevPath]
		if !ok {
			s.reportErr(TargetError(e.Clone()))
			// coldboot race window similar to an unexpected KOBJ_MOVE
			o = e.makeColdboot()
			s.uevent[e.DevPath] = o
		}
		if o.State != StateColdboot && o.State != StateNew {
			s.reportErr(BindStateError(o.Clone()))
		}
		o.State = StateBound
		o.merge(e.Env)
		s.dispatchIter(o)
		return

	case uevent.KOBJ_UNBIND:
		o, ok := s.uevent[e.DevPath]
		if !ok {
			s.reportErr(TargetError(e.Clone()))
			// coldboot race window similar to an unexpected KOBJ_MOVE, but does
			// not result in inconsistent state if dropped
			return
		}
		if o.State != StateBound {
			s.reportErr(UnbindStateError(o.Clone()))
		}
		o.State = StateNew
		o.Driver = ""
		s.dispatchIter(o)
		return

	default: // not reached
		s.reportErr(fmt.Errorf("invalid action %d", e.Action))
		return
	}
}

// BadSequenceError is reported by [State.Consume] for an unexpected SEQNUM.
type BadSequenceError struct{ Got, Want uint64 }

func (e BadSequenceError) Error() string {
	return "SEQNUM=" + strconv.FormatUint(e.Got, 10) +
		", want " + strconv.FormatUint(e.Want, 10)
}

// Consume receives uevent messages and updates s to reflect state of kernel.
func (s *State) Consume(ctx context.Context, events <-chan *uevent.Message) {
	if s.uevent == nil {
		s.uevent = make(map[string]*Object)
	}
	if s.reportErr == nil {
		s.reportErr = func(error) {}
	}

	var e Event
	for {
		select {
		case <-ctx.Done():
			return

		case m, ok := <-events:
			if !ok {
				return
			}
			e.Populate(s.reportErr, m)

			// skip external synthetic event
			if e.Synth != nil && *e.Synth != s.coldboot {
				continue
			}

			if s.seq == 0 {
				s.seq = e.Sequence
			}
			if s.seq != e.Sequence {
				s.reportErr(BadSequenceError{e.Sequence, s.seq})
			}
			s.seq++
			s.processEvent(&e)
		}
	}
}