#!/usr/bin/env python3
"""
Extract prompts from awesome-nano-banana-pro-prompts repository

Parses README.md and extracts prompts by category for use as
reference examples in the video generation pipeline.
"""

import json
import re
import os
from pathlib import Path
from typing import Dict, List, Optional


# Categories relevant for product video generation
RELEVANT_CATEGORIES = [
    "product-marketing",
    "ecommerce-main-image",
    "photography",
    "cinematic-film-still",
    "portrait-selfie",
    "lifestyle",
    "product",
    "food-drink",
]


def parse_readme(readme_path: str) -> List[Dict]:
    """
    Parse README.md and extract all prompts

    Returns list of {id, title, description, prompt, images, author, categories}
    """
    with open(readme_path, 'r', encoding='utf-8') as f:
        content = f.read()

    prompts = []

    # Split by prompt headers (### No. X: Title)
    prompt_sections = re.split(r'(?=### No\. \d+:)', content)

    for section in prompt_sections:
        if not section.strip().startswith('### No.'):
            continue

        prompt_data = parse_prompt_section(section)
        if prompt_data:
            prompts.append(prompt_data)

    return prompts


def parse_prompt_section(section: str) -> Optional[Dict]:
    """
    Parse a single prompt section from README

    Returns dict with prompt data or None if parsing fails
    """
    try:
        # Extract ID and title
        title_match = re.search(r'### No\. (\d+): (.+)', section)
        if not title_match:
            return None

        prompt_id = title_match.group(1)
        title = title_match.group(2).strip()

        # Extract description (after #### 📖 Description)
        desc_match = re.search(r'#### 📖 Description\s*\n+(.+?)(?=\n####|\n```)', section, re.DOTALL)
        description = desc_match.group(1).strip() if desc_match else ""

        # Extract prompt (in code block after #### 📝 Prompt)
        prompt_match = re.search(r'#### 📝 Prompt\s*\n+```\n?([\s\S]+?)\n```', section)
        prompt_text = prompt_match.group(1).strip() if prompt_match else ""

        if not prompt_text:
            return None

        # Extract image URLs
        images = re.findall(r'<img src="([^"]+)"', section)

        # Extract author
        author_match = re.search(r'\*\*Author:\*\*\s*\[([^\]]+)\]', section)
        author = author_match.group(1) if author_match else "Unknown"

        # Detect categories from badges and content
        categories = detect_categories(title, description, prompt_text)

        return {
            "id": prompt_id,
            "title": title,
            "description": description[:500],  # Limit description length
            "prompt": prompt_text[:2000],  # Limit prompt length
            "images": images[:4],  # Limit images
            "author": author,
            "categories": categories
        }

    except Exception as e:
        print(f"Error parsing section: {e}")
        return None


def detect_categories(title: str, description: str, prompt: str) -> List[str]:
    """
    Detect categories from prompt content
    """
    text = f"{title} {description} {prompt}".lower()
    categories = []

    category_keywords = {
        "product-marketing": ["product", "commercial", "advertising", "marketing", "luxury", "premium"],
        "ecommerce-main-image": ["e-commerce", "ecommerce", "main image", "hero image", "product shot"],
        "photography": ["photography", "photorealistic", "photo", "camera", "lens", "aperture", "f/"],
        "cinematic-film-still": ["cinematic", "film", "movie", "scene", "dramatic lighting"],
        "portrait-selfie": ["portrait", "selfie", "face", "headshot", "close-up"],
        "lifestyle": ["lifestyle", "natural", "everyday", "casual", "authentic"],
        "product": ["product", "item", "object", "merchandise"],
        "food-drink": ["food", "drink", "beverage", "coffee", "tea", "meal", "dish"],
    }

    for category, keywords in category_keywords.items():
        if any(kw in text for kw in keywords):
            categories.append(category)

    return categories


