package localcontrol

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"

	"iop/packages/go/agenttask"
	iop "iop/proto/gen/iop"

	"google.golang.org/protobuf/proto"
)

const (
	ledgerSchemaVersion   uint32 = 1
	ledgerRecordKey              = "local-control/ledger"
	defaultEventRetention        = 256
	maxEventRetention            = 256
	maxCommandRecords            = 4096
	maxLedgerCASAttempts         = 32
)

type ledgerStore interface {
	LoadIntegrationRecord(
		context.Context,
		string,
	) ([]byte, string, bool, error)
	CompareAndSwapIntegrationRecord(
		context.Context,
		string,
		string,
		[]byte,
	) (string, error)
}

type durableLedger struct {
	SchemaVersion uint32                    `json:"schema_version"`
	DaemonID      string                    `json:"daemon_id"`
	StateRevision uint64                    `json:"state_revision"`
	NextSequence  uint64                    `json:"next_sequence"`
	ReplayFloor   uint64                    `json:"replay_floor"`
	Commands      map[string]durableCommand `json:"commands,omitempty"`
	Events        []durableEvent            `json:"events,omitempty"`
}

type durableCommand struct {
	RequestHash   string `json:"request_hash"`
	StateRevision uint64 `json:"state_revision"`
	Complete      bool   `json:"complete"`
	Response      []byte `json:"response"`
}

type durableEvent struct {
	Sequence uint64 `json:"sequence"`
	Envelope []byte `json:"envelope"`
}

// Ledger stores the command-id acceptance record and replay window in one
// checksum-covered agentstate integration record.
type Ledger struct {
	store       ledgerStore
	daemonID    string
	retention   int
	mu          sync.Mutex
	flights     map[string]*commandFlight
	subMu       sync.Mutex
	nextSubID   uint64
	subscribers map[uint64]chan *iop.AgentLocalEnvelope
}

type commandOwnership uint8

const (
	commandCompleted commandOwnership = iota
	commandOwner
	commandWaiter
)

type commandFlight struct {
	done chan struct{}
	once sync.Once
}

type commandBegin struct {
	response  *iop.AgentLocalEnvelope
	ownership commandOwnership
	flight    *commandFlight
}

type LedgerView struct {
	DaemonID      string
	StateRevision uint64
	ReplayFloor   uint64
	ReplayCursor  uint64
}

func NewLedger(
	ctx context.Context,
	store ledgerStore,
	daemonID string,
	eventRetention int,
) (*Ledger, error) {
	if store == nil {
		return nil, fmt.Errorf("localcontrol: ledger store is required")
	}
	if !validIdentifier(daemonID, maxIdentifierBytes) {
		return nil, fmt.Errorf("localcontrol: daemon identity is invalid")
	}
	if eventRetention == 0 {
		eventRetention = defaultEventRetention
	}
	if eventRetention < 1 || eventRetention > maxEventRetention {
		return nil, fmt.Errorf(
			"localcontrol: event retention must be between 1 and %d",
			maxEventRetention,
		)
	}
	ledger := &Ledger{
		store:       store,
		daemonID:    daemonID,
		retention:   eventRetention,
		flights:     make(map[string]*commandFlight),
		subscribers: make(map[uint64]chan *iop.AgentLocalEnvelope),
	}
	if _, _, err := ledger.load(ctx); err != nil {
		return nil, err
	}
	return ledger, nil
}

