package acceptanceemulator

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"hash/crc32"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
	_ "golang.org/x/image/webp"
)

const maxProtocolBodyBytes = 128 << 20

type Config struct {
	Now        func() time.Time
	HTTPClient *http.Client
	Wait       func(context.Context, time.Duration) error
}

type Server struct {
	now        func() time.Time
	httpClient *http.Client
	wait       func(context.Context, time.Duration) error
	mu         sync.RWMutex
	tasks      map[string]videoTask
	fixtures   map[string]fixture
	nextID     atomic.Uint64
	report     Report
	callbacks  map[string]map[int64]int
	pngSmall   string
	pngLarge   string
	pngPeak    string
}

type Report struct {
	GeminiRequests       int64            `json:"geminiRequests"`
	GeminiInvalid        int64            `json:"geminiInvalid"`
	GeminiInputBytes     int64            `json:"geminiInputBytes"`
	GeminiOutputBytes    int64            `json:"geminiOutputBytes"`
	VideoSubmissions     int64            `json:"videoSubmissions"`
	VideoPolls           int64            `json:"videoPolls"`
	VideoInvalid         int64            `json:"videoInvalid"`
	VideoReferenceCounts map[string]int64 `json:"videoReferenceCounts"`
	VideoRoleCounts      map[string]int64 `json:"videoRoleCounts"`
	UniqueImageHashes    int              `json:"uniqueImageHashes"`
	UniqueVideoTasks     int              `json:"uniqueVideoTasks"`
	ForcedConversions    int64            `json:"forcedConversionRequests"`
	VerifiedConversions  int64            `json:"verifiedConversions"`
	CallbackEvents       int64            `json:"callbackEvents"`
	DuplicateCallbacks   int64            `json:"duplicateCallbacks"`
}

type videoTask struct {
	ID        string
	Model     string
	CreatedAt time.Time
	ReadyAt   time.Time
	ImageRefs []imageReference
}

type imageReference struct {
	Role        string
	SHA256      string
	ContentType string
	Width       int
	Height      int
}

type fixture struct {
	ContentType string
	Payload     []byte
}

func New(config Config) *Server {
	now := config.Now
	if now == nil {
		now = time.Now
	}
	client := config.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: 30 * time.Second}
	}
	waiter := config.Wait
	if waiter == nil {
		waiter = wait
	}
	return &Server{
		now: now, httpClient: client, wait: waiter, tasks: map[string]videoTask{}, fixtures: buildFixtures(),
		callbacks: map[string]map[int64]int{},
		report: Report{
			VideoReferenceCounts: map[string]int64{},
			VideoRoleCounts:      map[string]int64{},
		},
		pngSmall: base64.StdEncoding.EncodeToString(paddedPNG(256 << 10)),
		pngLarge: base64.StdEncoding.EncodeToString(paddedPNG(4 << 20)),
		pngPeak:  base64.StdEncoding.EncodeToString(paddedPNG(8 << 20)),
	}
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", s.health)
	mux.HandleFunc("GET /report", s.getReport)
	mux.HandleFunc("POST /v1beta/models/", s.geminiGenerateContent)
	mux.HandleFunc("POST /v1/models/", s.geminiGenerateContent)
	mux.HandleFunc("POST /contents/generations/tasks", s.submitVideo)
	mux.HandleFunc("GET /contents/generations/tasks/{taskID}", s.getVideo)
	mux.HandleFunc("DELETE /contents/generations/tasks/{taskID}", s.deleteVideo)
	mux.HandleFunc("GET /media/{asset}", s.getMedia)
	mux.HandleFunc("GET /fixtures/{asset}", s.getFixture)
	mux.HandleFunc("POST /callbacks", s.collectCallback)
	return mux
}

