package llamaswap

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"

	majaudio "gitea.stevedudenhoeffer.com/steve/majordomo/audio"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// SpeechModel implements audio.SpeechProvider, binding a text-to-speech model
// served by llama-swap (routed to a kokoro/chatterbox-style OpenAI-compatible
// upstream). The id is passed through verbatim and selects which upstream
// llama-swap loads.
func (p *Provider) SpeechModel(id string, opts ...majaudio.SpeechModelOption) (majaudio.SpeechModel, error) {
	if p.baseURL == "" {
		return nil, fmt.Errorf("llama-swap provider %q: no base URL configured (set one via WithBaseURL or an LLM_* env DSN)", p.name)
	}
	_ = majaudio.ApplySpeechModelOptions(opts)
	return &speechModel{p: p, id: id}, nil
}

type speechModel struct {
	p  *Provider
	id string
}

// speechRequest is the OpenAI /v1/audio/speech request shape. llama-swap
// routes by the `model` field in the body.
type speechRequest struct {
	Model          string  `json:"model"`
	Input          string  `json:"input"`
	Voice          string  `json:"voice,omitempty"`
	ResponseFormat string  `json:"response_format,omitempty"`
	Speed          float64 `json:"speed,omitempty"`
}

// Speak implements audio.SpeechModel via POST {base}/v1/audio/speech.
func (m *speechModel) Speak(ctx context.Context, req majaudio.SpeechRequest, opts ...majaudio.SpeechOption) (*majaudio.SpeechResult, error) {
	req = req.Apply(opts...)
	if strings.TrimSpace(req.Input) == "" {
		return nil, fmt.Errorf("%w: speech synthesis requires input text", llm.ErrUnsupported)
	}
	wire := speechRequest{
		Model:          m.id,
		Input:          req.Input,
		Voice:          req.Voice,
		ResponseFormat: req.Format,
		Speed:          req.Speed,
	}
	body, err := json.Marshal(wire)
	if err != nil {
		return nil, fmt.Errorf("llama-swap: encode speech request: %w", err)
	}
	audioBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/speech", m.id, "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	if len(audioBytes) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "speech response contained no audio"}
	}
	return &majaudio.SpeechResult{Audio: audioBytes, MIME: speechMIME(contentType, req.Format)}, nil
}

// speechMIME resolves the result MIME type: the response Content-Type when it
// is a concrete audio type, else a mapping from the requested format, else
// audio/mpeg (the OpenAI endpoint's default container is mp3).
func speechMIME(contentType, format string) string {
	if mt, _, err := mime.ParseMediaType(contentType); err == nil {
		if strings.HasPrefix(mt, "audio/") || strings.HasPrefix(mt, "video/") {
			return mt
		}
	}
	switch strings.ToLower(strings.TrimSpace(format)) {
	case "", "mp3":
		return "audio/mpeg"
	case "wav":
		return "audio/wav"
	case "opus":
		return "audio/ogg"
	case "aac":
		return "audio/aac"
	case "flac":
		return "audio/flac"
	case "pcm":
		return "audio/pcm"
	default:
		return "audio/mpeg"
	}
}

// TranscriptionModel implements audio.TranscriptionProvider, binding a
// speech-to-text model served by llama-swap (routed to a whisper.cpp-style
// upstream).
func (p *Provider) TranscriptionModel(id string, opts ...majaudio.TranscriptionModelOption) (majaudio.TranscriptionModel, error) {
	if p.baseURL == "" {
		return nil, fmt.Errorf("llama-swap provider %q: no base URL configured (set one via WithBaseURL or an LLM_* env DSN)", p.name)
	}
	_ = majaudio.ApplyTranscriptionModelOptions(opts)
	return &transcriptionModel{p: p, id: id}, nil
}

type transcriptionModel struct {
	p  *Provider
	id string
}

// Transcribe implements audio.TranscriptionModel via POST
// {base}/v1/audio/transcriptions (multipart/form-data — llama-swap routes by
// the `model` form field).
func (m *transcriptionModel) Transcribe(ctx context.Context, req majaudio.TranscriptionRequest, opts ...majaudio.TranscriptionOption) (*majaudio.TranscriptionResult, error) {
	req = req.Apply(opts...)
	if len(req.Audio) == 0 {
		return nil, fmt.Errorf("%w: transcription requires audio bytes", llm.ErrUnsupported)
	}

	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	fw, err := w.CreateFormFile("file", transcriptionFilename(req))
	if err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}
	if _, err := fw.Write(req.Audio); err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}
	fields := map[string]string{
		"model":           m.id,
		"language":        req.Language,
		"prompt":          req.Prompt,
		"response_format": "json",
	}
	for k, v := range fields {
		if v == "" {
			continue
		}
		if err := w.WriteField(k, v); err != nil {
			return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
		}
	}
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}

	raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, w.FormDataContentType(), &buf)
	if err != nil {
		return nil, err
	}
	var out struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llama-swap: decode transcription response: %w", err)
	}
	return &majaudio.TranscriptionResult{Text: out.Text, Raw: json.RawMessage(raw)}, nil
}

