package openai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	edgeservice "iop/apps/edge/internal/service"
)

type singleRequestChatPolicy struct {
	status       int
	finishReason string
	errorType    string
	message      string
	silent       bool
	error        bool
}

func singleRequestChatTerminalPolicy(disposition edgeservice.SingleRequestTerminalDisposition) singleRequestChatPolicy {
	if disposition.Kind == "" {
		disposition.Kind = edgeservice.SingleRequestTerminalEndTurn
	}
	if disposition.Validate() != nil {
		disposition = edgeservice.SingleRequestTerminalDisposition{
			Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider,
		}
	}
	switch disposition.Kind {
	case edgeservice.SingleRequestTerminalEndTurn:
		return singleRequestChatPolicy{status: http.StatusOK, finishReason: "stop"}
	case edgeservice.SingleRequestTerminalLength:
		return singleRequestChatPolicy{status: http.StatusOK, finishReason: "length"}
	case edgeservice.SingleRequestTerminalCancelled:
		return singleRequestChatPolicy{silent: true}
	case edgeservice.SingleRequestTerminalError:
		switch disposition.ErrorClass {
		case edgeservice.SingleRequestTerminalErrorValidation, edgeservice.SingleRequestTerminalErrorContext:
			return singleRequestChatPolicy{status: http.StatusBadRequest, errorType: "invalid_request_error", message: "single-request execution was rejected", error: true}
		case edgeservice.SingleRequestTerminalErrorTimeout:
			return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution timed out", error: true}
		default:
			return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution failed", error: true}
		}
	default:
		return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution failed", error: true}
	}
}

func (s *Server) handleChatSingleRequestStream(
	w http.ResponseWriter,
	r *http.Request,
	capability singleRequestService,
	dispatch routeDispatch,
	body []byte,
) {
	requestID, err := newLogicalRequestRandomID()
	if err != nil {
		writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
		return
	}
	requestID = "req_" + requestID
	stream, err := newSingleRequestChatStream(w, requestID, dispatch.SingleRequest.PublicModel)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "run_error", "single-request streaming is unavailable")
		return
	}
	execution, err := capability.StartSingleRequest(r.Context(), edgeservice.SingleRequestRequest{
		RequestID: requestID,
		Binding:   dispatch.SingleRequest.Clone(),
		Prompt:    string(append([]byte(nil), body...)),
	})
	if err != nil || execution == nil {
		if errors.Is(err, edgeservice.ErrSingleRequestExecutorUnavailable) {
			writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
			return
		}
		writeError(w, http.StatusBadGateway, "run_error", "single-request execution could not be started")
		return
	}
	defer execution.Cancel()
	_ = pumpSingleRequestChatStream(r.Context(), execution, stream, newWallClockSingleRequestAnthropicTicker)
}

func (s *Server) handleChatSingleRequest(
	w http.ResponseWriter,
	r *http.Request,
	capability singleRequestService,
	dispatch routeDispatch,
	body []byte,
) {
	requestID, err := newLogicalRequestRandomID()
	if err != nil {
		writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
		return
	}
	requestID = "req_" + requestID
	execution, err := capability.StartSingleRequest(r.Context(), edgeservice.SingleRequestRequest{
		RequestID: requestID,
		Binding:   dispatch.SingleRequest.Clone(),
		Prompt:    string(append([]byte(nil), body...)),
	})
	if err != nil || execution == nil {
		if errors.Is(err, edgeservice.ErrSingleRequestExecutorUnavailable) {
			writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
			return
		}
		writeError(w, http.StatusBadGateway, "run_error", "single-request execution could not be started")
		return
	}
	defer execution.Cancel()
	for {
		select {
		case <-r.Context().Done():
			execution.Cancel()
			return
		case progress, ok := <-execution.Progress():
			if !ok {
				if r.Context().Err() == nil && execution.State() != edgeservice.SingleRequestStateCancelled {
					writeError(w, http.StatusBadGateway, "run_error", "single-request execution failed")
				}
				return
			}
			switch progress.Stage {
			case edgeservice.SingleRequestStateFinalizing:
				if progress.Result == nil {
					_ = execution.AcknowledgeTerminal(false)
					writeError(w, http.StatusBadGateway, "run_error", "single-request execution failed")
					return
				}
				writeErr := writeChatSingleRequestTerminal(w, requestID, dispatch.SingleRequest.PublicModel, *progress.Result)
				_ = execution.AcknowledgeTerminal(writeErr == nil)
				return
			case edgeservice.SingleRequestStateFailed:
				policy := singleRequestChatTerminalPolicy(singleRequestProgressTerminal(progress, edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}))
				writeError(w, policy.status, policy.errorType, policy.message)
				return
			case edgeservice.SingleRequestStateCancelled:
				return
			}
		}
	}
}

func writeChatSingleRequestTerminal(w http.ResponseWriter, requestID, model string, result edgeservice.SingleRequestResult) error {
	policy := singleRequestChatTerminalPolicy(result.Terminal)
	if policy.silent || policy.error || policy.finishReason == "" {
		return errors.New("single-request result has no Chat terminal")
	}
	content := result.Output
	if policy.finishReason == "length" {
		content = ""
	}
	response := chatCompletionResponse{
		ID: "chatcmpl_iop_" + strings.TrimPrefix(requestID, "req_"), Object: "chat.completion",
		Created: time.Now().Unix(), Model: model,
		Choices: []chatCompletionChoice{{Index: 0, Message: chatMessage{Role: "assistant", Content: content}, FinishReason: policy.finishReason}},
	}
	encoded, err := json.Marshal(response)
	if err != nil {
		return err
	}
	encoded = append(encoded, '\n')
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(policy.status)
	n, err := w.Write(encoded)
	if err != nil {
		return err
	}
	if n != len(encoded) {
		return io.ErrShortWrite
	}
	return nil
}