func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
	if !strings.HasSuffix(r.URL.Path, ":generateContent") {
		http.NotFound(w, r)
		return
	}
	var body map[string]any
	if err := decodeJSON(r, &body); err != nil {
		s.recordGeminiInvalid()
		writeProtocolError(w, http.StatusBadRequest, err.Error())
		return
	}
	inputBytes, err := validateGeminiImageRequest(body)
	if err != nil {
		s.recordGeminiInvalid()
		writeProtocolError(w, http.StatusBadRequest, err.Error())
		return
	}
	outputBytes, outputBase64, delay := s.geminiProfile(inputBytes)
	if err := s.wait(r.Context(), delay); err != nil {
		return
	}
	s.mu.Lock()
	s.report.GeminiRequests++
	s.report.GeminiInputBytes += int64(inputBytes)
	s.report.GeminiOutputBytes += int64(outputBytes)
	s.mu.Unlock()
	w.Header().Set("X-Request-ID", "acceptance-gemini-"+strconv.FormatUint(s.nextID.Add(1), 10))
	writeJSON(w, http.StatusOK, map[string]any{
		"candidates": []any{map[string]any{
			"index": 0,
			"content": map[string]any{
				"role": "model",
				"parts": []any{map[string]any{"inlineData": map[string]any{
					"mimeType": "image/png",
					"data":     outputBase64,
				}}},
			},
			"finishReason": "STOP",
		}},
		"usageMetadata": map[string]any{
			"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2,
		},
	})
}

func (s *Server) geminiProfile(inputBytes int) (int, string, time.Duration) {
	switch {
	case inputBytes > acceptanceworkload.GeminiLarge.InputBytes:
		profile := acceptanceworkload.GeminiPeak
		return profile.OutputBytes, s.pngPeak, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
	case inputBytes > acceptanceworkload.GeminiBaseline.InputBytes:
		profile := acceptanceworkload.GeminiLarge
		return profile.OutputBytes, s.pngLarge, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
	default:
		profile := acceptanceworkload.GeminiBaseline
		return profile.OutputBytes, s.pngSmall, profile.DelayMin
	}
}

func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
	var body map[string]any
	if err := decodeJSON(r, &body); err != nil {
		s.recordVideoInvalid()
		writeProtocolError(w, http.StatusBadRequest, err.Error())
		return
	}
	refs, longRun, forceConversion, err := s.validateVideoRequest(r.Context(), body)
	if err != nil {
		s.recordVideoInvalid()
		writeProtocolError(w, http.StatusBadRequest, err.Error())
		return
	}
	delay := videoDelay(body, longRun)
	id := "acceptance-video-" + strconv.FormatUint(s.nextID.Add(1), 10)
	task := videoTask{
		ID: id, Model: strings.TrimSpace(stringValue(body["model"])),
		CreatedAt: s.now(), ReadyAt: s.now().Add(delay), ImageRefs: refs,
	}
	s.mu.Lock()
	s.tasks[id] = task
	s.report.VideoSubmissions++
	s.report.VideoReferenceCounts[strconv.Itoa(len(refs))]++
	if forceConversion {
		s.report.ForcedConversions++
		s.report.VerifiedConversions++
	}
	for _, ref := range refs {
		s.report.VideoRoleCounts[ref.Role]++
	}
	s.mu.Unlock()
	w.Header().Set("X-Request-ID", id)
	writeJSON(w, http.StatusOK, map[string]any{
		"id": id, "model": task.Model, "status": "queued",
		"created_at": task.CreatedAt.Unix(),
	})
}

func (s *Server) getVideo(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimSpace(r.PathValue("taskID"))
	s.mu.Lock()
	task, ok := s.tasks[id]
	if ok {
		s.report.VideoPolls++
	}
	s.mu.Unlock()
	if !ok {
		writeProtocolError(w, http.StatusNotFound, "video task not found")
		return
	}
	status := "running"
	content := map[string]any{}
	if !s.now().Before(task.ReadyAt) {
		status = "succeeded"
		content["video_url"] = requestOrigin(r) + "/media/" + url.PathEscape(id) + ".mp4"
	}
	w.Header().Set("X-Request-ID", id)
	writeJSON(w, http.StatusOK, map[string]any{
		"id": id, "model": task.Model, "status": status, "content": content,
		"created_at": task.CreatedAt.Unix(),
		"usage":      map[string]any{"completion_tokens": 1, "total_tokens": 1},
	})
}