def filter_relevant_prompts(prompts: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Filter and organize prompts by relevant categories
    """
    categorized = {cat: [] for cat in RELEVANT_CATEGORIES}

    for prompt in prompts:
        for category in prompt.get("categories", []):
            if category in categorized:
                # Create a slim version for storage
                slim_prompt = {
                    "id": prompt["id"],
                    "title": prompt["title"],
                    "prompt": prompt["prompt"],
                    "description": prompt["description"][:200],
                }
                categorized[category].append(slim_prompt)

    # Remove empty categories and limit prompts per category
    result = {}
    for category, prompts_list in categorized.items():
        if prompts_list:
            # Sort by ID (newer prompts first) and limit
            sorted_prompts = sorted(prompts_list, key=lambda x: int(x["id"]), reverse=True)
            result[category] = sorted_prompts[:20]  # Max 20 per category

    return result


def extract_and_save(readme_path: str, output_dir: str):
    """
    Extract prompts and save to JSON files
    """
    print(f"Parsing {readme_path}...")
    prompts = parse_readme(readme_path)
    print(f"Found {len(prompts)} prompts")

    # Filter and categorize
    categorized = filter_relevant_prompts(prompts)

    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Save each category to separate file
    for category, prompts_list in categorized.items():
        output_file = os.path.join(output_dir, f"{category.replace('-', '_')}.json")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(prompts_list, f, indent=2, ensure_ascii=False)
        print(f"Saved {len(prompts_list)} prompts to {output_file}")

    # Save combined file for quick reference
    combined_file = os.path.join(output_dir, "all_relevant.json")
    with open(combined_file, 'w', encoding='utf-8') as f:
        json.dump(categorized, f, indent=2, ensure_ascii=False)
    print(f"Saved combined file to {combined_file}")

    return categorized


def get_prompts_for_product_type(prompts_dir: str, product_type: str) -> List[Dict]:
    """
    Get relevant prompts for a product type

    Args:
        prompts_dir: Directory containing extracted prompt JSON files
        product_type: Type of product (e.g., "kitchenware", "fashion", "electronics")

    Returns:
        List of relevant prompts for Claude to use as reference
    """
    # Map product types to relevant prompt categories
    type_to_categories = {
        "home_appliances": ["product-marketing", "product", "photography"],
        "kitchenware": ["food-drink", "product", "photography", "lifestyle"],
        "fashion": ["portrait-selfie", "lifestyle", "photography", "cinematic-film-still"],
        "cosmetics": ["portrait-selfie", "product-marketing", "photography"],
        "electronics": ["product-marketing", "product", "photography"],
        "food": ["food-drink", "lifestyle", "photography"],
        "cleaning": ["product", "photography", "lifestyle"],
        "other": ["product-marketing", "product", "photography"],
    }

    categories = type_to_categories.get(product_type, type_to_categories["other"])

    relevant_prompts = []
    for category in categories:
        category_file = os.path.join(prompts_dir, f"{category.replace('-', '_')}.json")
        if os.path.exists(category_file):
            with open(category_file, 'r', encoding='utf-8') as f:
                prompts = json.load(f)
                # Take top 3 from each category
                relevant_prompts.extend(prompts[:3])

    # Limit total and deduplicate by ID
    seen_ids = set()
    unique_prompts = []
    for p in relevant_prompts:
        if p["id"] not in seen_ids:
            seen_ids.add(p["id"])
            unique_prompts.append(p)

    return unique_prompts[:10]  # Max 10 reference prompts


if __name__ == "__main__":
    import sys

    # Default paths
    default_readme = os.path.expanduser("~/awesome-nano-banana-pro-prompts/README.md")
    default_output = os.path.join(os.path.dirname(__file__), "..", "prompts", "awesome")

    if len(sys.argv) >= 2:
        readme_path = sys.argv[1]
    else:
        readme_path = default_readme

    if len(sys.argv) >= 3:
        output_dir = sys.argv[2]
    else:
        output_dir = default_output

    if not os.path.exists(readme_path):
        print(f"README not found: {readme_path}")
        sys.exit(1)

    extract_and_save(readme_path, output_dir)
    print("\nDone! Prompts extracted and saved.")
