package server

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

// Pictures for the assistant: the model asks for one in its answer - a
// ```image block, a new picture or an edit of a photo - and the panel makes
// it through the Vertex API (https://s2.x-grk.com/api/vertex/docs). The
// work is a job there: it is put, then asked about until it is done. The
// API key is the admins', set in Настройки and never sent to the browser.

const (
	cfgImageKey = "ai_image_key"
	cfgImageURL = "ai_image_url" // not in the settings: the service's address, for tests

	// DefaultImageURL is where the Vertex API is.
	DefaultImageURL = "https://s2.x-grk.com"

	imageMaxBytes = 20 << 20 // what the service takes of one photo
)

// imageConf is the service's address and the key (none: pictures are off).
func (s *Server) imageConf() (base, key string) {
	base, _ = s.db.GetConfig(cfgImageURL)
	if base == "" {
		base = DefaultImageURL
	}
	key, _ = s.db.GetConfig(cfgImageKey)
	return strings.TrimRight(base, "/"), strings.TrimSpace(key)
}

// imagePrompt tells the model it can make pictures, and how - when there
// is a key.
func (s *Server) imagePrompt() string {
	if _, key := s.imageConf(); key == "" {
		return ""
	}
	return "\n\nКартинки. Ты умеешь рисовать и править фото: сам, по смыслу разговора, без просьбы «вызови инструмент» и без лишних вопросов, " +
		"когда пользователь просит нарисовать, сгенерировать, изобразить, сделать картинку, арт, логотип, обои, иллюстрацию " +
		"или изменить, дорисовать, убрать, заменить что-то на фото.\n" +
		"```image\nописание картинки\n``` - создать новую. Опиши сцену, стиль, свет, детали; язык описания любой.\n" +
		"```image edit\nчто изменить на фото\n``` - править фото пользователя: то, что приложено к его сообщению " +
		"(приложено несколько - после edit назови файл: ```image edit имя.jpg```); фото нет - последнюю созданную картинку.\n" +
		"В первой строке блока можно задать формат: 1:1 (по умолчанию), 16:9, 9:16, 3:2, 2:3, 4:3, 3:4, auto - и качество: " +
		"fast (по умолчанию, 1 кредит), artist (5, лучше), 0.2plus (10, лучшее); повышай качество, только если пользователь просит. " +
		"Например: ```image 16:9 artist```.\n" +
		"Сервис по умолчанию сам улучшает описание (дописывает детали и стиль) - это хорошо для короткого запроса. " +
		"Если пользователь дал готовое, точное описание (промпт целиком, логотип «ровно как сказано», текст на картинке, строгая композиция) " +
		"или просит ничего не добавлять - добавь noenhance: ```image noenhance 3:2```. Для правки фото это не применяется.\n" +
		"Картинку покажет сама панель: ссылку в ответ не вставляй, скажи коротко, что получилось, и предложи поправить. " +
		"Делается она до пары минут. Одно действие за ответ."
}

// imageError says in Russian what an error code of the service means.
func imageError(code, msg string) string {
	switch code {
	case "unauthorized":
		return "сервис не принял API-ключ - проверьте его в Настройки › ИИ-ассистент › Генерация изображений"
	case "banned":
		return "аккаунт сервиса картинок заблокирован"
	case "insufficient_credits":
		return "на аккаунте сервиса картинок не хватает кредитов"
	case "daily_limit":
		return "исчерпан дневной лимит картинок"
	case "job_busy":
		return "уже делается другая картинка - дождитесь её"
	case "invalid_prompt":
		return "сервис не принял описание (пустое или длиннее 4000 символов)"
	case "invalid_image":
		return "сервис не принял фото (нужен png, jpg или webp)"
	case "image_too_large":
		return "фото больше 20 МБ"
	case "tokens_exhausted":
		return "сервис картинок перегружен - попробуйте позже"
	case "generation_failed", "edit_failed":
		return "сервис не справился с картинкой (кредиты вернутся)"
	case "source_missing":
		return "исходное фото потерялось на сервисе - загрузите его заново"
	}
	if msg == "" {
		msg = "сервис картинок ответил ошибкой"
	}
	return msg
}