func (s *Server) deleteVideo(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimSpace(r.PathValue("taskID"))
	s.mu.Lock()
	_, ok := s.tasks[id]
	delete(s.tasks, id)
	s.mu.Unlock()
	if !ok {
		writeProtocolError(w, http.StatusNotFound, "video task not found")
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"id": id, "status": "cancelled"})
}

func (s *Server) getMedia(w http.ResponseWriter, r *http.Request) {
	asset := strings.TrimSuffix(strings.TrimSpace(r.PathValue("asset")), ".mp4")
	s.mu.RLock()
	_, ok := s.tasks[asset]
	s.mu.RUnlock()
	if !ok {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "video/mp4")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(minimalMP4)
}

func (s *Server) getFixture(w http.ResponseWriter, r *http.Request) {
	fixture, ok := s.fixtures[strings.TrimSpace(r.PathValue("asset"))]
	if !ok {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", fixture.ContentType)
	w.Header().Set("Cache-Control", "public, max-age=3600")
	w.Header().Set("X-Content-SHA256", payloadSHA256(fixture.Payload))
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(fixture.Payload)
}

func (s *Server) getReport(w http.ResponseWriter, _ *http.Request) {
	s.mu.RLock()
	report := s.report
	tasks := make([]videoTask, 0, len(s.tasks))
	for _, task := range s.tasks {
		tasks = append(tasks, task)
	}
	s.mu.RUnlock()
	hashes := map[string]struct{}{}
	for _, task := range tasks {
		for _, ref := range task.ImageRefs {
			hashes[ref.SHA256] = struct{}{}
		}
	}
	report.UniqueImageHashes = len(hashes)
	report.UniqueVideoTasks = len(tasks)
	writeJSON(w, http.StatusOK, report)
}

func (s *Server) collectCallback(w http.ResponseWriter, r *http.Request) {
	var payload map[string]any
	if err := decodeJSON(r, &payload); err != nil {
		writeProtocolError(w, http.StatusBadRequest, err.Error())
		return
	}
	taskID := strings.TrimSpace(stringValue(payload["taskId"]))
	seq := int64(0)
	switch value := payload["seq"].(type) {
	case json.Number:
		seq, _ = value.Int64()
	case float64:
		seq = int64(value)
	}
	if taskID == "" || seq <= 0 {
		writeProtocolError(w, http.StatusBadRequest, "callback taskId and seq are required")
		return
	}
	s.mu.Lock()
	seen := s.callbacks[taskID]
	if seen == nil {
		seen = map[int64]int{}
		s.callbacks[taskID] = seen
	}
	seen[seq]++
	s.report.CallbackEvents++
	if seen[seq] > 1 {
		s.report.DuplicateCallbacks++
	}
	s.mu.Unlock()
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

func (s *Server) validateVideoRequest(ctx context.Context, body map[string]any) ([]imageReference, bool, bool, error) {
	content, _ := body["content"].([]any)
	refs := make([]imageReference, 0, 9)
	firstFrames := 0
	lastFrames := 0
	longRun := false
	forceConversion := false
	for _, raw := range content {
		item, _ := raw.(map[string]any)
		switch strings.TrimSpace(stringValue(item["type"])) {
		case "text":
			if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-long-recovery") {
				longRun = true
			}
			if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-force-conversion") {
				forceConversion = true
			}
		case "image_url":
			role := strings.TrimSpace(stringValue(item["role"]))
			switch role {
			case "reference_image":
			case "first_frame":
				firstFrames++
			case "last_frame":
				lastFrames++
			default:
				return nil, false, false, fmt.Errorf("unsupported image role %q", role)
			}
			nested, _ := item["image_url"].(map[string]any)
			rawURL := strings.TrimSpace(stringValue(nested["url"]))
			ref, err := s.fetchImageReference(ctx, role, rawURL)
			if err != nil {
				return nil, false, false, err
			}
			refs = append(refs, ref)
		}
	}
	if len(refs) != 3 && len(refs) != 6 && len(refs) != 9 {
		return nil, false, false, fmt.Errorf("expected 3, 6, or 9 image references, got %d", len(refs))
	}
	if firstFrames > 1 || lastFrames > 1 {
		return nil, false, false, errors.New("first_frame and last_frame must be unique")
	}
	if forceConversion {
		for _, ref := range refs {
			if ref.Width >= 3840 || ref.Height >= 2160 {
				return nil, false, false, errors.New("forced 4K image did not pass through automatic normalization")
			}
		}
	}
	return refs, longRun, forceConversion, nil
}

