package memoir

import (
	"bytes"
	"context"
	"data-terminal-cloud/internal/extraction"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"time"
)

type ArkIllustrator struct {
	BaseURL, APIKey, Model string
	Client                 *http.Client
}

func NewIllustrator(base, key, model string) *ArkIllustrator {
	return &ArkIllustrator{strings.TrimRight(base, "/"), key, model, &http.Client{Timeout: 120 * time.Second}}
}
func (a *ArkIllustrator) Generate(ctx context.Context, prompt string) ([]byte, error) {
	body, _ := json.Marshal(map[string]any{"model": a.Model, "prompt": prompt + "。温暖的水彩情景插画，米色纸张肌理，禁止文字和水印以外的标记。这是基于叙述的想象性插画，不是真实历史照片。", "size": "2K", "output_format": "jpeg", "response_format": "b64_json", "watermark": true})
	req, e := http.NewRequestWithContext(ctx, "POST", a.BaseURL+"/images/generations", bytes.NewReader(body))
	if e != nil {
		return nil, e
	}
	req.Header.Set("Authorization", "Bearer "+a.APIKey)
	req.Header.Set("Content-Type", "application/json")
	resp, e := a.Client.Do(req)
	if e != nil {
		return nil, errors.New("MEMOIR_IMAGE_UNAVAILABLE")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		var failure struct {
			Error struct {
				Code string `json:"code"`
			} `json:"error"`
		}
		_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&failure)
		if failure.Error.Code == "AccountOverdueError" {
			return nil, extraction.ErrModelAccountOverdue
		}
		return nil, errors.New("MEMOIR_IMAGE_REJECTED")
	}
	var out struct {
		Data []struct {
			Base64 string `json:"b64_json"`
		} `json:"data"`
	}
	if json.NewDecoder(io.LimitReader(resp.Body, 16<<20)).Decode(&out) != nil || len(out.Data) != 1 {
		return nil, errors.New("MEMOIR_IMAGE_INVALID")
	}
	b, e := base64.StdEncoding.DecodeString(out.Data[0].Base64)
	if e != nil || len(b) > 10<<20 || http.DetectContentType(b) != "image/jpeg" {
		return nil, errors.New("MEMOIR_IMAGE_INVALID")
	}
	return b, nil
}