var imageClient = &http.Client{Timeout: 90 * time.Second}

// imageCall sends the request to the service with the key and reads its
// JSON; an error is words for the user.
func imageCall(ctx context.Context, req *http.Request, key string, v any) error {
	req = req.WithContext(ctx)
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Accept", "application/json")
	resp, err := imageClient.Do(req)
	if err != nil {
		return errors.New("сервис картинок не отвечает: " + err.Error())
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	var e struct {
		OK        *bool  `json:"ok"`
		ErrorCode string `json:"error_code"`
		Error     string `json:"error"`
	}
	if json.Unmarshal(b, &e) != nil {
		return errors.New("сервис картинок ответил " + resp.Status)
	}
	if (e.OK != nil && !*e.OK) || resp.StatusCode/100 != 2 {
		return errors.New(imageError(e.ErrorCode, e.Error))
	}
	return json.Unmarshal(b, v)
}

var (
	imageRatios = map[string]bool{"1:1": true, "16:9": true, "9:16": true, "3:2": true, "2:3": true, "4:3": true, "3:4": true, "auto": true}
	imageModels = map[string]bool{"fast": true, "artist": true, "0.2plus": true}
	imageJobRe  = regexp.MustCompile(`^[A-Za-z0-9_-]{1,64}$`)
	imageExt    = map[string]string{"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg", "image/webp": "webp"}
)

// POST /api/panel/ai/image - {mode: generate|edit, prompt, model, ratio,
// image: {mime, data (base64)} | image_url}: puts the job, answers its id.
func (s *Server) handleAIImage(w http.ResponseWriter, r *http.Request) {
	fail := func(msg string) { writeJSON(w, map[string]any{"ok": false, "error": msg}) }
	if !s.aiOn(s.currentAdmin(r).ID) {
		fail("ассистент выключен в профиле")
		return
	}
	base, key := s.imageConf()
	if key == "" {
		fail("не задан API-ключ сервиса картинок (Настройки › ИИ-ассистент › Генерация изображений)")
		return
	}
	var in struct {
		Mode     string `json:"mode"`
		Prompt   string `json:"prompt"`
		Model    string `json:"model"`
		Ratio    string `json:"ratio"`
		Enhance  *bool  `json:"enhance"` // the service improves the description; on unless false
		ImageURL string `json:"image_url"`
		Image    struct {
			MIME string `json:"mime"`
			Data string `json:"data"`
		} `json:"image"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 32<<20)).Decode(&in); err != nil {
		fail("плохой запрос")
		return
	}
	in.Prompt = strings.TrimSpace(in.Prompt)
	if in.Prompt == "" {
		fail("пустое описание картинки")
		return
	}
	if r := []rune(in.Prompt); len(r) > 4000 {
		in.Prompt = string(r[:4000])
	}
	var job struct {
		JobID string `json:"job_id"`
		Cost  int    `json:"cost"`
	}
	ctx := r.Context()
	if in.Mode != "edit" {
		if !imageModels[in.Model] {
			in.Model = "fast"
		}
		if !imageRatios[in.Ratio] {
			in.Ratio = "1:1"
		}
		enhance := in.Enhance == nil || *in.Enhance
		b, _ := json.Marshal(map[string]any{"prompt": in.Prompt, "model": in.Model, "ratio": in.Ratio, "enhance": enhance})
		req, _ := http.NewRequest(http.MethodPost, base+"/api/vertex/v1/generate", bytes.NewReader(b))
		req.Header.Set("Content-Type", "application/json")
		if err := imageCall(ctx, req, key, &job); err != nil {
			fail(err.Error())
			return
		}
	} else {
		mime, data, err := s.imageSource(ctx, base, in.Image.MIME, in.Image.Data, in.ImageURL)
		if err != nil {
			fail(err.Error())
			return
		}
		var body bytes.Buffer
		mw := multipart.NewWriter(&body)
		h := textproto.MIMEHeader{}
		h.Set("Content-Disposition", `form-data; name="image"; filename="photo.`+imageExt[mime]+`"`)
		h.Set("Content-Type", mime)
		part, _ := mw.CreatePart(h)
		_, _ = part.Write(data)
		_ = mw.WriteField("prompt", in.Prompt)
		_ = mw.Close()
		req, _ := http.NewRequest(http.MethodPost, base+"/api/vertex/v1/edit", &body)
		req.Header.Set("Content-Type", mw.FormDataContentType())
		if err := imageCall(ctx, req, key, &job); err != nil {
			fail(err.Error())
			return
		}
	}
	if job.JobID == "" {
		fail("сервис картинок не дал номер задачи")
		return
	}
	writeJSON(w, map[string]any{"ok": true, "job_id": job.JobID, "cost": job.Cost})
}

// imageSource is the photo to edit: sent by the browser, or a picture made
// before, taken from the service's own address (and from no other).
func (s *Server) imageSource(ctx context.Context, base, mime, b64, picURL string) (string, []byte, error) {
	if picURL != "" {
		u, err := url.Parse(picURL)
		bu, _ := url.Parse(base)
		if err != nil || bu == nil || u.Host != bu.Host || (u.Scheme != "https" && u.Scheme != "http") {
			return "", nil, errors.New("прежнюю картинку можно править только с сервиса картинок")
		}
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
		resp, err := imageClient.Do(req)
		if err != nil || resp.StatusCode/100 != 2 {
			if resp != nil {
				resp.Body.Close()
			}
			return "", nil, errors.New("прежняя картинка больше не открывается")
		}
		defer resp.Body.Close()
		data, _ := io.ReadAll(io.LimitReader(resp.Body, imageMaxBytes+1))
		mime = strings.ToLower(strings.SplitN(resp.Header.Get("Content-Type"), ";", 2)[0])
		if imageExt[mime] == "" {
			mime = "image/png"
		}
		return mime, data, nil
	}
	mime = strings.ToLower(mime)
	if imageExt[mime] == "" {
		return "", nil, errors.New("править можно png, jpg и webp")
	}
	data, err := base64.StdEncoding.DecodeString(b64)
	if err != nil || len(data) == 0 {
		return "", nil, errors.New("нет фото для правки: прикрепите его к сообщению")
	}
	if len(data) > imageMaxBytes {
		return "", nil, errors.New("фото больше 20 МБ")
	}
	return mime, data, nil
}

// GET /api/panel/ai/image/{job}?wait=1..45 - the job's state, waiting up to
// wait seconds for it to end.
func (s *Server) handleAIImageJob(w http.ResponseWriter, r *http.Request) {
	fail := func(msg string) { writeJSON(w, map[string]any{"ok": false, "error": msg}) }
	if !s.aiOn(s.currentAdmin(r).ID) {
		fail("ассистент выключен в профиле")
		return
	}
	base, key := s.imageConf()
	id := r.PathValue("job")
	if key == "" || !imageJobRe.MatchString(id) {
		fail("нет такой задачи")
		return
	}
	wait, _ := strconv.Atoi(r.URL.Query().Get("wait"))
	if wait < 0 {
		wait = 0
	}
	if wait > 45 {
		wait = 45
	}
	req, _ := http.NewRequest(http.MethodGet, base+"/api/vertex/v1/jobs/"+id+"?wait="+strconv.Itoa(wait), nil)
	var st struct {
		Status    string `json:"status"`
		PhotoURL  string `json:"photo_url"`
		Cost      int    `json:"cost"`
		Balance   int    `json:"balance"`
		ErrorCode string `json:"error_code"`
		Error     string `json:"error"`
	}
	if err := imageCall(r.Context(), req, key, &st); err != nil {
		fail(err.Error())
		return
	}
	out := map[string]any{"ok": true, "status": st.Status}
	switch st.Status {
	case "done":
		out["photo_url"], out["cost"], out["balance"] = st.PhotoURL, st.Cost, st.Balance
	case "failed":
		out["error"] = imageError(st.ErrorCode, st.Error)
	}
	writeJSON(w, out)
}