package openai

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

	"iop/packages/go/streamgate"
)

const (
	openAIRebuildEndpointChat      = "/v1/chat/completions"
	openAIRebuildEndpointResponses = "/v1/responses"
	openAIRebuildFamily            = "openai.json"
)

var openAIRebuiltSequence atomic.Uint64

var errOpenAIRecoveryPatchDuplicate = errors.New("duplicate OpenAI recovery patch")

type openAIRecoveryPatchEntry struct {
	cursor int
	value  []byte
	guard  *streamgate.IngressSnapshotRebuildGuard
	once   sync.Once
}

func (e *openAIRecoveryPatchEntry) release() {
	if e == nil {
		return
	}
	e.once.Do(func() {
		e.value = nil
		if e.guard != nil {
			e.guard.Close()
			e.guard = nil
		}
	})
}

// openAIRecoveryPatchStore is request-local. It retains only the typed subtree
// named by a safe directive reference; it never retains a second request tree,
// caller headers, or provider credentials.
type openAIRecoveryPatchStore struct {
	mu           sync.Mutex
	ingress      *openAIIngressSnapshot
	continuation map[string]*openAIRecoveryPatchEntry
	schema       map[string]*openAIRecoveryPatchEntry
	closed       bool
}

func newOpenAIRecoveryPatchStore(ingress *openAIIngressSnapshot) *openAIRecoveryPatchStore {
	return &openAIRecoveryPatchStore{
		ingress:      ingress,
		continuation: make(map[string]*openAIRecoveryPatchEntry),
		schema:       make(map[string]*openAIRecoveryPatchEntry),
	}
}

func marshalOpenAIPatchValue(value any) ([]byte, error) {
	if raw, ok := value.(json.RawMessage); ok {
		if !json.Valid(raw) {
			return nil, fmt.Errorf("invalid OpenAI recovery patch JSON")
		}
		return raw, nil
	}
	encoded, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}
	return encoded, nil
}

func (s *openAIRecoveryPatchStore) PutContinuation(snapshotRef string, cursor int, value any) error {
	if _, err := streamgate.NewStableTokenRequired("snapshotRef", snapshotRef); err != nil {
		return err
	}
	if cursor < 0 {
		return fmt.Errorf("continuation cursor must be non-negative")
	}
	encoded, err := marshalOpenAIPatchValue(value)
	if err != nil {
		return err
	}
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return streamgate.ErrIngressSnapshotClosed
	}
	if _, duplicate := s.continuation[snapshotRef]; duplicate {
		s.mu.Unlock()
		return errOpenAIRecoveryPatchDuplicate
	}
	ingress := s.ingress
	s.mu.Unlock()
	guard, err := ingress.reserveRebuild(int64(len(encoded)))
	if err != nil {
		return err
	}
	entry := &openAIRecoveryPatchEntry{cursor: cursor, value: encoded, guard: guard}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		entry.release()
		return streamgate.ErrIngressSnapshotClosed
	}
	if _, duplicate := s.continuation[snapshotRef]; duplicate {
		entry.release()
		return errOpenAIRecoveryPatchDuplicate
	}
	s.continuation[snapshotRef] = entry
	return nil
}

func openAISchemaPatchKey(schemaRef, patchCode string) string {
	return schemaRef + ":" + patchCode
}

func (s *openAIRecoveryPatchStore) PutSchema(schemaRef, patchCode string, value any) error {
	if _, err := streamgate.NewStableTokenRequired("schemaRef", schemaRef); err != nil {
		return err
	}
	if _, err := streamgate.NewStableTokenRequired("patchCode", patchCode); err != nil {
		return err
	}
	encoded, err := marshalOpenAIPatchValue(value)
	if err != nil {
		return err
	}
	key := openAISchemaPatchKey(schemaRef, patchCode)
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return streamgate.ErrIngressSnapshotClosed
	}
	if _, duplicate := s.schema[key]; duplicate {
		s.mu.Unlock()
		return errOpenAIRecoveryPatchDuplicate
	}
	ingress := s.ingress
	s.mu.Unlock()
	guard, err := ingress.reserveRebuild(int64(len(encoded)))
	if err != nil {
		return err
	}
	entry := &openAIRecoveryPatchEntry{value: encoded, guard: guard}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		entry.release()
		return streamgate.ErrIngressSnapshotClosed
	}
	if _, duplicate := s.schema[key]; duplicate {
		entry.release()
		return errOpenAIRecoveryPatchDuplicate
	}
	s.schema[key] = entry
	return nil
}

