package multimodal

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"net/url"
	"strings"
	"time"
)

var ErrEmptyResult = errors.New("multimodal provider returned an empty result")

type SpeechClient struct {
	BaseURL string
	Client  *http.Client
}

type VisionClient struct {
	BaseURL string
	Client  *http.Client
}

type ImageGenerationClient struct {
	BaseURL string
	Client  *http.Client
}

type VideoUnderstandingClient struct {
	BaseURL string
	Client  *http.Client
}

type GeneratedImage struct {
	DataBase64  string
	SourceURL   string
	ContentType string
}

type ImageGenerator interface {
	Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error)
}

// ComfyImageGenerationClient is the replaceable transport for a ComfyUI
// image-generation skill. The workflow is deliberately kept in this adapter:
// the core only selects a capability and supplies the user's prompt.
type ComfyImageGenerationClient struct {
	BaseURL      string
	Client       *http.Client
	PollInterval time.Duration
	PollTimeout  time.Duration
}

func (c SpeechClient) Transcribe(ctx context.Context, filename, contentType string, data []byte) (string, error) {
	if len(data) == 0 {
		return "", errors.New("audio payload is empty")
	}
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	header := make(textproto.MIMEHeader)
	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, strings.ReplaceAll(filename, `"`, "")))
	header.Set("Content-Type", contentType)
	part, err := writer.CreatePart(header)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(data); err != nil {
		return "", err
	}
	if err := writer.WriteField("response_format", "json"); err != nil {
		return "", err
	}
	if err := writer.Close(); err != nil {
		return "", err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/inference", body)
	if err != nil {
		return "", err
	}
	request.Header.Set("Content-Type", writer.FormDataContentType())
	response, err := c.client().Do(request)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return "", fmt.Errorf("speech provider status %d", response.StatusCode)
	}
	var result struct {
		Text string `json:"text"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&result); err != nil {
		return "", err
	}
	if strings.TrimSpace(result.Text) == "" {
		return "", ErrEmptyResult
	}
	return strings.TrimSpace(result.Text), nil
}

func (c VisionClient) Describe(ctx context.Context, model, contentType string, data []byte, instruction string) (string, error) {
	if len(data) == 0 {
		return "", errors.New("image payload is empty")
	}
	payload := struct {
		Model    string `json:"model"`
		Messages []struct {
			Role    string `json:"role"`
			Content []struct {
				Type     string `json:"type"`
				Text     string `json:"text,omitempty"`
				ImageURL *struct {
					URL string `json:"url"`
				} `json:"image_url,omitempty"`
			} `json:"content"`
		} `json:"messages"`
		MaxTokens int `json:"max_tokens"`
	}{}
	payload.Model = model
	payload.MaxTokens = 1024
	payload.Messages = append(payload.Messages, struct {
		Role    string `json:"role"`
		Content []struct {
			Type     string `json:"type"`
			Text     string `json:"text,omitempty"`
			ImageURL *struct {
				URL string `json:"url"`
			} `json:"image_url,omitempty"`
		} `json:"content"`
	}{Role: "user", Content: []struct {
		Type     string `json:"type"`
		Text     string `json:"text,omitempty"`
		ImageURL *struct {
			URL string `json:"url"`
		} `json:"image_url,omitempty"`
	}{
		{Type: "text", Text: instruction},
		{Type: "image_url", ImageURL: &struct {
			URL string `json:"url"`
		}{URL: "data:" + contentType + ";base64," + encodeBase64(data)},
		}}})
	encoded, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/chat/completions", bytes.NewReader(encoded))
	if err != nil {
		return "", err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := c.client().Do(request)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return "", fmt.Errorf("vision provider status %d", response.StatusCode)
	}
	var result struct {
		Choices []struct {
			Message struct {
				Content          string `json:"content"`
				ReasoningContent string `json:"reasoning_content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&result); err != nil {
		return "", err
	}
	if len(result.Choices) == 0 {
		return "", ErrEmptyResult
	}
	text := strings.TrimSpace(result.Choices[0].Message.Content)
	if text == "" {
		text = strings.TrimSpace(result.Choices[0].Message.ReasoningContent)
	}
	if text == "" {
		return "", ErrEmptyResult
	}
	return text, nil
}

