// restore.go implements the imagegen.ColorizeProvider and
// imagegen.FaceRestoreProvider surfaces against the mediautils shim reached
// through llama-swap's /upstream passthrough (ADR-0023):
//
//	colorize       POST /upstream/<id>/v1/colorize        (DDColor)
//	restore_faces  POST /upstream/<id>/v1/restore_faces   (GFPGAN)
//
// Both are one-file multipart in, one PNG out, mirroring mediautil.go.
package llamaswap

import (
	"context"
	"fmt"
	"net/http"
	"strconv"

	"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// --- colorize ---

// ColorizeModel implements imagegen.ColorizeProvider against the mediautils
// shim's POST /v1/colorize. The id selects which upstream llama-swap loads.
func (p *Provider) ColorizeModel(id string, opts ...imagegen.ColorizeModelOption) (imagegen.Colorizer, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = imagegen.ApplyColorizeModelOptions(opts)
	return &colorizeModel{p: p, id: id}, nil
}

type colorizeModel struct {
	p  *Provider
	id string
}

// Colorize implements imagegen.Colorizer.
func (m *colorizeModel) Colorize(ctx context.Context, req imagegen.ColorizeRequest, opts ...imagegen.ColorizeOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Image.Data) == 0 {
		return nil, fmt.Errorf("%w: colorization requires an image", llm.ErrUnsupported)
	}
	path, err := upstreamPath(m.id, "/v1/colorize")
	if err != nil {
		return nil, err
	}
	body, contentType, err := buildMultipart("build colorize form",
		filePart{field: "file", filename: "image.png", data: req.Image.Data},
		nil)
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
	if err != nil {
		return nil, err
	}
	return singleImageResult(m.p.name, m.id, "colorize", raw, respType)
}

// --- face restoration ---

// FaceRestoreModel implements imagegen.FaceRestoreProvider against the
// mediautils shim's POST /v1/restore_faces.
func (p *Provider) FaceRestoreModel(id string, opts ...imagegen.FaceRestoreModelOption) (imagegen.FaceRestorer, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = imagegen.ApplyFaceRestoreModelOptions(opts)
	return &faceRestoreModel{p: p, id: id}, nil
}

type faceRestoreModel struct {
	p  *Provider
	id string
}

// RestoreFaces implements imagegen.FaceRestorer.
func (m *faceRestoreModel) RestoreFaces(ctx context.Context, req imagegen.FaceRestoreRequest, opts ...imagegen.FaceRestoreOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Image.Data) == 0 {
		return nil, fmt.Errorf("%w: face restoration requires an image", llm.ErrUnsupported)
	}
	if req.Upscale != 0 && req.Upscale != 1 && req.Upscale != 2 {
		return nil, fmt.Errorf("%w: face-restore upscale must be 1 or 2, got %d", llm.ErrUnsupported, req.Upscale)
	}
	path, err := upstreamPath(m.id, "/v1/restore_faces")
	if err != nil {
		return nil, err
	}
	upscale := ""
	if req.Upscale != 0 {
		upscale = strconv.Itoa(req.Upscale)
	}
	body, contentType, err := buildMultipart("build restore-faces form",
		filePart{field: "file", filename: "image.png", data: req.Image.Data},
		[]formField{{"upscale", upscale, false}})
	if err != nil {
		return nil, err
	}
	raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes)
	if err != nil {
		return nil, err
	}
	return singleImageResult(m.p.name, m.id, "face restoration", raw, respType)
}