package openai

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"
)

const (
	defaultLogicalRequestCapacity         = 1024
	defaultLogicalRequestTTL              = 30 * time.Minute
	defaultLogicalRequestFrontierCapacity = 64
	defaultLogicalRequestMappingCapacity  = 512
)

var (
	errLogicalRequestNotFound        = errors.New("logical request state is unavailable")
	errLogicalRequestOwnerMismatch   = errors.New("logical request owner mismatch")
	errLogicalRequestPrincipal       = errors.New("logical request principal mismatch")
	errLogicalRequestLineage         = errors.New("logical request lineage mismatch")
	errLogicalRequestFrontier        = errors.New("logical request frontier mismatch")
	errLogicalRequestNoFrontier      = errors.New("logical request has no unconsumed frontier")
	errLogicalRequestActiveStage     = errors.New("logical request already has an active stage")
	errLogicalRequestCapacityReached = errors.New("logical request coordinator capacity reached")
)

type logicalRequestState string

const (
	logicalRequestStateAccepted logicalRequestState = "accepted"
	logicalRequestStateActive   logicalRequestState = "active"
	logicalRequestStateWaiting  logicalRequestState = "agent_tool_wait"
	logicalRequestStateResumed  logicalRequestState = "resumed"
	logicalRequestStateCleanup  logicalRequestState = "cleanup_pending"
	logicalRequestStateDetached logicalRequestState = "disconnected"
)

type logicalRequestCoordinatorOptions struct {
	Capacity         int
	TTL              time.Duration
	FrontierCapacity int
	MappingCapacity  int
	Now              func() time.Time
	IDSource         func() (string, error)
}

type logicalRequestAdmission struct {
	OwnerEdgeID      string
	PrincipalRef     string
	Lineage          logicalRequestLineage
	PresetGeneration string
}

type logicalRequestExpectedTool struct {
	PublicCallID   string
	ProviderCallID string
}

type logicalRequestToolResult struct {
	PublicCallID string
}

type logicalRequestContinuation struct {
	RequestID    string
	OwnerEdgeID  string
	PrincipalRef string
	Lineage      logicalRequestContinuationLineage
	Results      []logicalRequestToolResult
}

// logicalRequestSnapshot is a deliberately payload-free view suitable for
// handlers and tests. It is copied while the coordinator lock is held.
type logicalRequestSnapshot struct {
	ID               string
	State            logicalRequestState
	OwnerEdgeID      string
	PrincipalRef     string
	PresetGeneration string
	ActiveStageID    string
	ExpectedCallIDs  []string
	TerminalClass    string
	CreatedAt        time.Time
	UpdatedAt        time.Time
}

type logicalRequestRecord struct {
	id                     string
	ownerEdgeID            string
	principalRef           string
	lineage                logicalRequestLineage
	presetGeneration       string
	state                  logicalRequestState
	activeStageID          string
	expected               map[string]string // public tool-call id -> provider tool-call id
	expectedIssuedCallHash string
	publicToProvider       map[string]string
	providerToPublic       map[string]string
	cleanup                bool
	terminalClass          string
	createdAt              time.Time
	updatedAt              time.Time
}

// logicalRequestCoordinator owns the transient Edge-local continuation state.
// It is intentionally independent of HTTP handlers so wire-specific callers
// can supply their canonical immutable prefix and result frontier.
type logicalRequestCoordinator struct {
	mu               sync.Mutex
	capacity         int
	ttl              time.Duration
	frontierCapacity int
	mappingCapacity  int
	now              func() time.Time
	idSource         func() (string, error)
	requests         map[string]*logicalRequestRecord
}

func newLogicalRequestCoordinator(options logicalRequestCoordinatorOptions) *logicalRequestCoordinator {
	capacity := options.Capacity
	if capacity <= 0 {
		capacity = defaultLogicalRequestCapacity
	}
	ttl := options.TTL
	if ttl <= 0 {
		ttl = defaultLogicalRequestTTL
	}
	frontierCapacity := options.FrontierCapacity
	if frontierCapacity <= 0 {
		frontierCapacity = defaultLogicalRequestFrontierCapacity
	}
	mappingCapacity := options.MappingCapacity
	if mappingCapacity <= 0 {
		mappingCapacity = defaultLogicalRequestMappingCapacity
	}
	now := options.Now
	if now == nil {
		now = time.Now
	}
	idSource := options.IDSource
	if idSource == nil {
		idSource = newLogicalRequestRandomID
	}
	return &logicalRequestCoordinator{
		capacity: capacity, ttl: ttl, frontierCapacity: frontierCapacity, mappingCapacity: mappingCapacity, now: now, idSource: idSource,
		requests: make(map[string]*logicalRequestRecord),
	}
}