// BeginCommand durably accepts a command before the host mutator is called.
// A matching completed record returns its final response. A matching
// incomplete record joins the active process-local flight, or becomes the
// recovery owner when the flight was lost with a prior process. A conflicting
// command ID returns command_id_conflict without changing the record.
func (l *Ledger) BeginCommand(
	ctx context.Context,
	commandID string,
	requestHash string,
	accepted *iop.AgentLocalEnvelope,
) (commandBegin, *ProtocolError) {
	if !validIdentifier(commandID, maxIdentifierBytes) ||
		!validIdentifier(requestHash, maxIdentifierBytes) ||
		accepted == nil {
		return commandBegin{}, malformed("command acceptance input is invalid")
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	for attempt := 0; attempt < maxLedgerCASAttempts; attempt++ {
		state, revision, err := l.load(ctx)
		if err != nil {
			return commandBegin{}, protocolError(
				ErrorInternal,
				"durable command state is unavailable",
			)
		}
		if record, exists := state.Commands[commandID]; exists {
			if record.RequestHash != requestHash {
				return commandBegin{}, protocolError(
					ErrorCommandIDConflict,
					"command_id was already used with different immutable arguments",
				)
			}
			if record.Complete {
				response, err := decodeStoredEnvelope(record.Response)
				if err != nil {
					return commandBegin{}, protocolError(
						ErrorInternal,
						"durable command state is invalid",
					)
				}
				return commandBegin{
					response:  response,
					ownership: commandCompleted,
				}, nil
			}
			if flight, active := l.flights[commandID]; active {
				return commandBegin{
					ownership: commandWaiter,
					flight:    flight,
				}, nil
			}
			flight := newCommandFlight()
			l.flights[commandID] = flight
			return commandBegin{
				ownership: commandOwner,
				flight:    flight,
			}, nil
		}
		if len(state.Commands) >= maxCommandRecords {
			return commandBegin{}, protocolError(
				ErrorInvalidState,
				"durable command ledger capacity is exhausted",
			)
		}
		if state.StateRevision == ^uint64(0) {
			return commandBegin{}, protocolError(
				ErrorInvalidState,
				"local-control state revision is exhausted",
			)
		}

		state.StateRevision++
		response := cloneEnvelope(accepted)
		applyResponseMetadata(response, state, commandID)
		responsePayload, err := marshalStoredEnvelope(response)
		if err != nil {
			return commandBegin{}, protocolError(
				ErrorInternal,
				"command response could not be recorded",
			)
		}
		state.Commands[commandID] = durableCommand{
			RequestHash:   requestHash,
			StateRevision: state.StateRevision,
			Response:      responsePayload,
		}
		if err := l.commit(ctx, revision, state); errors.Is(err, agenttask.ErrRevisionConflict) {
			continue
		} else if err != nil {
			return commandBegin{}, protocolError(
				ErrorInternal,
				"durable command state could not be recorded",
			)
		}
		flight := newCommandFlight()
		l.flights[commandID] = flight
		return commandBegin{
			ownership: commandOwner,
			flight:    flight,
		}, nil
	}
	return commandBegin{}, protocolError(
		ErrorInternal,
		"durable command state is busy",
	)
}

func newCommandFlight() *commandFlight {
	return &commandFlight{done: make(chan struct{})}
}

// WaitCommand waits only for an owner flight. Cancelling this wait never
// cancels the owner or changes durable command state.
func (l *Ledger) WaitCommand(
	ctx context.Context,
	flight *commandFlight,
) *ProtocolError {
	if flight == nil {
		return protocolError(ErrorInternal, "command completion flight is unavailable")
	}
	select {
	case <-flight.done:
		return nil
	case <-ctx.Done():
		return protocolError(ErrorInternal, "command completion wait was cancelled")
	}
}

// ReleaseCommandOwner unblocks same-process waiters after a durable completion
// or an owner failure. Waiters always reload the durable record before they
// can replay or recover ownership.
func (l *Ledger) ReleaseCommandOwner(commandID string, flight *commandFlight) {
	if flight == nil {
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.flights[commandID] == flight {
		delete(l.flights, commandID)
	}
	flight.once.Do(func() { close(flight.done) })
}

// FinishCommand records the final response and optional event atomically. The
// event becomes replayable only after the command was already durably accepted.
func (l *Ledger) FinishCommand(
	ctx context.Context,
	commandID string,
	requestHash string,
	response *iop.AgentLocalEnvelope,
	event *iop.AgentLocalEvent,
) (*iop.AgentLocalEnvelope, *ProtocolError) {
	if response == nil {
		return nil, malformed("command completion response is required")
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	for attempt := 0; attempt < maxLedgerCASAttempts; attempt++ {
		state, revision, err := l.load(ctx)
		if err != nil {
			return nil, protocolError(ErrorInternal, "durable command state is unavailable")
		}
		record, exists := state.Commands[commandID]
		if !exists || record.RequestHash != requestHash {
			return nil, protocolError(ErrorInvalidState, "durable command acceptance is missing")
		}
		if record.Complete {
			stored, err := decodeStoredEnvelope(record.Response)
			if err != nil {
				return nil, protocolError(ErrorInternal, "durable command state is invalid")
			}
			return stored, nil
		}

		finalResponse := cloneEnvelope(response)
		applyResponseMetadataAtRevision(
			finalResponse,
			state,
			commandID,
			record.StateRevision,
		)
		if event != nil {
			if state.NextSequence == ^uint64(0) {
				return nil, protocolError(
					ErrorInvalidState,
					"local-control event sequence is exhausted",
				)
			}
			sequence := state.NextSequence
			state.NextSequence++
			eventCopy := proto.Clone(event).(*iop.AgentLocalEvent)
			eventCopy.EventSequence = sequence
			eventCopy.StateRevision = record.StateRevision
			eventEnvelope := &iop.AgentLocalEnvelope{
				ProtocolVersion: ProtocolVersion,
				Kind:            iop.AgentLocalKind_AGENT_LOCAL_KIND_EVENT,
				MessageId: derivedMessageID(
					"event",
					l.daemonID+":"+strconv.FormatUint(sequence, 10),
				),
				EventSequence: sequence,
				Operation:     finalResponse.GetOperation(),
				Payload: &iop.AgentLocalEnvelope_Event{
					Event: eventCopy,
				},
			}
			eventPayload, err := marshalStoredEnvelope(eventEnvelope)
			if err != nil {
				return nil, protocolError(ErrorInternal, "local-control event could not be recorded")
			}
			state.Events = append(state.Events, durableEvent{
				Sequence: sequence,
				Envelope: eventPayload,
			})
			if excess := len(state.Events) - l.retention; excess > 0 {
				state.ReplayFloor = state.Events[excess-1].Sequence
				state.Events = append([]durableEvent(nil), state.Events[excess:]...)
			}
			if payload := finalResponse.GetResponse(); payload != nil {
				payload.ReplayCursor = sequence
			}
			event = eventCopy
		}

		responsePayload, err := marshalStoredEnvelope(finalResponse)
		if err != nil {
			return nil, protocolError(ErrorInternal, "command response could not be recorded")
		}
		record.Response = responsePayload
		record.Complete = true
		state.Commands[commandID] = record
		if err := l.commit(ctx, revision, state); errors.Is(err, agenttask.ErrRevisionConflict) {
			continue
		} else if err != nil {
			return nil, protocolError(
				ErrorInternal,
				"durable command result could not be recorded",
			)
		}
		if event != nil {
			l.publish(&iop.AgentLocalEnvelope{
				ProtocolVersion: ProtocolVersion,
				Kind:            iop.AgentLocalKind_AGENT_LOCAL_KIND_EVENT,
				MessageId: derivedMessageID(
					"event",
					l.daemonID+":"+strconv.FormatUint(event.GetEventSequence(), 10),
				),
				EventSequence: event.GetEventSequence(),
				Operation:     finalResponse.GetOperation(),
				Payload: &iop.AgentLocalEnvelope_Event{
					Event: proto.Clone(event).(*iop.AgentLocalEvent),
				},
			})
		}
		return finalResponse, nil
	}
	return nil, protocolError(ErrorInternal, "durable command state is busy")
}

// Subscribe returns best-effort live delivery of committed events. Retained
// replay remains authoritative: an overflowing live subscriber is closed so it
// must reconnect with its last contiguous cursor.
func (l *Ledger) Subscribe() (<-chan *iop.AgentLocalEnvelope, func()) {
	l.subMu.Lock()
	l.nextSubID++
	id := l.nextSubID
	channel := make(chan *iop.AgentLocalEnvelope, l.retention)
	l.subscribers[id] = channel
	l.subMu.Unlock()
	var once sync.Once
	return channel, func() {
		once.Do(func() {
			l.subMu.Lock()
			if current, exists := l.subscribers[id]; exists {
				delete(l.subscribers, id)
				close(current)
			}
			l.subMu.Unlock()
		})
	}
}

func (l *Ledger) publish(event *iop.AgentLocalEnvelope) {
	l.subMu.Lock()
	defer l.subMu.Unlock()
	for id, subscriber := range l.subscribers {
		select {
		case subscriber <- cloneEnvelope(event):
		default:
			delete(l.subscribers, id)
			close(subscriber)
		}
	}
}

func (l *Ledger) View(ctx context.Context) (LedgerView, *ProtocolError) {
	l.mu.Lock()
	defer l.mu.Unlock()
	state, _, err := l.load(ctx)
	if err != nil {
		return LedgerView{}, protocolError(ErrorInternal, "durable command state is unavailable")
	}
	return viewOf(state), nil
}

func (l *Ledger) Replay(
	ctx context.Context,
	daemonID string,
	after uint64,
) ([]*iop.AgentLocalEvent, LedgerView, *ProtocolError) {
	l.mu.Lock()
	defer l.mu.Unlock()
	state, _, err := l.load(ctx)
	if err != nil {
		return nil, LedgerView{}, protocolError(
			ErrorInternal,
			"durable event state is unavailable",
		)
	}
	view := viewOf(state)
	if daemonID != state.DaemonID ||
		after < state.ReplayFloor ||
		after > view.ReplayCursor {
		return nil, view, &ProtocolError{
			Code:             ErrorReplayUnavailable,
			SafeMessage:      "requested event replay is unavailable; request a fresh snapshot",
			ReplayFloor:      state.ReplayFloor,
			SnapshotRequired: true,
			SnapshotMarker:   snapshotMarker(state.DaemonID, state.StateRevision),
		}
	}
	events := make([]*iop.AgentLocalEvent, 0, len(state.Events))
	expected := after + 1
	for _, record := range state.Events {
		if record.Sequence <= after {
			continue
		}
		if record.Sequence != expected {
			return nil, view, &ProtocolError{
				Code:             ErrorReplayUnavailable,
				SafeMessage:      "requested event replay is not contiguous; request a fresh snapshot",
				ReplayFloor:      state.ReplayFloor,
				SnapshotRequired: true,
				SnapshotMarker:   snapshotMarker(state.DaemonID, state.StateRevision),
			}
		}
		envelope, err := decodeStoredEnvelope(record.Envelope)
		if err != nil || envelope.GetEvent() == nil {
			return nil, view, protocolError(ErrorInternal, "durable event state is invalid")
		}
		events = append(events, proto.Clone(envelope.GetEvent()).(*iop.AgentLocalEvent))
		expected++
	}
	if expected != view.ReplayCursor+1 {
		return nil, view, &ProtocolError{
			Code:             ErrorReplayUnavailable,
			SafeMessage:      "requested event replay is not contiguous; request a fresh snapshot",
			ReplayFloor:      state.ReplayFloor,
			SnapshotRequired: true,
			SnapshotMarker:   snapshotMarker(state.DaemonID, state.StateRevision),
		}
	}
	return events, view, nil
}

func (l *Ledger) load(
	ctx context.Context,
) (durableLedger, string, error) {
	payload, revision, found, err := l.store.LoadIntegrationRecord(
		ctx,
		ledgerRecordKey,
	)
	if err != nil {
		return durableLedger{}, "", err
	}
	if !found {
		return durableLedger{
			SchemaVersion: ledgerSchemaVersion,
			DaemonID:      l.daemonID,
			NextSequence:  1,
			Commands:      make(map[string]durableCommand),
		}, "", nil
	}
	state, err := decodeLedger(payload, l.daemonID)
	if err != nil {
		return durableLedger{}, "", err
	}
	return state, revision, nil
}

func (l *Ledger) commit(
	ctx context.Context,
	revision string,
	state durableLedger,
) error {
	payload, err := json.Marshal(state)
	if err != nil {
		return err
	}
	_, err = l.store.CompareAndSwapIntegrationRecord(
		ctx,
		ledgerRecordKey,
		revision,
		payload,
	)
	return err
}

func decodeLedger(payload []byte, daemonID string) (durableLedger, error) {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	var state durableLedger
	if err := decoder.Decode(&state); err != nil {
		return durableLedger{}, fmt.Errorf("localcontrol: decode ledger: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		return durableLedger{}, fmt.Errorf("localcontrol: trailing ledger data")
	}
	if state.SchemaVersion != ledgerSchemaVersion ||
		state.DaemonID != daemonID ||
		state.NextSequence == 0 ||
		state.ReplayFloor >= state.NextSequence {
		return durableLedger{}, fmt.Errorf("localcontrol: invalid ledger identity or revision")
	}
	if state.Commands == nil {
		state.Commands = make(map[string]durableCommand)
	}
	if len(state.Commands) > maxCommandRecords {
		return durableLedger{}, fmt.Errorf("localcontrol: command ledger exceeds its bound")
	}
	for commandID, record := range state.Commands {
		if !validIdentifier(commandID, maxIdentifierBytes) ||
			!validIdentifier(record.RequestHash, maxIdentifierBytes) ||
			record.StateRevision == 0 ||
			record.StateRevision > state.StateRevision {
			return durableLedger{}, fmt.Errorf("localcontrol: invalid command record")
		}
		envelope, err := decodeStoredEnvelope(record.Response)
		if err != nil ||
			(envelope.GetKind() != iop.AgentLocalKind_AGENT_LOCAL_KIND_RESPONSE &&
				envelope.GetKind() != iop.AgentLocalKind_AGENT_LOCAL_KIND_ERROR) {
			return durableLedger{}, fmt.Errorf("localcontrol: invalid command response")
		}
	}
	if len(state.Events) > maxEventRetention {
		return durableLedger{}, fmt.Errorf("localcontrol: event ledger exceeds its bound")
	}
	expected := state.ReplayFloor + 1
	for _, record := range state.Events {
		if record.Sequence != expected {
			return durableLedger{}, fmt.Errorf("localcontrol: event sequence is not contiguous")
		}
		envelope, err := decodeStoredEnvelope(record.Envelope)
		if err != nil ||
			envelope.GetKind() != iop.AgentLocalKind_AGENT_LOCAL_KIND_EVENT ||
			envelope.GetEventSequence() != record.Sequence ||
			envelope.GetEvent() == nil ||
			envelope.GetEvent().GetEventSequence() != record.Sequence {
			return durableLedger{}, fmt.Errorf("localcontrol: invalid retained event")
		}
		expected++
	}
	if expected != state.NextSequence {
		return durableLedger{}, fmt.Errorf("localcontrol: next event sequence is invalid")
	}
	return state, nil
}

func marshalStoredEnvelope(envelope *iop.AgentLocalEnvelope) ([]byte, error) {
	if envelope == nil ||
		envelope.GetProtocolVersion() != ProtocolVersion ||
		proto.Size(envelope) <= 0 ||
		proto.Size(envelope) > MaxFrameBytes ||
		hasUnknownFields(envelope.ProtoReflect()) {
		return nil, fmt.Errorf("localcontrol: invalid stored envelope")
	}
	return (proto.MarshalOptions{Deterministic: true}).Marshal(envelope)
}

func decodeStoredEnvelope(payload []byte) (*iop.AgentLocalEnvelope, error) {
	envelope, err := ParseEnvelope(payload)
	if err != nil {
		return nil, err
	}
	if envelope.GetProtocolVersion() != ProtocolVersion ||
		hasUnknownFields(envelope.ProtoReflect()) ||
		!validIdentifier(envelope.GetMessageId(), maxIdentifierBytes) {
		return nil, fmt.Errorf("localcontrol: invalid stored envelope")
	}
	return envelope, nil
}

func applyResponseMetadata(
	envelope *iop.AgentLocalEnvelope,
	state durableLedger,
	commandID string,
) {
	applyResponseMetadataAtRevision(
		envelope,
		state,
		commandID,
		state.StateRevision,
	)
}

func applyResponseMetadataAtRevision(
	envelope *iop.AgentLocalEnvelope,
	state durableLedger,
	commandID string,
	revision uint64,
) {
	response := envelope.GetResponse()
	if response == nil {
		return
	}
	response.CommandId = commandID
	response.StateRevision = revision
	response.SnapshotMarker = snapshotMarker(state.DaemonID, revision)
	response.ReplayDaemonId = state.DaemonID
	response.ReplayCursor = state.NextSequence - 1
}

func viewOf(state durableLedger) LedgerView {
	return LedgerView{
		DaemonID:      state.DaemonID,
		StateRevision: state.StateRevision,
		ReplayFloor:   state.ReplayFloor,
		ReplayCursor:  state.NextSequence - 1,
	}
}

func snapshotMarker(daemonID string, revision uint64) string {
	return daemonID + ":" + strconv.FormatUint(revision, 10)
}

func cloneEnvelope(envelope *iop.AgentLocalEnvelope) *iop.AgentLocalEnvelope {
	if envelope == nil {
		return nil
	}
	return proto.Clone(envelope).(*iop.AgentLocalEnvelope)
}