func (s *openAIRecoveryPatchStore) takeContinuation(snapshotRef string, cursor int) (*openAIRecoveryPatchEntry, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	entry := s.continuation[snapshotRef]
	if entry == nil || entry.cursor != cursor {
		return nil, fmt.Errorf("OpenAI continuation patch is unavailable")
	}
	delete(s.continuation, snapshotRef)
	return entry, nil
}

func (s *openAIRecoveryPatchStore) takeSchema(schemaRef, patchCode string) (*openAIRecoveryPatchEntry, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	key := openAISchemaPatchKey(schemaRef, patchCode)
	entry := s.schema[key]
	if entry == nil {
		return nil, fmt.Errorf("OpenAI schema patch is unavailable")
	}
	delete(s.schema, key)
	return entry, nil
}

func (s *openAIRecoveryPatchStore) close() {
	if s == nil {
		return
	}
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	continuation := s.continuation
	schema := s.schema
	s.continuation = make(map[string]*openAIRecoveryPatchEntry)
	s.schema = make(map[string]*openAIRecoveryPatchEntry)
	s.ingress = nil
	s.mu.Unlock()
	for _, entry := range continuation {
		entry.release()
	}
	for _, entry := range schema {
		entry.release()
	}
}

type openAIRebuiltLease struct {
	mu        sync.Mutex
	ingress   *openAIIngressSnapshot
	guard     *streamgate.IngressSnapshotRebuildGuard
	rebuilt   *streamgate.IngressSnapshot
	bodyAlias []byte
	released  bool
	once      sync.Once
}

func (l *openAIRebuiltLease) body() ([]byte, error) {
	if l == nil {
		return nil, fmt.Errorf("rebuilt OpenAI request is unavailable")
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.released || l.bodyAlias == nil {
		return nil, fmt.Errorf("rebuilt OpenAI request is unavailable")
	}
	return l.bodyAlias, nil
}

func (l *openAIRebuiltLease) release() {
	if l == nil {
		return
	}
	l.once.Do(func() {
		l.mu.Lock()
		rebuilt := l.rebuilt
		guard := l.guard
		l.bodyAlias = nil
		l.rebuilt = nil
		l.guard = nil
		l.ingress = nil
		l.released = true
		l.mu.Unlock()
		if rebuilt != nil {
			rebuilt.Close()
		}
		if guard != nil {
			guard.Close()
		}
	})
}

func (l *openAIRebuiltLease) isReleased() bool {
	if l == nil {
		return true
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.released
}

type openAIRebuiltRequestStore struct {
	mu     sync.Mutex
	leases map[string]*openAIRebuiltLease
	closed bool
}

func newOpenAIRebuiltRequestStore() *openAIRebuiltRequestStore {
	return &openAIRebuiltRequestStore{leases: make(map[string]*openAIRebuiltLease)}
}

func (s *openAIRebuiltRequestStore) put(ref string, lease *openAIRebuiltLease) error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return streamgate.ErrIngressSnapshotClosed
	}
	if _, exists := s.leases[ref]; exists {
		s.mu.Unlock()
		return fmt.Errorf("duplicate rebuilt OpenAI request reference")
	}
	s.leases[ref] = lease
	s.mu.Unlock()
	return nil
}

func (s *openAIRebuiltRequestStore) take(ref string) (*openAIRebuiltLease, error) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	lease := s.leases[ref]
	if lease == nil {
		s.mu.Unlock()
		return nil, fmt.Errorf("rebuilt OpenAI request reference is unavailable")
	}
	delete(s.leases, ref)
	s.mu.Unlock()
	return lease, nil
}