// transcriptionFilename picks the multipart filename hint: the caller's, else
// one derived from the MIME subtype ("audio.mp3"), else "audio".
func transcriptionFilename(req majaudio.TranscriptionRequest) string {
	if req.Filename != "" {
		return req.Filename
	}
	switch strings.ToLower(req.MIME) {
	case "audio/mpeg", "audio/mp3":
		return "audio.mp3"
	case "audio/wav", "audio/x-wav", "audio/wave":
		return "audio.wav"
	case "audio/ogg", "audio/opus":
		return "audio.ogg"
	case "audio/flac", "audio/x-flac":
		return "audio.flac"
	case "audio/mp4", "audio/m4a", "audio/x-m4a":
		return "audio.m4a"
	case "audio/webm", "video/webm":
		return "audio.webm"
	default:
		return "audio"
	}
}

// ListVoices returns the voices a TTS model offers (GET
// /v1/audio/voices?model=...). The decode is tolerant: upstreams answer with
// either a bare string list or a list of {id|name} objects.
func (p *Provider) ListVoices(ctx context.Context, model string) ([]string, error) {
	if model == "" {
		return nil, fmt.Errorf("llama-swap: ListVoices requires a model id")
	}
	raw, _, err := p.doRaw(ctx, http.MethodGet, "/v1/audio/voices?model="+url.QueryEscape(model), model, "", nil)
	if err != nil {
		return nil, err
	}
	return parseVoices(raw)
}

// parseVoices extracts voice names from the various shapes upstreams use:
// {"voices":[...]} or {"data":[...]} envelopes (or a bare array), holding
// either strings or objects keyed by id/name/voice_id.
func parseVoices(raw []byte) ([]string, error) {
	var env struct {
		Voices json.RawMessage `json:"voices"`
		Data   json.RawMessage `json:"data"`
	}
	list := json.RawMessage(raw)
	if err := json.Unmarshal(raw, &env); err == nil {
		if len(env.Voices) > 0 {
			list = env.Voices
		} else if len(env.Data) > 0 {
			list = env.Data
		}
	}

	var names []string
	if err := json.Unmarshal(list, &names); err == nil {
		return names, nil
	}
	// A failed decode above may have partially populated names — start fresh.
	names = nil
	var objs []struct {
		ID      string `json:"id"`
		Name    string `json:"name"`
		VoiceID string `json:"voice_id"`
	}
	if err := json.Unmarshal(list, &objs); err == nil {
		for _, o := range objs {
			switch {
			case o.ID != "":
				names = append(names, o.ID)
			case o.Name != "":
				names = append(names, o.Name)
			case o.VoiceID != "":
				names = append(names, o.VoiceID)
			}
		}
		return names, nil
	}
	return nil, fmt.Errorf("llama-swap: unrecognized voices payload shape")
}

// doRaw performs a request to a llama-swap endpoint and returns the raw
// response body and its Content-Type — the sibling of doJSON for endpoints
// whose success payload is not JSON (audio bytes) or whose shape varies.
// contentType sets the request Content-Type when body is non-nil.
func (p *Provider) doRaw(ctx context.Context, method, path, model, contentType string, body io.Reader) ([]byte, string, error) {
	if p.baseURL == "" {
		return nil, "", fmt.Errorf("llama-swap provider %q: no base URL configured (set one via WithBaseURL or an LLM_* env DSN)", p.name)
	}
	req, err := http.NewRequestWithContext(ctx, method, p.baseURL+path, body)
	if err != nil {
		return nil, "", fmt.Errorf("llama-swap: build request: %w", err)
	}
	if body != nil && contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if p.token != "" {
		req.Header.Set("Authorization", "Bearer "+p.token)
	}
	resp, err := p.client.Do(req)
	if err != nil {
		return nil, "", fmt.Errorf("llama-swap: do request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return nil, "", p.apiError(resp, model)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
	if err != nil {
		return nil, "", fmt.Errorf("llama-swap: read response: %w", err)
	}
	return data, resp.Header.Get("Content-Type"), nil
}