Jobmaker-LdPost/art/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"image"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/disintegration/imaging"
	"github.com/spf13/cobra"
	"ldpost/shared/config"
	"ldpost/shared/gemini"
	"ldpost/shared/groq"
	"ldpost/shared/state"
	"ldpost/shared/telegram"
	"ldpost/shared/workspace"
)

// ─── Types ────────────────────────────────────────────────────────────────────

type slidePrompts struct {
	DescricaoOriginal string `json:"descricao_original"`
	PromptGemini      string `json:"prompt_gemini"`
}

type artPromptsFile struct {
	Slide1 slidePrompts `json:"slide1"`
	Slide2 slidePrompts `json:"slide2"`
}

type artChoices struct {
	Slide1 string // "bottom" or "right"
	Slide2 string
}

// ─── Text parsing ─────────────────────────────────────────────────────────────

// parseImageSection extracts content under "## Imagem N" heading.
func parseImageSection(text, heading string) string {
	lines := strings.Split(text, "\n")
	inSection := false
	var sb strings.Builder
	for _, line := range lines {
		if strings.HasPrefix(line, "## ") {
			title := strings.TrimSpace(strings.TrimPrefix(line, "## "))
			if strings.EqualFold(title, heading) {
				inSection = true
				continue
			} else if inSection {
				break
			}
		}
		if inSection {
			sb.WriteString(line + "\n")
		}
	}
	return strings.TrimSpace(sb.String())
}

// firstNWords returns first n words from a string.
func firstNWords(s string, n int) string {
	words := strings.Fields(s)
	if len(words) <= n {
		return s
	}
	return strings.Join(words[:n], " ")
}

// ─── Groq prompt builder ──────────────────────────────────────────────────────

const systemPromptImgTranslate = `Você converte descrições de imagens em prompts otimizados para geração com Gemini Image Generation.

Regras:
- Prompt em inglês
- Estilo: clean, professional, tech-focused, suitable for LinkedIn carousel
- Incluir: composição, estilo visual, paleta de cores sugerida (azul/branco/cinza para tech)
- Evitar: rostos humanos, texto em destaque (o Gemini erra texto), logos de marcas reais
- Máximo 100 palavras
- Retornar APENAS o prompt, sem explicação`

func buildGeminiPrompt(gc *groq.GroqClient, postContext, desc string) (string, error) {
	system := systemPromptImgTranslate + "\n\nContexto do post: " + firstNWords(postContext, 50)
	user := "Descrição original: " + desc
	prompt, err := gc.Chat(groq.TextModel, groq.TextMessages(system, user), 0.4, 150)
	if err != nil {
		return "", fmt.Errorf("traduzir prompt: %w", err)
	}
	return strings.TrimSpace(prompt), nil
}

// ─── Image helpers ────────────────────────────────────────────────────────────

// findInputImage looks for input/imagemN in common formats.
func findInputImage(inputPath, base string) string {
	for _, ext := range []string{".png", ".jpg", ".jpeg", ".webp"} {
		p := filepath.Join(inputPath, base+ext)
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}
	return ""
}

func cropBottomVariant(img image.Image, px int) image.Image {
	b := img.Bounds()
	newH := b.Max.Y - b.Min.Y - px
	if newH <= 0 {
		return img
	}
	return imaging.Crop(img, image.Rect(b.Min.X, b.Min.Y, b.Max.X, b.Min.Y+newH))
}

func cropRightVariant(img image.Image, px int) image.Image {
	b := img.Bounds()
	newW := b.Max.X - b.Min.X - px
	if newW <= 0 {
		return img
	}
	return imaging.Crop(img, image.Rect(b.Min.X, b.Min.Y, b.Min.X+newW, b.Max.Y))
}