func (s *Server) fetchImageReference(ctx context.Context, role string, rawURL string) (imageReference, error) {
	if rawURL == "" {
		return imageReference{}, errors.New("image reference URL is required")
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return imageReference{}, err
	}
	response, err := s.httpClient.Do(request)
	if err != nil {
		return imageReference{}, fmt.Errorf("fetch image reference: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return imageReference{}, fmt.Errorf("fetch image reference: status %d", response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, 32<<20))
	if err != nil {
		return imageReference{}, err
	}
	if len(payload) == 0 {
		return imageReference{}, errors.New("image reference is empty")
	}
	config, format, err := image.DecodeConfig(bytes.NewReader(payload))
	if err != nil {
		return imageReference{}, fmt.Errorf("decode image reference: %w", err)
	}
	sum := sha256.Sum256(payload)
	return imageReference{
		Role: role, SHA256: hex.EncodeToString(sum[:]), ContentType: "image/" + format,
		Width: config.Width, Height: config.Height,
	}, nil
}

func validateGeminiImageRequest(body map[string]any) (int, error) {
	generationConfig, _ := body["generationConfig"].(map[string]any)
	modalities, _ := generationConfig["responseModalities"].([]any)
	hasImageModality := false
	for _, modality := range modalities {
		if strings.EqualFold(strings.TrimSpace(stringValue(modality)), "IMAGE") {
			hasImageModality = true
		}
	}
	if !hasImageModality {
		return 0, errors.New("generationConfig.responseModalities must contain IMAGE")
	}
	total := 0
	contents, _ := body["contents"].([]any)
	for _, rawContent := range contents {
		content, _ := rawContent.(map[string]any)
		parts, _ := content["parts"].([]any)
		for _, rawPart := range parts {
			part, _ := rawPart.(map[string]any)
			inline, _ := part["inlineData"].(map[string]any)
			if inline == nil {
				inline, _ = part["inline_data"].(map[string]any)
			}
			encoded := strings.TrimSpace(stringValue(inline["data"]))
			if encoded == "" {
				continue
			}
			decoded, err := base64.StdEncoding.DecodeString(encoded)
			if err != nil {
				return 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err)
			}
			if len(decoded) == 0 {
				return 0, errors.New("Gemini inlineData is empty")
			}
			total += len(decoded)
		}
	}
	if total == 0 {
		return 0, errors.New("Gemini request has no inlineData image")
	}
	return total, nil
}

func (s *Server) recordGeminiInvalid() {
	s.mu.Lock()
	s.report.GeminiInvalid++
	s.mu.Unlock()
}

func (s *Server) recordVideoInvalid() {
	s.mu.Lock()
	s.report.VideoInvalid++
	s.mu.Unlock()
}

func decodeJSON(r *http.Request, target any) error {
	reader := io.LimitReader(r.Body, maxProtocolBodyBytes+1)
	payload, err := io.ReadAll(reader)
	if err != nil {
		return err
	}
	if len(payload) > maxProtocolBodyBytes {
		return errors.New("protocol request body is too large")
	}
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	return decoder.Decode(target)
}

func writeJSON(w http.ResponseWriter, status int, value any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(value)
}

func writeProtocolError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]any{"error": map[string]any{
		"code": "acceptance_protocol_invalid", "message": message,
	}})
}

