#!/usr/bin/env python3
"""
concept_grid.py — Generic 2x3 NBP concept grid generator.

Freeform concepting tool for exploring visual variants of any subject
(vehicles, props, locations, abstract concepts). Generates a single 2x3
grid image via NBP (Gemini 3 Pro Image), splits into 6 panels, and
writes to assets/concepts/{slug}/.

Unlike prep_character_angles.py Path B (character-specific casting), this
tool uses a neutral diegetic frame and accepts any description — no
character/synthetic detection, no codename mapping, no gender anchoring.

For character concepting, prefer prep_character_angles.py --description.

Usage:
    python -m tools.concept_grid \\
        --project driver-beware \\
        --slug new_car_pool \\
        --kind object \\
        --description "Compact 2D cartoon cars, varied silhouettes..."

    python -m tools.concept_grid \\
        --project driver-beware \\
        --slug rural_sunny \\
        --kind location \\
        --aspect 16:9 \\
        --description "Rural two-lane country road, cornfields, barns..."

Cost: ~$0.134 per grid (single NBP call).
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Optional

_PROJECT_ROOT = Path(__file__).parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from recoil.core.paths import ProjectPaths
from recoil.core.model_profiles import get_model
from tools.prep_character_angles import _split_grid_into_panels

logger = logging.getLogger("pipeline.concept_grid")

NBP_COST = 0.134

# Diegetic frames by subject kind. Controls how NBP frames the grid.
DIEGETIC_FRAMES = {
    "object": (
        "A production design reference sheet. 6 DIFFERENT design interpretations "
        "of the same subject category. Isolated on a flat 18% neutral gray "
        "background. Studio lighting, macro lens. No borders, no frames."
    ),
    "location": (
        "An environment concept art sheet. 6 DIFFERENT establishing-shot takes "
        "on the same location, each showing a distinct vantage, composition, "
        "or time-of-day interpretation. No characters. No borders, no frames."
    ),
    "freeform": (
        "A visual concept reference sheet. 6 DIFFERENT interpretations of the "
        "same description. Clean layout, no borders, no frames, no sprocket holes."
    ),
}


def generate_concept_grid(
    slug: str,
    description: str,
    project: str,
    kind: str = "freeform",
    aspect_ratio: str = "2:3",
    dry_run: bool = False,
) -> dict:
    """Generate a 2x3 NBP concept grid and split into panels.

    Args:
        slug: short identifier for the concept pool (e.g. "new_car_pool")
        description: freeform text description of the subject
        project: project name (used for output path)
        kind: "object" | "location" | "freeform" — picks diegetic frame
        aspect_ratio: grid aspect ratio (2:3 portrait, 3:2 / 16:9 landscape)
        dry_run: if True, print the prompt and output path and exit

    Returns:
        dict with grid_path, panels (list of paths), cost.
    """
    out_dir = ProjectPaths.for_project(project).assets_dir / "concepts" / slug
    out_dir.mkdir(parents=True, exist_ok=True)

    diegetic = DIEGETIC_FRAMES.get(kind, DIEGETIC_FRAMES["freeform"])

    prompt = (
        f"CRITICAL DIRECTIVE: Generate a single image containing a 2x3 grid "
        f"(3 columns by 2 rows) of distinct panels.\n"
        f"DIEGETIC FRAMING: {diegetic}\n\n"
        f"SUBJECT: {description}\n\n"
        f"GRID STRUCTURE: 3 columns by 2 rows. Strictly isolated panels, "
        f"no overlapping elements. Each of the 6 panels must feature a DIFFERENT "
        f"interpretation of the subject description. Vary design, proportions, "
        f"color, and styling — but keep the core category consistent so all six "
        f"read as the same type of thing.\n\n"
        f"NEGATIVE CONSTRAINTS: no photorealism, no AI artifacts, no text overlays, "
        f"no watermarks, no logos, no borders, no film frames, no sprocket holes."
    )

    if dry_run:
        print(f"[DRY RUN] Would generate concept grid for slug={slug} kind={kind}")
        print(f"  Output dir: {out_dir}")
        print(f"  Aspect ratio: {aspect_ratio}")
        print(f"  Prompt:\n{prompt}")
        return {
            "slug": slug,
            "grid_path": None,
            "panels": [],
            "cost": 0.0,
            "dry_run": True,
        }

    from recoil.execution.api_client import get_client
    from google.genai import types as genai_types

    model_id = get_model("production_hq", "image")
    client = get_client(model_id)

    config = genai_types.GenerateContentConfig(
        temperature=0.4,
        response_modalities=["IMAGE", "TEXT"],
        image_config=genai_types.ImageConfig(aspect_ratio=aspect_ratio),
    )

    grid_path = out_dir / f"{slug}_concept_grid.png"

    try:
        api_client = client._get_client()
        response = api_client.models.generate_content(
            model=model_id,
            contents=prompt,
            config=config,
        )
    except Exception as e:
        logger.error("Concept grid generation failed for %s: %s", slug, e)
        return {"slug": slug, "grid_path": None, "panels": [], "cost": NBP_COST}

    if response and response.candidates:
        for candidate in response.candidates:
            if candidate.content and candidate.content.parts:
                for part in candidate.content.parts:
                    if hasattr(part, "inline_data") and part.inline_data:
                        grid_path.write_bytes(part.inline_data.data)
                        logger.info("Concept grid saved: %s", grid_path)
                        break

    if not grid_path.exists():
        logger.error("NBP returned no image data for slug=%s", slug)
        return {"slug": slug, "grid_path": None, "panels": [], "cost": NBP_COST}

    # Reuse the splitter from prep_character_angles (handles label-band detection).
    panels = _split_grid_into_panels(grid_path, out_dir, slug)

    return {
        "slug": slug,
        "grid_path": str(grid_path),
        "panels": [str(p) for p in panels],
        "cost": NBP_COST,
    }


def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )
    parser = argparse.ArgumentParser(description="Generic NBP concept grid generator")
    parser.add_argument("--project", required=True, help="Project name (e.g. driver-beware)")
    parser.add_argument("--slug", required=True, help="Short identifier (e.g. new_car_pool)")
    parser.add_argument("--description", required=True, help="Freeform text description")
    parser.add_argument(
        "--kind",
        default="freeform",
        choices=["object", "location", "freeform"],
        help="Diegetic frame: object, location, or freeform (default: freeform)",
    )
    parser.add_argument(
        "--aspect",
        default="2:3",
        help="Aspect ratio: 2:3 (portrait), 3:2, 16:9, 9:16, 1:1 (default: 2:3)",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print prompt and exit")
    args = parser.parse_args()

    result = generate_concept_grid(
        slug=args.slug,
        description=args.description,
        project=args.project,
        kind=args.kind,
        aspect_ratio=args.aspect,
        dry_run=args.dry_run,
    )

    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()