// saveVariants crops bottom and right variants, returns (bottomPath, rightPath).
func saveVariants(rawPath, artRawDir, name string, cropBottom, cropRight int) (string, string, error) {
	img, err := imaging.Open(rawPath)
	if err != nil {
		return "", "", fmt.Errorf("abrir %s: %w", rawPath, err)
	}

	bottomPath := filepath.Join(artRawDir, name+"-crop-bottom.png")
	rightPath := filepath.Join(artRawDir, name+"-crop-right.png")

	if err := imaging.Save(cropBottomVariant(img, cropBottom), bottomPath); err != nil {
		return "", "", fmt.Errorf("salvar crop-bottom: %w", err)
	}
	if err := imaging.Save(cropRightVariant(img, cropRight), rightPath); err != nil {
		return "", "", fmt.Errorf("salvar crop-right: %w", err)
	}
	return bottomPath, rightPath, nil
}

// ─── Telegram flow ────────────────────────────────────────────────────────────

func sendArtPreview(bot *telegram.Bot, slug string,
	s1bottom, s1right, s2bottom, s2right string) error {

	if _, err := bot.SendMessage(fmt.Sprintf(
		"🎨 <b>ldpost-art</b> | <code>%s</code>\n\nGerei 2 imagens para o carrossel.\nEscolha a variante de cada uma:",
		slug)); err != nil {
		return fmt.Errorf("msg intro: %w", err)
	}

	if err := bot.SendMediaGroup([]telegram.MediaFile{
		{Path: s1bottom, Caption: "Slide 1 — A) Corte de baixo"},
		{Path: s1right, Caption: "Slide 1 — B) Corte da direita"},
	}); err != nil {
		log.Printf("[WARN] media group slide1: %v", err)
	}

	if err := bot.SendMediaGroup([]telegram.MediaFile{
		{Path: s2bottom, Caption: "Slide 2 — A) Corte de baixo"},
		{Path: s2right, Caption: "Slide 2 — B) Corte da direita"},
	}); err != nil {
		log.Printf("[WARN] media group slide2: %v", err)
	}

	buttons := [][]telegram.InlineButton{
		{
			{Text: "Slide 1: Corte de baixo", CallbackData: "s1_bottom"},
			{Text: "Slide 1: Corte da direita", CallbackData: "s1_right"},
		},
		{
			{Text: "Slide 2: Corte de baixo", CallbackData: "s2_bottom"},
			{Text: "Slide 2: Corte da direita", CallbackData: "s2_right"},
		},
		{
			{Text: "🔄 Regenerar tudo", CallbackData: "regenerar"},
			{Text: "✅ Confirmar escolhas", CallbackData: "confirmar"},
		},
	}
	if _, err := bot.SendMessageWithKeyboard("Selecione as variantes preferidas:", buttons); err != nil {
		return fmt.Errorf("botões: %w", err)
	}
	return nil
}

// waitChoices polls Telegram until both slides are chosen and confirmed.
// Returns ("regenerar", {}) if user requests full regeneration.
func waitChoices(bot *telegram.Bot) (action string, choices artChoices) {
	all := []string{"s1_bottom", "s1_right", "s2_bottom", "s2_right", "regenerar", "confirmar"}
	for {
		cb, err := bot.WaitForCallback(all, 0)
		if err != nil {
			log.Printf("[WARN] WaitForCallback: %v", err)
			continue
		}
		switch cb {
		case "s1_bottom":
			choices.Slide1 = "bottom"
			bot.SendMessage("✓ Slide 1: corte de baixo selecionado")
		case "s1_right":
			choices.Slide1 = "right"
			bot.SendMessage("✓ Slide 1: corte da direita selecionado")
		case "s2_bottom":
			choices.Slide2 = "bottom"
			bot.SendMessage("✓ Slide 2: corte de baixo selecionado")
		case "s2_right":
			choices.Slide2 = "right"
			bot.SendMessage("✓ Slide 2: corte da direita selecionado")
		case "regenerar":
			return "regenerar", artChoices{}
		case "confirmar":
			if choices.Slide1 == "" || choices.Slide2 == "" {
				bot.SendMessage("⚠️ Selecione variante para ambos os slides antes de confirmar.")
				continue
			}
			return "confirmar", choices
		}
	}
}

// ─── Main ─────────────────────────────────────────────────────────────────────