type singleRequestChatStream struct {
	mu          sync.Mutex
	w           http.ResponseWriter
	flusher     http.Flusher
	id          string
	model       string
	created     int64
	started     bool
	terminal    bool
	terminalErr error
}

func newSingleRequestChatStream(w http.ResponseWriter, requestID, model string) (*singleRequestChatStream, error) {
	flusher, ok := w.(http.Flusher)
	if !ok || strings.TrimSpace(requestID) == "" || strings.TrimSpace(model) == "" {
		return nil, errors.New("single-request Chat stream is unavailable")
	}
	return &singleRequestChatStream{w: w, flusher: flusher, id: "chatcmpl_iop_" + strings.TrimPrefix(requestID, "req_"), model: model, created: time.Now().Unix()}, nil
}

func (s *singleRequestChatStream) startLocked() error {
	if s.started {
		return nil
	}
	s.w.Header().Set("Content-Type", "text/event-stream")
	s.w.Header().Set("Cache-Control", "no-cache")
	s.w.WriteHeader(http.StatusOK)
	s.flusher.Flush()
	s.started = true
	return nil
}

func (s *singleRequestChatStream) Start() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	return s.startLocked()
}

func (s *singleRequestChatStream) Ping() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	_, err := io.WriteString(s.w, ": ping\n\n")
	if err == nil {
		s.flusher.Flush()
	}
	return err
}

func (s *singleRequestChatStream) writeSSELocked(value any) error {
	payload, err := json.Marshal(value)
	if err != nil {
		return err
	}
	if _, err = fmt.Fprintf(s.w, "data: %s\n\n", payload); err != nil {
		return err
	}
	s.flusher.Flush()
	return nil
}

func (s *singleRequestChatStream) Final(result edgeservice.SingleRequestResult) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	policy := singleRequestChatTerminalPolicy(result.Terminal)
	if policy.silent || policy.error || policy.finishReason == "" {
		return errors.New("single-request result has no Chat stream terminal")
	}
	s.terminal = true
	if policy.finishReason != "length" && result.Output != "" {
		if err := s.writeSSELocked(chatCompletionChunk{ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model, Choices: []chatCompletionChunkChoice{{Index: 0, Delta: chatDelta{Content: result.Output}}}}); err != nil {
			s.terminalErr = err
			return err
		}
	}
	if err := s.writeSSELocked(chatCompletionChunk{ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model, Choices: []chatCompletionChunkChoice{{Index: 0, Delta: chatDelta{}, FinishReason: policy.finishReason}}}); err != nil {
		s.terminalErr = err
		return err
	}
	if _, err := io.WriteString(s.w, "data: [DONE]\n\n"); err != nil {
		s.terminalErr = err
		return err
	}
	s.flusher.Flush()
	return nil
}

func (s *singleRequestChatStream) TerminalError(disposition edgeservice.SingleRequestTerminalDisposition) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	policy := singleRequestChatTerminalPolicy(disposition)
	if policy.silent {
		s.terminal = true
		return nil
	}
	if !policy.error {
		policy = singleRequestChatTerminalPolicy(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	s.terminal = true
	if err := s.writeSSELocked(errorResponse{Error: errorBody{Type: policy.errorType, Message: policy.message}}); err != nil {
		s.terminalErr = err
		return err
	}
	if _, err := io.WriteString(s.w, "data: [DONE]\n\n"); err != nil {
		s.terminalErr = err
		return err
	}
	s.flusher.Flush()
	return nil
}

func pumpSingleRequestChatStream(ctx context.Context, execution edgeservice.SingleRequestExecution, stream *singleRequestChatStream, tickerFactory singleRequestAnthropicTickerFactory) error {
	if execution == nil || stream == nil || tickerFactory == nil {
		return errors.New("single-request Chat stream is unavailable")
	}
	if err := ctx.Err(); err != nil {
		execution.Cancel()
		return err
	}
	if err := stream.Start(); err != nil {
		execution.Cancel()
		return err
	}
	ticker := tickerFactory()
	if ticker == nil {
		execution.Cancel()
		return errors.New("single-request Chat stream is unavailable")
	}
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			execution.Cancel()
			return ctx.Err()
		case <-ticker.Ticks():
			if err := stream.Ping(); err != nil {
				execution.Cancel()
				return err
			}
		case progress, ok := <-execution.Progress():
			if !ok {
				if ctx.Err() != nil {
					return ctx.Err()
				}
				if execution.State() == edgeservice.SingleRequestStateCompleted {
					return nil
				}
				return stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
			}
			switch progress.Stage {
			case edgeservice.SingleRequestStateFinalizing:
				if progress.Result == nil {
					err := stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
					return errors.Join(err, execution.AcknowledgeTerminal(false))
				}
				err := stream.Final(*progress.Result)
				return errors.Join(err, execution.AcknowledgeTerminal(err == nil))
			case edgeservice.SingleRequestStateFailed:
				return stream.TerminalError(singleRequestProgressTerminal(progress, edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}))
			case edgeservice.SingleRequestStateCancelled:
				return stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalCancelled})
			}
		}
	}
}