func (s *openAIRebuiltRequestStore) close() {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	leases := s.leases
	s.leases = make(map[string]*openAIRebuiltLease)
	s.mu.Unlock()
	for _, lease := range leases {
		lease.release()
	}
}

// openAIRequestRebuilder implements the Core raw-free RequestRebuilder seam.
// Its endpoint and patch sources are fixed for one HTTP request.
// Lifecycle is strictly linear: open → in-flight rebuilds → closed.
// Close blocks until every in-flight rebuild has released its patch/lease
// reservation; nil receivers return ErrIngressSnapshotClosed instead of
// panicking.
type openAIRequestRebuilder struct {
	mu           sync.Mutex
	cond         *sync.Cond
	inFlight     int
	closeWaiters int
	closed       bool
	closeDone    bool
	ingress      *openAIIngressSnapshot
	endpoint     string
	patches      *openAIRecoveryPatchStore
	rebuilt      *openAIRebuiltRequestStore
}

func newOpenAIRequestRebuilder(ingress *openAIIngressSnapshot, endpoint string) (*openAIRequestRebuilder, error) {
	switch endpoint {
	case openAIRebuildEndpointChat, openAIRebuildEndpointResponses:
	default:
		return nil, fmt.Errorf("unsupported OpenAI rebuild endpoint")
	}
	if ingress == nil || ingress.isClosed() {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	r := &openAIRequestRebuilder{
		ingress:  ingress,
		endpoint: endpoint,
		patches:  newOpenAIRecoveryPatchStore(ingress),
		rebuilt:  newOpenAIRebuiltRequestStore(),
	}
	r.cond = sync.NewCond(&r.mu)
	return r, nil
}

func (r *openAIRequestRebuilder) PatchStore() *openAIRecoveryPatchStore    { return r.patches }
func (r *openAIRequestRebuilder) RebuiltStore() *openAIRebuiltRequestStore { return r.rebuilt }

func (r *openAIRequestRebuilder) Close() {
	if r == nil {
		return
	}
	r.mu.Lock()
	if r.closeDone {
		r.mu.Unlock()
		return
	}
	r.closeWaiters++
	r.cond.Broadcast()
	defer r.leaveCloseWaiter()
	r.closed = true
	for r.inFlight > 0 && !r.closeDone {
		r.cond.Wait()
	}
	if r.closeDone {
		r.mu.Unlock()
		return
	}
	r.rebuilt.close()
	r.patches.close()
	r.closeDone = true
	r.cond.Broadcast()
	r.mu.Unlock()
}

func (r *openAIRequestRebuilder) RebuildRequest(ctx context.Context, snapshotRef streamgate.RecoveryRequestSnapshotRef, plan streamgate.RecoveryPlan) (streamgate.RebuiltRequestDraft, error) {
	if r == nil {
		return streamgate.RebuiltRequestDraft{}, streamgate.ErrIngressSnapshotClosed
	}
	if err := ctx.Err(); err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return streamgate.RebuiltRequestDraft{}, streamgate.ErrIngressSnapshotClosed
	}
	r.inFlight++
	r.mu.Unlock()
	defer func() {
		r.mu.Lock()
		r.inFlight--
		if r.inFlight == 0 {
			r.cond.Broadcast()
		} else {
			r.cond.Signal()
		}
		r.mu.Unlock()
	}()
	if r.ingress == nil || r.ingress.isClosed() {
		return streamgate.RebuiltRequestDraft{}, streamgate.ErrIngressSnapshotClosed
	}
	if err := snapshotRef.Validate(); err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}
	if err := plan.Validate(); err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}
	if !plan.ReadyForRebuild() {
		return streamgate.RebuiltRequestDraft{}, streamgate.ErrRecoveryPlanNotReady
	}
	currentRef, err := r.ingress.recoveryRef()
	if err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}
	if snapshotRef.SnapshotRef() != currentRef.SnapshotRef() {
		return streamgate.RebuiltRequestDraft{}, fmt.Errorf("OpenAI ingress snapshot reference mismatch")
	}

	directive := plan.Directive()
	var patchEntry *openAIRecoveryPatchEntry
	switch directive.Kind() {
	case streamgate.RecoveryDirectiveKindExact:
		if directive.RequestRef() != snapshotRef.SnapshotRef() {
			return streamgate.RebuiltRequestDraft{}, fmt.Errorf("OpenAI exact request reference mismatch")
		}
	case streamgate.RecoveryDirectiveKindContinuation:
		patchEntry, err = r.patches.takeContinuation(directive.SnapshotRef(), directive.Cursor())
	case streamgate.RecoveryDirectiveKindSchema:
		patchEntry, err = r.patches.takeSchema(directive.SchemaRef(), directive.PatchCode())
	default:
		err = fmt.Errorf("unsupported OpenAI recovery directive")
	}
	if err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}
	if patchEntry != nil {
		defer patchEntry.release()
	}
	if err := ctx.Err(); err != nil {
		return streamgate.RebuiltRequestDraft{}, err
	}

	requestRef := fmt.Sprintf("openai.rebuilt.%d", openAIRebuiltSequence.Add(1))
	lease := &openAIRebuiltLease{ingress: r.ingress}
	retained, peak, maxBytes := currentRef.RetainedBytes(), currentRef.PeakBytes(), currentRef.MaxBytes()
	rebuiltPromptTokens := 0

	if patchEntry == nil {
		body, err := r.ingress.canonicalBody()
		if err != nil {
			return streamgate.RebuiltRequestDraft{}, err
		}
		lease.bodyAlias = body
		rebuiltPromptTokens = estimateInputTokensBytes(body, nil, nil, nil)
	} else {
		field := "messages"
		if r.endpoint == openAIRebuildEndpointResponses {
			field = "input"
		}
		body, err := r.ingress.canonicalBody()
		if err != nil {
			return streamgate.RebuiltRequestDraft{}, err
		}
		patchPlan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: field, value: patchEntry.value}})
		if err != nil {
			return streamgate.RebuiltRequestDraft{}, err
		}
		guard, err := r.ingress.reserveRebuild(int64(patchPlan.outputSize))
		if err != nil {
			return streamgate.RebuiltRequestDraft{}, err
		}
		output := patchPlan.apply()
		rebuilt, err := guard.CommitOwnedTyped(openAIRebuiltBodyViewName, output)
		if err != nil {
			guard.Close()
			return streamgate.RebuiltRequestDraft{}, err
		}
		lease.guard = guard
		lease.rebuilt = rebuilt
		lease.bodyAlias = output
		accessor := rebuilt.Accessor()
		retained = uint64(accessor.RetainedBytes())
		maxBytes = uint64(accessor.MaxBytes())
		peak = retained
		if current, refErr := r.ingress.recoveryRef(); refErr == nil && current.PeakBytes() > peak {
			peak = current.PeakBytes()
		}
		rebuiltPromptTokens = estimateInputTokensBytes(output, nil, nil, nil)
	}

	draft, err := streamgate.NewRebuiltRequestDraftWithIdempotency(
		plan.PlanID(), plan.IdempotencyKey(), requestRef, r.endpoint, openAIRebuildFamily,
		retained, peak, maxBytes, rebuiltPromptTokens, plan.RequiredCapabilities(),
	)
	if err != nil {
		lease.release()
		return streamgate.RebuiltRequestDraft{}, err
	}
	if err := r.rebuilt.put(requestRef, lease); err != nil {
		lease.release()
		return streamgate.RebuiltRequestDraft{}, err
	}
	return draft, nil
}

func (r *openAIRequestRebuilder) leaveCloseWaiter() {
	r.mu.Lock()
	r.closeWaiters--
	r.cond.Broadcast()
	r.mu.Unlock()
}

var _ streamgate.RequestRebuilder = (*openAIRequestRebuilder)(nil)

func isOpenAIRebuildOverflow(err error) bool {
	return errors.Is(err, streamgate.ErrIngressSnapshotRebuildOverflow) ||
		errors.Is(err, streamgate.ErrRecoverySnapshotLimitExceeded)
}