// Generate uses an OpenAI-compatible image-generation skill. The provider is
// selected outside the core by the capability binding; this adapter contains
// transport only and never decides what the user meant.
func (c ImageGenerationClient) Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error) {
	if strings.TrimSpace(prompt) == "" {
		return GeneratedImage{}, errors.New("image prompt is empty")
	}
	payload := struct {
		Model          string `json:"model"`
		Prompt         string `json:"prompt"`
		Size           string `json:"size,omitempty"`
		ResponseFormat string `json:"response_format,omitempty"`
	}{Model: model, Prompt: prompt, Size: size, ResponseFormat: "b64_json"}
	encoded, err := json.Marshal(payload)
	if err != nil {
		return GeneratedImage{}, err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/images/generations", bytes.NewReader(encoded))
	if err != nil {
		return GeneratedImage{}, err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := c.client().Do(request)
	if err != nil {
		return GeneratedImage{}, err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode > 299 {
		return GeneratedImage{}, fmt.Errorf("image provider status %d", response.StatusCode)
	}
	var result struct {
		Data []struct {
			Base64 string `json:"b64_json"`
			URL    string `json:"url"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 32<<20)).Decode(&result); err != nil {
		return GeneratedImage{}, err
	}
	if len(result.Data) == 0 || (strings.TrimSpace(result.Data[0].Base64) == "" && strings.TrimSpace(result.Data[0].URL) == "") {
		return GeneratedImage{}, ErrEmptyResult
	}
	return GeneratedImage{DataBase64: result.Data[0].Base64, SourceURL: result.Data[0].URL, ContentType: "image/png"}, nil
}

func (c ComfyImageGenerationClient) Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error) {
	if strings.TrimSpace(prompt) == "" {
		return GeneratedImage{}, errors.New("image prompt is empty")
	}
	width, height := imageSize(size)
	workflow := map[string]any{
		"3":  map[string]any{"inputs": map[string]any{"seed": time.Now().UnixNano(), "steps": 8, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0, "model": []any{"16", 0}, "positive": []any{"6", 0}, "negative": []any{"7", 0}, "latent_image": []any{"13", 0}}, "class_type": "KSampler"},
		"6":  map[string]any{"inputs": map[string]any{"text": prompt, "clip": []any{"18", 0}}, "class_type": "CLIPTextEncode"},
		"7":  map[string]any{"inputs": map[string]any{"text": "blurry, ugly, distorted", "clip": []any{"18", 0}}, "class_type": "CLIPTextEncode"},
		"8":  map[string]any{"inputs": map[string]any{"samples": []any{"3", 0}, "vae": []any{"17", 0}}, "class_type": "VAEDecode"},
		"9":  map[string]any{"inputs": map[string]any{"filename_prefix": "domovoy", "images": []any{"8", 0}}, "class_type": "SaveImage"},
		"13": map[string]any{"inputs": map[string]any{"width": width, "height": height, "batch_size": 1}, "class_type": "EmptySD3LatentImage"},
		"16": map[string]any{"inputs": map[string]any{"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
		"17": map[string]any{"inputs": map[string]any{"vae_name": "ae.safetensors"}, "class_type": "VAELoader"},
		"18": map[string]any{"inputs": map[string]any{"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}, "class_type": "CLIPLoader"},
	}
	payload := map[string]any{"prompt": workflow, "client_id": fmt.Sprintf("domovoy-%d", time.Now().UnixNano())}
	encoded, err := json.Marshal(payload)
	if err != nil {
		return GeneratedImage{}, err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/prompt", bytes.NewReader(encoded))
	if err != nil {
		return GeneratedImage{}, err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := c.client().Do(request)
	if err != nil {
		return GeneratedImage{}, err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode > 299 {
		return GeneratedImage{}, fmt.Errorf("comfy provider status %d", response.StatusCode)
	}
	var accepted struct {
		PromptID string `json:"prompt_id"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&accepted); err != nil {
		return GeneratedImage{}, err
	}
	if accepted.PromptID == "" {
		return GeneratedImage{}, ErrEmptyResult
	}
	// A caller timeout must remove only this prompt. ComfyUI's global
	// /interrupt endpoint can otherwise stop unrelated users' work and leave
	// queued requests in an ambiguous state.
	defer func() {
		if ctx.Err() != nil {
			cancelCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			if err := c.cancelPrompt(cancelCtx, accepted.PromptID); err != nil {
				// The original context error is more useful to the caller; the
				// provider cleanup failure remains observable in the server log.
				log.Printf("comfy prompt cleanup failed for %s: %v", accepted.PromptID, err)
			}
		}
	}()

	interval := c.PollInterval
	if interval <= 0 {
		interval = time.Second
	}
	timeout := c.PollTimeout
	if timeout <= 0 {
		timeout = 10 * time.Minute
	}
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		image, done, err := c.historyImage(ctx, accepted.PromptID)
		if err != nil {
			return GeneratedImage{}, err
		}
		if done {
			return image, nil
		}
		select {
		case <-ctx.Done():
			return GeneratedImage{}, ctx.Err()
		case <-deadline.C:
			return GeneratedImage{}, context.DeadlineExceeded
		case <-ticker.C:
		}
	}
}

func (c ComfyImageGenerationClient) cancelPrompt(ctx context.Context, promptID string) error {
	if strings.TrimSpace(promptID) == "" {
		return errors.New("comfy prompt id is empty")
	}
	payload, err := json.Marshal(map[string][]string{"delete": {promptID}})
	if err != nil {
		return err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/queue", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := c.client().Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode > 299 {
		return fmt.Errorf("comfy prompt cleanup status %d", response.StatusCode)
	}
	return nil
}

func (c ComfyImageGenerationClient) historyImage(ctx context.Context, promptID string) (GeneratedImage, bool, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+"/history/"+url.PathEscape(promptID), nil)
	if err != nil {
		return GeneratedImage{}, false, err
	}
	response, err := c.client().Do(request)
	if err != nil {
		return GeneratedImage{}, false, err
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusNotFound {
		return GeneratedImage{}, false, nil
	}
	if response.StatusCode < 200 || response.StatusCode > 299 {
		return GeneratedImage{}, false, fmt.Errorf("comfy history status %d", response.StatusCode)
	}
	var history map[string]struct {
		Outputs map[string]struct {
			Images []struct {
				Filename  string `json:"filename"`
				Subfolder string `json:"subfolder"`
				Type      string `json:"type"`
			} `json:"images"`
		} `json:"outputs"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 4<<20)).Decode(&history); err != nil {
		return GeneratedImage{}, false, err
	}
	record, ok := history[promptID]
	if !ok {
		return GeneratedImage{}, false, nil
	}
	for _, output := range record.Outputs {
		for _, image := range output.Images {
			if image.Filename == "" {
				continue
			}
			query := url.Values{"filename": {image.Filename}, "subfolder": {image.Subfolder}, "type": {image.Type}}
			imageRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+"/view?"+query.Encode(), nil)
			if err != nil {
				return GeneratedImage{}, false, err
			}
			imageResponse, err := c.client().Do(imageRequest)
			if err != nil {
				return GeneratedImage{}, false, err
			}
			data, readErr := io.ReadAll(io.LimitReader(imageResponse.Body, 32<<20))
			imageResponse.Body.Close()
			if readErr != nil {
				return GeneratedImage{}, false, readErr
			}
			if imageResponse.StatusCode < 200 || imageResponse.StatusCode > 299 || len(data) == 0 {
				return GeneratedImage{}, false, ErrEmptyResult
			}
			contentType := imageResponse.Header.Get("Content-Type")
			if contentType == "" {
				contentType = "image/png"
			}
			return GeneratedImage{DataBase64: encodeBase64(data), ContentType: contentType}, true, nil
		}
	}
	return GeneratedImage{}, true, ErrEmptyResult
}

func imageSize(size string) (int, int) {
	parts := strings.Split(strings.TrimSpace(size), "x")
	if len(parts) == 2 {
		var width, height int
		if _, err := fmt.Sscanf(size, "%dx%d", &width, &height); err == nil && width >= 256 && height >= 256 && width <= 2048 && height <= 2048 {
			return width, height
		}
	}
	return 512, 512
}

// Describe accepts a video skill endpoint that performs its own bounded frame
// sampling. Frame extraction is provider-owned so the core does not acquire a
// media-processing dependency or hide a policy in orchestration code.
func (c VideoUnderstandingClient) Describe(ctx context.Context, model, filename, contentType string, data []byte, instruction string) (string, error) {
	if len(data) == 0 {
		return "", errors.New("video payload is empty")
	}
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	header := make(textproto.MIMEHeader)
	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, strings.ReplaceAll(filename, `"`, "")))
	header.Set("Content-Type", contentType)
	part, err := writer.CreatePart(header)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(data); err != nil {
		return "", err
	}
	_ = writer.WriteField("model", model)
	_ = writer.WriteField("instruction", instruction)
	if err := writer.Close(); err != nil {
		return "", err
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/understand", body)
	if err != nil {
		return "", err
	}
	request.Header.Set("Content-Type", writer.FormDataContentType())
	response, err := c.client().Do(request)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode > 299 {
		return "", fmt.Errorf("video provider status %d", response.StatusCode)
	}
	var result struct {
		Text string `json:"text"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&result); err != nil {
		return "", err
	}
	if strings.TrimSpace(result.Text) == "" {
		return "", ErrEmptyResult
	}
	return strings.TrimSpace(result.Text), nil
}

func (c SpeechClient) client() *http.Client {
	if c.Client != nil {
		return c.Client
	}
	return http.DefaultClient
}

func (c VisionClient) client() *http.Client {
	if c.Client != nil {
		return c.Client
	}
	return http.DefaultClient
}

func (c ImageGenerationClient) client() *http.Client {
	if c.Client != nil {
		return c.Client
	}
	return http.DefaultClient
}

func (c ComfyImageGenerationClient) client() *http.Client {
	if c.Client != nil {
		return c.Client
	}
	return http.DefaultClient
}

func (c VideoUnderstandingClient) client() *http.Client {
	if c.Client != nil {
		return c.Client
	}
	return http.DefaultClient
}

func encodeBase64(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}