func wait(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func scaledDelay(minimum time.Duration, maximum time.Duration, seed int) time.Duration {
	if maximum <= minimum {
		return minimum
	}
	windowSeconds := int64((maximum - minimum) / time.Second)
	if windowSeconds <= 0 {
		return minimum
	}
	return minimum + time.Duration(int64(seed)%(windowSeconds+1))*time.Second
}

func videoDelay(body map[string]any, longRun bool) time.Duration {
	seed := int64(0)
	switch value := body["seed"].(type) {
	case json.Number:
		seed, _ = value.Int64()
	case float64:
		seed = int64(value)
	case int:
		seed = int64(value)
	}
	if seed < 0 {
		seed = -seed
	}
	if longRun {
		return 2*time.Minute + time.Duration(seed%61)*time.Second
	}
	return 5*time.Second + time.Duration(seed%11)*time.Second
}

func requestOrigin(r *http.Request) string {
	scheme := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
	if scheme == "" {
		scheme = "http"
	}
	host := strings.TrimSpace(r.Host)
	return scheme + "://" + host
}

func stringValue(value any) string {
	switch typed := value.(type) {
	case nil:
		return ""
	case string:
		return typed
	case json.Number:
		return typed.String()
	default:
		return fmt.Sprint(value)
	}
}

func paddedPNG(size int) []byte {
	base, _ := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")
	if size <= len(base) {
		return base
	}
	iendOffset := len(base) - 12
	paddingLength := size - len(base) - 12
	if paddingLength < 0 {
		paddingLength = 0
	}
	chunk := make([]byte, 12+paddingLength)
	binary.BigEndian.PutUint32(chunk[:4], uint32(paddingLength))
	copy(chunk[4:8], []byte("teST"))
	binary.BigEndian.PutUint32(chunk[8+paddingLength:], crc32.ChecksumIEEE(chunk[4:8+paddingLength]))
	out := make([]byte, 0, len(base)+len(chunk))
	out = append(out, base[:iendOffset]...)
	out = append(out, chunk...)
	out = append(out, base[iendOffset:]...)
	return out
}

func buildFixtures() map[string]fixture {
	fixtures := map[string]fixture{}
	for index := 0; index < 4; index++ {
		imageValue := patternedImage(768, 512, index+1)
		var encoded bytes.Buffer
		_ = png.Encode(&encoded, imageValue)
		fixtures[fmt.Sprintf("image-%02d.png", index)] = fixture{ContentType: "image/png", Payload: encoded.Bytes()}
	}
	for index := 0; index < 4; index++ {
		imageValue := patternedImage(1024, 768, index+11)
		var encoded bytes.Buffer
		_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 88})
		fixtures[fmt.Sprintf("image-%02d.jpg", index+4)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
	}
	webpPayload, _ := base64.StdEncoding.DecodeString("UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEALmk0mk0iIiIiIgBoSygABc6zbAAA")
	for index := 0; index < 4; index++ {
		payload := append([]byte(nil), webpPayload...)
		payload = append(payload, byte(index))
		fixtures[fmt.Sprintf("image-%02d.webp", index+8)] = fixture{ContentType: "image/webp", Payload: payload}
	}
	for index := 0; index < 4; index++ {
		imageValue := patternedImage(4096, 2160, index+21)
		var encoded bytes.Buffer
		_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 90})
		fixtures[fmt.Sprintf("image-%02d-4k.jpg", index+12)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
	}
	return fixtures
}

func patternedImage(width int, height int, seed int) image.Image {
	out := image.NewNRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			out.SetNRGBA(x, y, color.NRGBA{
				R: uint8((x + seed*31) % 256),
				G: uint8((y + seed*47) % 256),
				B: uint8((x/8 + y/8 + seed*59) % 256),
				A: 255,
			})
		}
	}
	return out
}

func payloadSHA256(payload []byte) string {
	sum := sha256.Sum256(payload)
	return hex.EncodeToString(sum[:])
}

var minimalMP4 = []byte{
	0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm',
	0, 0, 0, 0, 'i', 's', 'o', 'm', 'm', 'p', '4', '2',
	0, 0, 0, 8, 'm', 'd', 'a', 't',
}