Fix n7: honor comic page and panel counts

5 files changed, 285 insertions, 48 deletions

diff --git a/internal/comic/artist.go b/internal/comic/artist.go
index fd5dc60..248e8c4 100644
--- a/internal/comic/artist.go
+++ b/internal/comic/artist.go
@@ -56,9 +56,9 @@ func NewArtist(cfg *ArtistConfig) *Artist {
 		outputDir:      ".",
 		language:       "Bulgarian",
 		script:         "Cyrillic",
-		storyPages:     storyPagesInScript,
+		storyPages:     defaultStoryPagesInScript,
 		galleryPages:   5,
-		panelsPerPage:  storyPanelsPerPage,
+		panelsPerPage:  defaultStoryPanelsPerPage,
 		ultraRealistic: true,
 	}
 	if cfg == nil {
@@ -294,7 +294,11 @@ func (a *Artist) storyPagePromptData(section string, pageNum int, style, bible s
 		"Words":                   buildWordList(entries, ""),
 		"PageNum":                 pageNum,
 		"TotalPages":              a.storyPages,
-		"PanelLayout":             buildPanelLayout(section, pageScriptForPage(panelScript, pageNum-1)),
+		"PanelsPerPage":           a.panelsPerPage,
+		"RequiredDialoguePanels":  requiredDialoguePanels(a.panelsPerPage),
+		"TotalPanels":             a.storyPages * a.panelsPerPage,
+		"PanelLabelsText":         panelLabelsText(a.panelsPerPage),
+		"PanelLayout":             buildPanelLayout(section, pageScriptForPage(panelScript, pageNum-1), a.panelsPerPage),
 		"RenderingRequirement":    a.renderingRequirement(),
 		"RenderingRequirementEnd": a.renderingRequirementEnd(),
 	}
@@ -357,15 +361,17 @@ func pageScriptForPage(panelScript [][]string, idx int) []string {
 	return panelScript[idx]
 }
 
-func buildPanelLayout(section string, pagePanels []string) string {
-	if len(pagePanels) == 4 && pagePanels[0] != "" && pagePanels[1] != "" && pagePanels[2] != "" && pagePanels[3] != "" {
+func buildPanelLayout(section string, pagePanels []string, panelCount int) string {
+	panelCount = normalizePositive(panelCount, defaultStoryPanelsPerPage)
+	if len(pagePanels) == panelCount && allPanelsPresent(pagePanels) {
 		var sb strings.Builder
-		sb.WriteString("Divide the image into exactly 4 distinct panels in a 2x2 grid. The panels must tell the scene in sequence and remain visually distinct from each other.\n")
+		sb.WriteString(panelLayoutLead(panelCount))
+		sb.WriteString(" The panels must tell the scene in sequence and remain visually distinct from each other.\n")
 		for i, panel := range pagePanels {
 			if panel == "" {
 				continue
 			}
-			fmt.Fprintf(&sb, "Panel %d: %s\n", i+1, panel)
+			fmt.Fprintf(&sb, "Panel %s: %s\n", panelLabel(i), panel)
 		}
 		return sb.String()
 	}
@@ -378,10 +384,34 @@ func buildPanelLayout(section string, pagePanels []string) string {
 		}
 		excerpt += "…"
 	}
-	return "Divide the image into exactly 4 distinct panels in a 2x2 grid. The panels must tell the story in sequence from beginning to end and remain visually distinct.\n" +
+	return panelLayoutLead(panelCount) + " The panels must tell the story in sequence from beginning to end and remain visually distinct.\n" +
 		"Story excerpt:\n\n" + excerpt + "\n"
 }
 
+func allPanelsPresent(pagePanels []string) bool {
+	for _, panel := range pagePanels {
+		if strings.TrimSpace(panel) == "" {
+			return false
+		}
+	}
+	return true
+}
+
+func panelLayoutLead(panelCount int) string {
+	switch panelCount {
+	case 1:
+		return "Divide the image into exactly 1 distinct panel."
+	case 2:
+		return "Divide the image into exactly 2 distinct panels in a balanced two-panel layout."
+	case 3:
+		return "Divide the image into exactly 3 distinct panels in a balanced three-panel layout."
+	case 4:
+		return "Divide the image into exactly 4 distinct panels in a 2x2 grid."
+	default:
+		return fmt.Sprintf("Divide the image into exactly %d distinct panels in a balanced grid.", panelCount)
+	}
+}
+
 func blurbBoxInstruction(blurb string) string {
 	if strings.TrimSpace(blurb) == "" {
 		return "a rectangular text box near the bottom with a white or cream background and a thin black border, styled like a classic back-cover synopsis box"
diff --git a/internal/comic/comic_test.go b/internal/comic/comic_test.go
index ca1ef71..5502bc6 100644
--- a/internal/comic/comic_test.go
+++ b/internal/comic/comic_test.go
@@ -46,12 +46,51 @@ func TestParseGenerateResult(t *testing.T) {
 func TestBuildPanelLayoutUsesFallbackExcerpt(t *testing.T) {
 	t.Parallel()
 
-	got := buildPanelLayout("one two three four five", nil)
+	got := buildPanelLayout("one two three four five", nil, 2)
+	if !strings.Contains(got, "exactly 2 distinct panels") {
+		t.Fatalf("buildPanelLayout() = %q, want 2-panel layout instruction", got)
+	}
 	if !strings.Contains(got, "Story excerpt") {
 		t.Fatalf("buildPanelLayout() = %q", got)
 	}
 }
 
+func TestParseGenerateResultUsesConfiguredDimensions(t *testing.T) {
+	t.Parallel()
+
+	lines := []string{
+		"story text",
+		storyBibleSeparator,
+		"bible text",
+		storyTitleSeparator,
+		"My Comic",
+		storyPanelSeparator,
+		"P1-A: first",
+		"P1-B: second",
+		"P7-A: last",
+		"P7-B: end",
+		"P7-C: ignored",
+	}
+	got := parseGenerateResultWithDimensions(strings.Join(lines, "\n"), 7, 2)
+	if got.StoryText != "story text" || got.Bible != "bible text" || got.Title != "My Comic" {
+		t.Fatalf("parseGenerateResultWithDimensions() = %#v", got)
+	}
+	if got.PanelScript == nil || len(got.PanelScript) != 7 {
+		t.Fatalf("panel script pages = %d, want 7", len(got.PanelScript))
+	}
+	for i, page := range got.PanelScript {
+		if len(page) != 2 {
+			t.Fatalf("page %d panels = %d, want 2", i+1, len(page))
+		}
+	}
+	if got.PanelScript[0][0] != "first" || got.PanelScript[0][1] != "second" {
+		t.Fatalf("first page panel script = %#v", got.PanelScript[0])
+	}
+	if got.PanelScript[6][0] != "last" || got.PanelScript[6][1] != "end" {
+		t.Fatalf("last page panel script = %#v", got.PanelScript[6])
+	}
+}
+
 func TestComicOutputDirDoesNotDuplicateComicsSegment(t *testing.T) {
 	t.Parallel()
 
@@ -133,6 +172,78 @@ func TestGeneratorGenerateFull(t *testing.T) {
 	}
 }
 
+func TestGeneratorGenerateFullUsesConfiguredDimensions(t *testing.T) {
+	t.Parallel()
+
+	renderer := &recordingPromptRenderer{}
+	text := strings.Join([]string{
+		"история",
+		storyBibleSeparator,
+		"библия",
+		storyTitleSeparator,
+		"Заглавие",
+		storyPanelSeparator,
+		"P1-A: а",
+		"P1-B: б",
+		"P2-A: в",
+		"P2-B: г",
+		"P3-A: д",
+		"P3-B: е",
+		"P4-A: ж",
+		"P4-B: з",
+		"P5-A: и",
+		"P5-B: й",
+		"P6-A: к",
+		"P6-B: л",
+		"P7-A: м",
+		"P7-B: н",
+	}, "\n")
+
+	generator := NewGenerator(&GeneratorConfig{
+		TextProvider:  fakeTextProvider{text: text},
+		Prompts:       renderer,
+		StoryPages:    7,
+		PanelsPerPage: 2,
+	})
+	got, err := generator.GenerateFull(context.Background(), []WordEntry{{Word: "ябълка"}})
+	if err != nil {
+		t.Fatalf("GenerateFull() error = %v", err)
+	}
+	if len(got.PanelScript) != 7 {
+		t.Fatalf("panel script pages = %d, want 7", len(got.PanelScript))
+	}
+	for i, page := range got.PanelScript {
+		if len(page) != 2 {
+			t.Fatalf("page %d panels = %d, want 2", i+1, len(page))
+		}
+	}
+	var promptData map[string]any
+	for _, call := range renderer.calls {
+		if call.name == storyFullPromptTemplate {
+			promptData = call.data
+			break
+		}
+	}
+	if promptData == nil {
+		t.Fatal("story full prompt render was not recorded")
+	}
+	if got, want := promptData["StoryPages"], 7; got != want {
+		t.Fatalf("StoryPages prompt data = %#v, want %d", got, want)
+	}
+	if got, want := promptData["PanelsPerPage"], 2; got != want {
+		t.Fatalf("PanelsPerPage prompt data = %#v, want %d", got, want)
+	}
+	if got, want := promptData["TotalPanels"], 14; got != want {
+		t.Fatalf("TotalPanels prompt data = %#v, want %d", got, want)
+	}
+	if got, want := promptData["PanelLabelsText"], "A or B"; got != want {
+		t.Fatalf("PanelLabelsText prompt data = %#v, want %q", got, want)
+	}
+	if got, want := promptData["RequiredDialoguePanels"], 1; got != want {
+		t.Fatalf("RequiredDialoguePanels prompt data = %#v, want %d", got, want)
+	}
+}
+
 func TestDrawComicPagesReturnsErrorWhenRenderFails(t *testing.T) {
 	originalSleep := sleep
 	sleep = func(time.Duration) {}
@@ -262,6 +373,7 @@ func TestDrawComicPagesUsesOneStyleAcrossTheWholePDF(t *testing.T) {
 		Prompts:        renderer,
 		OutputDir:      t.TempDir(),
 		UltraRealistic: false,
+		PanelsPerPage:  2,
 	})
 
 	if _, err := artist.DrawComicPages(context.Background(), "история", "библия", "slug", []WordEntry{{Word: "ябълка"}}, nil); err != nil {
@@ -290,6 +402,28 @@ func TestDrawComicPagesUsesOneStyleAcrossTheWholePDF(t *testing.T) {
 	if promptCount == 0 {
 		t.Fatal("no image prompts recorded")
 	}
+	var pagePrompt map[string]any
+	for _, call := range renderer.calls {
+		if call.name == storyPagePromptTemplate {
+			pagePrompt = call.data
+			break
+		}
+	}
+	if pagePrompt == nil {
+		t.Fatal("story page prompt was not recorded")
+	}
+	if got, want := pagePrompt["PanelsPerPage"], 2; got != want {
+		t.Fatalf("PanelsPerPage prompt data = %#v, want %d", got, want)
+	}
+	if got, want := pagePrompt["RequiredDialoguePanels"], 1; got != want {
+		t.Fatalf("RequiredDialoguePanels prompt data = %#v, want %d", got, want)
+	}
+	if got, want := pagePrompt["PanelLabelsText"], "A or B"; got != want {
+		t.Fatalf("PanelLabelsText prompt data = %#v, want %q", got, want)
+	}
+	if layout, ok := pagePrompt["PanelLayout"].(string); !ok || !strings.Contains(layout, "exactly 2 distinct panels") {
+		t.Fatalf("PanelLayout prompt data = %#v, want 2-panel layout", pagePrompt["PanelLayout"])
+	}
 }
 
 func TestDrawComicPagesChainsReferenceImages(t *testing.T) {
diff --git a/internal/comic/generator.go b/internal/comic/generator.go
index 89591c8..8a13c0d 100644
--- a/internal/comic/generator.go
+++ b/internal/comic/generator.go
@@ -11,23 +11,27 @@ import (
 
 // GeneratorConfig configures story generation.
 type GeneratorConfig struct {
-	TextProvider provider.TextProvider
-	Prompts      PromptRenderer
-	Language     string
-	Script       string
-	Theme        string
-	Genres       []string
+	TextProvider  provider.TextProvider
+	Prompts       PromptRenderer
+	Language      string
+	Script        string
+	Theme         string
+	Genres        []string
+	StoryPages    int
+	PanelsPerPage int
 }
 
 // Generator produces the comic story text, bible, title, and panel script.
 type Generator struct {
-	textProvider provider.TextProvider
-	prompts      PromptRenderer
-	language     string
-	script       string
-	theme        string
-	genres       []string
-	initErr      error
+	textProvider  provider.TextProvider
+	prompts       PromptRenderer
+	language      string
+	script        string
+	theme         string
+	genres        []string
+	storyPages    int
+	panelsPerPage int
+	initErr       error
 }
 
 var _ = (*Generator)(nil)
@@ -35,9 +39,11 @@ var _ = (*Generator)(nil)
 // NewGenerator creates a new story generator.
 func NewGenerator(cfg *GeneratorConfig) *Generator {
 	g := &Generator{
-		language: "Bulgarian",
-		script:   "Cyrillic",
-		genres:   defaultStoryGenres,
+		language:      "Bulgarian",
+		script:        "Cyrillic",
+		genres:        defaultStoryGenres,
+		storyPages:    defaultStoryPagesInScript,
+		panelsPerPage: defaultStoryPanelsPerPage,
 	}
 	if cfg == nil {
 		g.initErr = errors.New("generator config is required")
@@ -48,6 +54,12 @@ func NewGenerator(cfg *GeneratorConfig) *Generator {
 	g.language = orDefault(cfg.Language, g.language)
 	g.script = orDefault(cfg.Script, g.script)
 	g.theme = cfg.Theme
+	if cfg.StoryPages > 0 {
+		g.storyPages = cfg.StoryPages
+	}
+	if cfg.PanelsPerPage > 0 {
+		g.panelsPerPage = cfg.PanelsPerPage
+	}
 	if len(cfg.Genres) > 0 {
 		g.genres = append([]string(nil), cfg.Genres...)
 	}
@@ -107,7 +119,7 @@ func (g *Generator) GenerateFull(ctx context.Context, entries []WordEntry) (Gene
 	if text == "" {
 		return GenerateResult{}, fmt.Errorf("no content returned")
 	}
-	result := parseGenerateResult(text)
+	result := parseGenerateResultWithDimensions(text, g.storyPages, g.panelsPerPage)
 	if err := validateGeneratedResult(result, g.script); err != nil {
 		return GenerateResult{}, err
 	}
@@ -127,15 +139,20 @@ func (g *Generator) ready() error {
 func (g *Generator) renderStoryPrompt(templateName string, entries []WordEntry) (string, error) {
 	genre := resolveGenre(g.theme, g.genres)
 	data := map[string]any{
-		"Language":            g.language,
-		"LanguageName":        localizedLanguageName(g.language, g.script),
-		"Script":              g.script,
-		"ScriptName":          localizedScriptName(g.script),
-		"Genre":               genre,
-		"Words":               buildWordList(entries, ""),
-		"StoryBibleSeparator": storyBibleSeparator,
-		"StoryTitleSeparator": storyTitleSeparator,
-		"StoryPanelSeparator": storyPanelSeparator,
+		"Language":               g.language,
+		"LanguageName":           localizedLanguageName(g.language, g.script),
+		"Script":                 g.script,
+		"ScriptName":             localizedScriptName(g.script),
+		"Genre":                  genre,
+		"Words":                  buildWordList(entries, ""),
+		"StoryPages":             g.storyPages,
+		"PanelsPerPage":          g.panelsPerPage,
+		"TotalPanels":            g.storyPages * g.panelsPerPage,
+		"PanelLabelsText":        panelLabelsText(g.panelsPerPage),
+		"RequiredDialoguePanels": requiredDialoguePanels(g.panelsPerPage),
+		"StoryBibleSeparator":    storyBibleSeparator,
+		"StoryTitleSeparator":    storyTitleSeparator,
+		"StoryPanelSeparator":    storyPanelSeparator,
 	}
 
 	systemPrompt, err := g.prompts.RenderPrompt(storySystemPromptTemplate, map[string]any{
diff --git a/internal/comic/runner.go b/internal/comic/runner.go
index 00275b8..0b47a52 100644
--- a/internal/comic/runner.go
+++ b/internal/comic/runner.go
@@ -64,11 +64,13 @@ func NewRunner(cfg *RunnerConfig) *Runner {
 	}
 
 	r.generator = NewGenerator(&GeneratorConfig{
-		TextProvider: cfg.TextProvider,
-		Prompts:      cfg.Prompts,
-		Language:     cfg.Language,
-		Script:       cfg.Script,
-		Theme:        cfg.Theme,
+		TextProvider:  cfg.TextProvider,
+		Prompts:       cfg.Prompts,
+		Language:      cfg.Language,
+		Script:        cfg.Script,
+		Theme:         cfg.Theme,
+		StoryPages:    cfg.StoryPages,
+		PanelsPerPage: cfg.PanelsPerPage,
 	})
 	r.artist = NewArtist(&ArtistConfig{
 		ImageProvider:  cfg.ImageProvider,
diff --git a/internal/comic/types.go b/internal/comic/types.go
index 0bee759..f1b609c 100644
--- a/internal/comic/types.go
+++ b/internal/comic/types.go
@@ -33,8 +33,8 @@ const (
 	storyTitleSeparator = "---COMIC TITLE---"
 	storyPanelSeparator = "---PANEL SCRIPT---"
 
-	storyPagesInScript = 5
-	storyPanelsPerPage = 4
+	defaultStoryPagesInScript = 5
+	defaultStoryPanelsPerPage = 4
 
 	pageMaxRetries       = 5
 	pageRetryBase        = 15 * time.Second
@@ -160,6 +160,10 @@ func pickStyle(styles []string, ultraRealistic bool) string {
 }
 
 func parseGenerateResult(combined string) GenerateResult {
+	return parseGenerateResultWithDimensions(combined, defaultStoryPagesInScript, defaultStoryPanelsPerPage)
+}
+
+func parseGenerateResultWithDimensions(combined string, storyPages, panelsPerPage int) GenerateResult {
 	bibleIdx := strings.Index(combined, storyBibleSeparator)
 	if bibleIdx < 0 {
 		return GenerateResult{StoryText: strings.TrimSpace(combined)}
@@ -192,17 +196,23 @@ func parseGenerateResult(combined string) GenerateResult {
 		StoryText:   story,
 		Bible:       bible,
 		Title:       title,
-		PanelScript: parsePanelScript(panelText),
+		PanelScript: parsePanelScript(panelText, storyPages, panelsPerPage),
 	}
 }
 
-func parsePanelScript(text string) [][]string {
-	script := make([][]string, storyPagesInScript)
+func parsePanelScript(text string, storyPages, panelsPerPage int) [][]string {
+	storyPages = normalizePositive(storyPages, defaultStoryPagesInScript)
+	panelsPerPage = normalizePositive(panelsPerPage, defaultStoryPanelsPerPage)
+
+	script := make([][]string, storyPages)
 	for i := range script {
-		script[i] = make([]string, storyPanelsPerPage)
+		script[i] = make([]string, panelsPerPage)
 	}
 
-	panelIndex := map[byte]int{'A': 0, 'B': 1, 'C': 2, 'D': 3}
+	panelIndex := make(map[byte]int, panelsPerPage)
+	for i := 0; i < panelsPerPage; i++ {
+		panelIndex[byte('A'+i)] = i
+	}
 	for _, line := range strings.Split(text, "\n") {
 		line = strings.TrimSpace(line)
 		if len(line) < 7 || line[0] != 'P' || line[2] != '-' || line[4] != ':' {
@@ -210,7 +220,7 @@ func parsePanelScript(text string) [][]string {
 		}
 		page := int(line[1] - '1')
 		panel, ok := panelIndex[line[3]]
-		if !ok || page < 0 || page >= storyPagesInScript {
+		if !ok || page < 0 || page >= storyPages {
 			continue
 		}
 		script[page][panel] = strings.TrimSpace(line[5:])
@@ -247,3 +257,47 @@ func min(a, b int) int {
 	}
 	return b
 }
+
+func normalizePositive(value, fallback int) int {
+	if value > 0 {
+		return value
+	}
+	return fallback
+}
+
+func panelLabel(idx int) string {
+	if idx < 0 {
+		return ""
+	}
+	return string(rune('A' + idx))
+}
+
+func panelLabelsText(count int) string {
+	count = normalizePositive(count, defaultStoryPanelsPerPage)
+	labels := make([]string, 0, count)
+	for i := 0; i < count; i++ {
+		labels = append(labels, panelLabel(i))
+	}
+	return joinEnglishList(labels)
+}
+
+func requiredDialoguePanels(count int) int {
+	count = normalizePositive(count, defaultStoryPanelsPerPage)
+	if count <= 1 {
+		return 1
+	}
+	return (count + 1) / 2
+}
+
+func joinEnglishList(items []string) string {
+	switch len(items) {
+	case 0:
+		return ""
+	case 1:
+		return items[0]
+	case 2:
+		return items[0] + " or " + items[1]
+	default:
+		return strings.Join(items[:len(items)-1], ", ") + ", or " + items[len(items)-1]
+	}
+}