func newLogicalRequestRandomID() (string, error) {
	buf := make([]byte, 18) // 144 bits; the public ID is not an authorization secret.
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("read logical request random id: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(buf), nil
}

func (c *logicalRequestCoordinator) create(admission logicalRequestAdmission) (logicalRequestSnapshot, error) {
	if err := validateLogicalRequestAdmission(admission); err != nil {
		return logicalRequestSnapshot{}, err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	now := c.now()
	if len(c.requests) >= c.capacity {
		return logicalRequestSnapshot{}, errLogicalRequestCapacityReached
	}
	for attempts := 0; attempts < 32; attempts++ {
		id, err := c.allocateID("req")
		if err != nil {
			return logicalRequestSnapshot{}, err
		}
		if _, exists := c.requests[id]; exists {
			continue
		}
		record := &logicalRequestRecord{
			id: id, ownerEdgeID: admission.OwnerEdgeID, principalRef: admission.PrincipalRef,
			lineage: admission.Lineage, presetGeneration: admission.PresetGeneration,
			state: logicalRequestStateAccepted, publicToProvider: make(map[string]string),
			providerToPublic: make(map[string]string), createdAt: now, updatedAt: now,
		}
		c.requests[id] = record
		return record.snapshot(), nil
	}
	return logicalRequestSnapshot{}, fmt.Errorf("could not allocate unique logical request id")
}

// newStageID and newCallID issue endpoint-safe opaque identities. They do not
// carry authority; ownership remains enforced by the request record.
func (c *logicalRequestCoordinator) newStageID() (string, error) { return c.allocateID("stg") }

func (c *logicalRequestCoordinator) newCallID() (string, error) { return c.allocateID("call") }

func (c *logicalRequestCoordinator) allocateID(prefix string) (string, error) {
	rawID, err := c.idSource()
	if err != nil {
		return "", err
	}
	if !validLogicalRequestID(rawID) {
		return "", fmt.Errorf("invalid logical request id from source")
	}
	return prefix + "_" + rawID, nil
}

// activateStage gives the request exactly one active stage. A later handler
// must consume a frontier before it can activate a replacement stage.
func (c *logicalRequestCoordinator) activateStage(requestID, ownerEdgeID, stageID string) (logicalRequestSnapshot, error) {
	if !validLogicalRequestID(stageID) {
		return logicalRequestSnapshot{}, fmt.Errorf("invalid logical request stage id")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return logicalRequestSnapshot{}, err
	}
	if record.activeStageID != "" || record.expected != nil {
		return logicalRequestSnapshot{}, errLogicalRequestActiveStage
	}
	record.activeStageID = stageID
	record.state = logicalRequestStateActive
	record.updatedAt = c.now()
	return record.snapshot(), nil
}

// transitionStage commits a tool-free stage terminal and installs the next
// pinned stage without exposing an intermediate resumable state. This is the
// local-completion to review transaction boundary for the light flow.
func (c *logicalRequestCoordinator) transitionStage(requestID, ownerEdgeID, fromStageID, toStageID string) (logicalRequestSnapshot, error) {
	if !validLogicalRequestID(fromStageID) || !validLogicalRequestID(toStageID) || fromStageID == toStageID {
		return logicalRequestSnapshot{}, fmt.Errorf("invalid logical request stage transition")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return logicalRequestSnapshot{}, err
	}
	if record.state != logicalRequestStateActive || record.activeStageID != fromStageID || record.expected != nil {
		return logicalRequestSnapshot{}, errLogicalRequestActiveStage
	}
	record.activeStageID = toStageID
	record.state = logicalRequestStateActive
	record.updatedAt = c.now()
	return record.snapshot(), nil
}

// startCleanup transfers the active request to one cleanup stage. A primary
// artifact error may start from the resumed frontier, while review completion
// must name the exact active stage it is replacing.
func (c *logicalRequestCoordinator) startCleanup(requestID, ownerEdgeID, fromStageID, cleanupStageID, terminalClass string) (logicalRequestSnapshot, error) {
	if !validLogicalRequestID(cleanupStageID) || strings.TrimSpace(terminalClass) == "" {
		return logicalRequestSnapshot{}, fmt.Errorf("invalid logical request cleanup identity")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return logicalRequestSnapshot{}, err
	}
	if fromStageID == "" {
		if record.state != logicalRequestStateResumed || record.activeStageID != "" || record.expected != nil {
			return logicalRequestSnapshot{}, errLogicalRequestActiveStage
		}
	} else if record.state != logicalRequestStateActive || record.activeStageID != fromStageID || record.expected != nil {
		return logicalRequestSnapshot{}, errLogicalRequestActiveStage
	}
	record.activeStageID = cleanupStageID
	record.state = logicalRequestStateActive
	record.cleanup = true
	record.terminalClass = terminalClass
	record.updatedAt = c.now()
	return record.snapshot(), nil
}

// awaitToolResults pins the public/provider tool mapping and creates the sole
// next continuation frontier. All expected results must arrive in one call,
// but their order is intentionally irrelevant.
func (c *logicalRequestCoordinator) awaitToolResults(requestID, ownerEdgeID, stageID string, expected []logicalRequestExpectedTool, expectedIssuedCallHash string) (logicalRequestSnapshot, error) {
	if strings.TrimSpace(expectedIssuedCallHash) == "" {
		return logicalRequestSnapshot{}, fmt.Errorf("issued-call hash is required")
	}
	if len(expected) == 0 {
		return logicalRequestSnapshot{}, fmt.Errorf("logical request frontier is empty")
	}
	if len(expected) > c.frontierCapacity {
		return logicalRequestSnapshot{}, errLogicalRequestFrontier
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return logicalRequestSnapshot{}, err
	}
	if record.state != logicalRequestStateActive || record.activeStageID != stageID || record.expected != nil {
		return logicalRequestSnapshot{}, errLogicalRequestFrontier
	}
	if len(record.publicToProvider)+len(expected) > c.mappingCapacity {
		return logicalRequestSnapshot{}, errLogicalRequestFrontier
	}
	frontier := make(map[string]string, len(expected))
	providers := make(map[string]struct{}, len(expected))
	for _, item := range expected {
		if !validLogicalRequestID(item.PublicCallID) || !validLogicalRequestID(item.ProviderCallID) {
			return logicalRequestSnapshot{}, fmt.Errorf("invalid logical request tool id")
		}
		if _, duplicate := frontier[item.PublicCallID]; duplicate {
			return logicalRequestSnapshot{}, errLogicalRequestFrontier
		}
		if _, duplicate := providers[item.ProviderCallID]; duplicate {
			return logicalRequestSnapshot{}, errLogicalRequestFrontier
		}
		if _, exists := record.publicToProvider[item.PublicCallID]; exists {
			return logicalRequestSnapshot{}, errLogicalRequestFrontier
		}
		if _, exists := record.providerToPublic[item.ProviderCallID]; exists {
			return logicalRequestSnapshot{}, errLogicalRequestFrontier
		}
		frontier[item.PublicCallID] = item.ProviderCallID
		providers[item.ProviderCallID] = struct{}{}
	}
	for public, provider := range frontier {
		record.publicToProvider[public] = provider
		record.providerToPublic[provider] = public
	}
	record.expected = frontier
	record.expectedIssuedCallHash = expectedIssuedCallHash
	if record.cleanup {
		record.state = logicalRequestStateCleanup
	} else {
		record.state = logicalRequestStateWaiting
	}
	record.updatedAt = c.now()
	return record.snapshot(), nil
}

// consumeContinuation performs every validation before changing state. Holding
// the coordinator lock across validation and consume makes duplicate resumes
// deterministic: exactly one concurrent caller can consume a frontier.
func (c *logicalRequestCoordinator) consumeContinuation(continuation logicalRequestContinuation) (logicalRequestSnapshot, error) {
	if continuation.RequestID == "" {
		return c.consumeContinuationByLineage(continuation.OwnerEdgeID, continuation.PrincipalRef, continuation.Lineage)
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(continuation.RequestID, continuation.OwnerEdgeID, c.now())
	if err != nil {
		return logicalRequestSnapshot{}, err
	}
	if record.principalRef != continuation.PrincipalRef {
		return logicalRequestSnapshot{}, errLogicalRequestPrincipal
	}
	if record.expected == nil {
		return logicalRequestSnapshot{}, errLogicalRequestNoFrontier
	}
	if err := validateLogicalRequestContinuationLineage(record.lineage, record.expectedIssuedCallHash, record.expected, continuation.Lineage); err != nil {
		return logicalRequestSnapshot{}, err
	}
	if !sameLogicalRequestResultSet(record.expected, continuation.Results) {
		return logicalRequestSnapshot{}, errLogicalRequestFrontier
	}
	record.expected = nil
	record.expectedIssuedCallHash = ""
	record.lineage = continuation.Lineage.Committed
	record.activeStageID = ""
	record.state = logicalRequestStateResumed
	record.updatedAt = c.now()
	return record.snapshot(), nil
}

func (c *logicalRequestCoordinator) consumeContinuationByLineage(ownerEdgeID, principalRef string, lineage logicalRequestContinuationLineage) (logicalRequestSnapshot, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	now := c.now()

	var target *logicalRequestRecord
	for _, record := range c.requests {
		if record.ownerEdgeID == ownerEdgeID && record.principalRef == principalRef && record.state == logicalRequestStateWaiting {
			if record.lineage == lineage.Prefix && sameLogicalRequestResultIDs(record.expected, lineage.ResultIDs) {
				target = record
				break
			}
		}
	}
	if target == nil {
		for _, record := range c.requests {
			if record.state == logicalRequestStateWaiting && sameLogicalRequestResultIDs(record.expected, lineage.ResultIDs) {
				if record.ownerEdgeID != ownerEdgeID {
					return logicalRequestSnapshot{}, errLogicalRequestOwnerMismatch
				}
				if record.principalRef != principalRef {
					return logicalRequestSnapshot{}, errLogicalRequestPrincipal
				}
				if record.lineage != lineage.Prefix {
					return logicalRequestSnapshot{}, errLogicalRequestLineage
				}
			}
		}
		return logicalRequestSnapshot{}, errLogicalRequestNotFound
	}

	results := make([]logicalRequestToolResult, 0, len(lineage.ResultIDs))
	for _, id := range lineage.ResultIDs {
		results = append(results, logicalRequestToolResult{PublicCallID: id})
	}

	if target.expected == nil {
		return logicalRequestSnapshot{}, errLogicalRequestNoFrontier
	}
	if err := validateLogicalRequestContinuationLineage(target.lineage, target.expectedIssuedCallHash, target.expected, lineage); err != nil {
		return logicalRequestSnapshot{}, err
	}
	if !sameLogicalRequestResultSet(target.expected, results) {
		return logicalRequestSnapshot{}, errLogicalRequestFrontier
	}

	target.expected = nil
	target.expectedIssuedCallHash = ""
	target.lineage = lineage.Committed
	target.activeStageID = ""
	target.state = logicalRequestStateResumed
	target.updatedAt = now
	return target.snapshot(), nil
}

func (c *logicalRequestCoordinator) snapshot(requestID string) (logicalRequestSnapshot, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	record, ok := c.requests[requestID]
	if !ok || c.expiredForSweepLocked(record, c.now()) {
		return logicalRequestSnapshot{}, errLogicalRequestNotFound
	}
	return record.snapshot(), nil
}

func (c *logicalRequestCoordinator) publicToolID(requestID, providerCallID string) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	record, ok := c.requests[requestID]
	if !ok || c.expiredForSweepLocked(record, c.now()) {
		return "", errLogicalRequestNotFound
	}
	public, ok := record.providerToPublic[providerCallID]
	if !ok {
		return "", errLogicalRequestFrontier
	}
	return public, nil
}

func (c *logicalRequestCoordinator) terminal(requestID, ownerEdgeID string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return err
	}
	delete(c.requests, record.id)
	return nil
}

func (c *logicalRequestCoordinator) disconnect(requestID, ownerEdgeID, terminalClass string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	record, err := c.getOwnedLocked(requestID, ownerEdgeID, c.now())
	if err != nil {
		return err
	}
	record.state = logicalRequestStateDetached
	record.terminalClass = strings.TrimSpace(terminalClass)
	if record.terminalClass == "" {
		record.terminalClass = "cancelled"
	}
	record.updatedAt = c.now()
	return nil
}

func (c *logicalRequestCoordinator) removeOwned(requestID, ownerEdgeID string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	record, ok := c.requests[requestID]
	if !ok {
		return errLogicalRequestNotFound
	}
	if record.ownerEdgeID != ownerEdgeID {
		return errLogicalRequestOwnerMismatch
	}
	delete(c.requests, requestID)
	return nil
}

func (c *logicalRequestCoordinator) getOwnedLocked(requestID, ownerEdgeID string, now time.Time) (*logicalRequestRecord, error) {
	record, ok := c.requests[requestID]
	if !ok || c.expiredForSweepLocked(record, now) {
		return nil, errLogicalRequestNotFound
	}
	if record.ownerEdgeID != ownerEdgeID {
		return nil, errLogicalRequestOwnerMismatch
	}
	return record, nil
}

func (r *logicalRequestRecord) snapshot() logicalRequestSnapshot {
	expected := make([]string, 0, len(r.expected))
	for id := range r.expected {
		expected = append(expected, id)
	}
	return logicalRequestSnapshot{
		ID: r.id, State: r.state, OwnerEdgeID: r.ownerEdgeID, PrincipalRef: r.principalRef,
		PresetGeneration: r.presetGeneration, ActiveStageID: r.activeStageID,
		ExpectedCallIDs: expected, TerminalClass: r.terminalClass,
		CreatedAt: r.createdAt, UpdatedAt: r.updatedAt,
	}
}

func sameLogicalRequestResultSet(expected map[string]string, results []logicalRequestToolResult) bool {
	if len(expected) != len(results) {
		return false
	}
	seen := make(map[string]struct{}, len(results))
	for _, result := range results {
		if _, ok := expected[result.PublicCallID]; !ok {
			return false
		}
		if _, duplicate := seen[result.PublicCallID]; duplicate {
			return false
		}
		seen[result.PublicCallID] = struct{}{}
	}
	return true
}

func sameLogicalRequestResultIDs(expected map[string]string, resultIDs []string) bool {
	if len(expected) != len(resultIDs) {
		return false
	}
	seen := make(map[string]struct{}, len(resultIDs))
	for _, id := range resultIDs {
		if _, ok := expected[id]; !ok {
			return false
		}
		if _, duplicate := seen[id]; duplicate {
			return false
		}
		seen[id] = struct{}{}
	}
	return true
}

func validateLogicalRequestAdmission(admission logicalRequestAdmission) error {
	if strings.TrimSpace(admission.OwnerEdgeID) == "" || strings.TrimSpace(admission.PrincipalRef) == "" {
		return fmt.Errorf("logical request owner and principal are required")
	}
	if strings.TrimSpace(admission.PresetGeneration) == "" {
		return fmt.Errorf("logical request preset generation is required")
	}
	if admission.Lineage.Endpoint == "" || admission.Lineage.HistoryDigest == "" || admission.Lineage.ToolsetDigest == "" {
		return fmt.Errorf("logical request lineage is incomplete")
	}
	return nil
}

func validateLogicalRequestContinuationLineage(prefix logicalRequestLineage, expectedIssuedCallHash string, expected map[string]string, lineage logicalRequestContinuationLineage) error {
	if strings.TrimSpace(lineage.IssuedCallHash) == "" || lineage.IssuedCallHash != expectedIssuedCallHash {
		return errLogicalRequestLineage
	}
	if lineage.Prefix != prefix {
		return errLogicalRequestLineage
	}
	if lineage.Committed.Endpoint == "" || lineage.Committed.HistoryDigest == "" || lineage.Committed.ToolsetDigest == "" {
		return errLogicalRequestLineage
	}
	if lineage.Committed.Endpoint != prefix.Endpoint || lineage.Committed.ToolsetDigest != prefix.ToolsetDigest {
		return errLogicalRequestLineage
	}
	if lineage.Committed.HistoryDigest == prefix.HistoryDigest {
		return errLogicalRequestLineage
	}
	if len(lineage.ResultIDs) == 0 || !sameLogicalRequestResultIDs(expected, lineage.ResultIDs) {
		return errLogicalRequestFrontier
	}
	return nil
}

func validLogicalRequestID(value string) bool {
	if value == "" || len(value) > 256 {
		return false
	}
	for _, r := range value {
		if !((r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' || r == '-') {
			return false
		}
	}
	return true
}