func main() {
	log.SetFlags(log.Ltime)

	var (
		flagPost           string
		flagWorkspace      string
		flagDryRun         bool
		flagSkipGeneration bool
		flagCropBottom     int
		flagCropRight      int
	)

	root := &cobra.Command{
		Use:   "ldpost-art --post <slug>",
		Short: "Gera imagens do carrossel via Gemini e envia para aprovação no Telegram",
		RunE: func(cmd *cobra.Command, args []string) error {
			if flagPost == "" {
				return fmt.Errorf("--post é obrigatório")
			}
			cfg := config.Load()
			if flagWorkspace != "" {
				cfg.Workspace = flagWorkspace
			}
			// CLI flags override .env
			if !cmd.Flags().Changed("crop-bottom") {
				flagCropBottom = cfg.CropBottomPx
			}
			if !cmd.Flags().Changed("crop-right") {
				flagCropRight = cfg.CropRightPx
			}

			if !flagDryRun && !flagSkipGeneration {
				if err := cfg.Validate("groq", "gemini"); err != nil {
					return err
				}
			} else if !flagDryRun {
				if err := cfg.Validate("groq"); err != nil {
					return err
				}
			}

			// ── 1. Encontrar post e validar state ─────────────────────────
			postPath, err := workspace.FindPostBySlug(cfg.Workspace, flagPost)
			if err != nil {
				return fmt.Errorf("post %q: %w", flagPost, err)
			}
			s, err := state.LoadState(postPath)
			if err != nil {
				return fmt.Errorf("ler state: %w", err)
			}
			if !s.IsStatus(state.StatusWaitingArt) {
				log.Printf("[ERROR] status atual: %q — esperado: %q", s.Status, state.StatusWaitingArt)
				return fmt.Errorf("estado incorreto — rode o agente correto para o status atual")
			}

			inputPath := workspace.InputPath(postPath)
			artRawPath := workspace.ArtRawPath(postPath)
			outputPath := workspace.OutputPath(postPath)

			// ── 2. Ler input/texto.md ─────────────────────────────────────
			textoRaw, err := os.ReadFile(workspace.InputTextoPath(postPath))
			if err != nil {
				return fmt.Errorf("input/texto.md: %w", err)
			}
			textoStr := string(textoRaw)

			desc1 := parseImageSection(textoStr, "Imagem 1")
			desc2 := parseImageSection(textoStr, "Imagem 2")
			if desc1 == "" {
				desc1 = "Abstract tech illustration for LinkedIn post slide 1"
				log.Printf("[WARN] seção 'Imagem 1' não encontrada em texto.md — usando fallback")
			}
			if desc2 == "" {
				desc2 = "Abstract tech illustration for LinkedIn post slide 2"
				log.Printf("[WARN] seção 'Imagem 2' não encontrada em texto.md — usando fallback")
			}

			// ── 3. Ler post final para contexto ──────────────────────────
			postContext := ""
			if data, err := os.ReadFile(workspace.OutputPostPath(postPath)); err == nil {
				postContext = strings.TrimSpace(string(data))
			} else if finalPath := workspace.WorkPath(postPath) + "/editor-final.md"; func() bool {
				_, e := os.Stat(finalPath)
				return e == nil
			}() {
				if data, err := os.ReadFile(workspace.WorkPath(postPath) + "/editor-final.md"); err == nil {
					postContext = strings.TrimSpace(string(data))
				}
			}

			gc := groq.NewGroqClient(cfg.GroqAPIKey)

			// ── 4. Loop principal (suporta regeneração) ───────────────────
			for {
				// ── 4a. Checar imagens de input direto ────────────────────
				var raw1, raw2 string
				if flagSkipGeneration {
					raw1 = findInputImage(inputPath, "imagem1")
					raw2 = findInputImage(inputPath, "imagem2")
					if raw1 == "" || raw2 == "" {
						return fmt.Errorf("--skip-generation requer input/imagem1.* e input/imagem2.*")
					}
					log.Printf("[INFO] usando imagens de input: %s, %s", filepath.Base(raw1), filepath.Base(raw2))
				} else {
					// ── 4b. Construir prompts via Groq ────────────────────
					log.Printf("[INFO] gerando prompts Gemini via Groq...")
					prompt1, err := buildGeminiPrompt(gc, postContext, desc1)
					if err != nil {
						log.Printf("[WARN] prompt slide1: %v — usando descrição original", err)
						prompt1 = desc1
					}
					prompt2, err := buildGeminiPrompt(gc, postContext, desc2)
					if err != nil {
						log.Printf("[WARN] prompt slide2: %v — usando descrição original", err)
						prompt2 = desc2
					}

					apf := artPromptsFile{
						Slide1: slidePrompts{DescricaoOriginal: desc1, PromptGemini: prompt1},
						Slide2: slidePrompts{DescricaoOriginal: desc2, PromptGemini: prompt2},
					}
					apfData, _ := json.MarshalIndent(apf, "", "  ")
					if err := os.WriteFile(workspace.ArtPromptsPath(postPath), apfData, 0644); err != nil {
						log.Printf("[WARN] salvar art-prompts.json: %v", err)
					}
					log.Printf("[INFO] art-prompts.json salvo")

					if flagDryRun {
						fmt.Printf("=== DRY-RUN — ldpost-art ===\n")
						fmt.Printf("post:    %s\n", flagPost)
						fmt.Printf("modelo:  %s\n", gemini.ImageModelV2)
						fmt.Printf("crop:    bottom=%dpx right=%dpx\n", flagCropBottom, flagCropRight)
						fmt.Printf("\nSlide 1 — descrição: %s\n", desc1)
						fmt.Printf("Slide 1 — prompt:    %s\n\n", prompt1)
						fmt.Printf("Slide 2 — descrição: %s\n", desc2)
						fmt.Printf("Slide 2 — prompt:    %s\n", prompt2)
						return nil
					}

					// ── 4c. Gerar imagens via Gemini ──────────────────────
					gm := gemini.New(cfg.GeminiAPIKey)
					imageModel := gemini.ImageModelV2

					log.Printf("[INFO] gerando slide1 via Gemini...")
					img1Bytes, err := gm.GenerateImageSquare(context.Background(), imageModel, prompt1)
					if err != nil {
						log.Printf("[WARN] gemini slide1: %v", err)
						// fallback para input/imagem1.*
						raw1 = findInputImage(inputPath, "imagem1")
						if raw1 == "" {
							return fmt.Errorf("falha ao gerar slide1 e sem fallback em input/imagem1.*: %w", err)
						}
						log.Printf("[INFO] usando fallback: %s", filepath.Base(raw1))
					} else {
						raw1 = filepath.Join(artRawPath, "slide1-raw.png")
						if err := os.WriteFile(raw1, img1Bytes, 0644); err != nil {
							return fmt.Errorf("salvar slide1-raw.png: %w", err)
						}
					}

					log.Printf("[INFO] gerando slide2 via Gemini...")
					img2Bytes, err := gm.GenerateImageSquare(context.Background(), imageModel, prompt2)
					if err != nil {
						log.Printf("[WARN] gemini slide2: %v", err)
						raw2 = findInputImage(inputPath, "imagem2")
						if raw2 == "" {
							return fmt.Errorf("falha ao gerar slide2 e sem fallback em input/imagem2.*: %w", err)
						}
						log.Printf("[INFO] usando fallback: %s", filepath.Base(raw2))
					} else {
						raw2 = filepath.Join(artRawPath, "slide2-raw.png")
						if err := os.WriteFile(raw2, img2Bytes, 0644); err != nil {
							return fmt.Errorf("salvar slide2-raw.png: %w", err)
						}
					}
				}

				// ── 4d. Crop 2 variantes por imagem ──────────────────────
				log.Printf("[INFO] gerando variantes de crop (bottom=%dpx, right=%dpx)...", flagCropBottom, flagCropRight)
				s1bottom, s1right, err := saveVariants(raw1, artRawPath, "slide1", flagCropBottom, flagCropRight)
				if err != nil {
					return fmt.Errorf("crop slide1: %w", err)
				}
				s2bottom, s2right, err := saveVariants(raw2, artRawPath, "slide2", flagCropBottom, flagCropRight)
				if err != nil {
					return fmt.Errorf("crop slide2: %w", err)
				}
				log.Printf("[INFO] variantes salvas em %s", artRawPath)

				// ── 4e. Telegram: enviar preview e aguardar escolha ───────
				if cfg.TelegramBotToken == "" || cfg.TelegramChatID == "" {
					// Auto-approve first variant
					log.Printf("[INFO] Telegram não configurado — aprovando variante 'bottom' automaticamente")
					if err := copyFile(s1bottom, filepath.Join(outputPath, "slide1.png")); err != nil {
						return err
					}
					if err := copyFile(s2bottom, filepath.Join(outputPath, "slide2.png")); err != nil {
						return err
					}
					break
				}

				bot := telegram.NewBot(cfg.TelegramBotToken, cfg.TelegramChatID)
				if err := sendArtPreview(bot, flagPost, s1bottom, s1right, s2bottom, s2right); err != nil {
					log.Printf("[WARN] enviar preview: %v", err)
				}

				action, choices := waitChoices(bot)
				if action == "regenerar" {
					if flagSkipGeneration {
						return fmt.Errorf("--skip-generation ativo: não é possível regenerar")
					}
					log.Printf("[INFO] regenerando imagens...")
					bot.SendMessage("🔄 Regenerando imagens...")
					continue // restart loop
				}

				// Copiar variantes escolhidas para output/
				slide1src := s1bottom
				if choices.Slide1 == "right" {
					slide1src = s1right
				}
				slide2src := s2bottom
				if choices.Slide2 == "right" {
					slide2src = s2right
				}
				if err := copyFile(slide1src, filepath.Join(outputPath, "slide1.png")); err != nil {
					return err
				}
				if err := copyFile(slide2src, filepath.Join(outputPath, "slide2.png")); err != nil {
					return err
				}

				bot.SendMessage(fmt.Sprintf(
					"✅ Imagens salvas em output/\n\nPróximo: <code>ldpost-director --post %s</code>",
					flagPost))
				break
			}

			// ── 5. Atualizar state ────────────────────────────────────────
			s.Status = state.StatusWaitingDirector
			s.SetEtapa("art", state.EtapaDone)
			s.SetEtapa("director", state.EtapaWaiting)
			if err := state.SaveState(postPath, s); err != nil {
				return fmt.Errorf("state: %w", err)
			}

			fmt.Printf("✅ Arte concluída.\n")
			fmt.Printf("   output/slide1.png e output/slide2.png\n")
			fmt.Printf("   Próximo: ldpost-director --post %s\n", flagPost)
			return nil
		},
	}

	root.Flags().StringVar(&flagPost, "post", "", "Slug do post (obrigatório)")
	root.Flags().StringVar(&flagWorkspace, "workspace", "", "Override de LDPOST_WORKSPACE")
	root.Flags().BoolVar(&flagDryRun, "dry-run", false, "Mostra prompts sem chamar Gemini")
	root.Flags().BoolVar(&flagSkipGeneration, "skip-generation", false, "Usa input/imagem1.* diretamente")
	root.Flags().IntVar(&flagCropBottom, "crop-bottom", 48, "Pixels a cortar de baixo (override LDPOST_CROP_BOTTOM)")
	root.Flags().IntVar(&flagCropRight, "crop-right", 48, "Pixels a cortar da direita (override LDPOST_CROP_RIGHT)")

	if err := root.Execute(); err != nil {
		os.Exit(1)
	}
}

// ─── Helpers ──────────────────────────────────────────────────────────────────

func copyFile(src, dst string) error {
	data, err := os.ReadFile(src)
	if err != nil {
		return fmt.Errorf("ler %s: %w", src, err)
	}
	if err := os.WriteFile(dst, data, 0644); err != nil {
		return fmt.Errorf("escrever %s: %w", dst, err)
	}
	return nil
}