#!/usr/bin/env python3
"""
preview.py — universal media preview tool.

Write anything. The script figures out what to make:
- video idea / product ad → full scenario + storyboard grid image
- portrait trend ("аниме", "я в средневековье") → professional portrait prompt + image
- any image idea → direct image generation

Usage:
    python story-gen/preview.py "котёнок и снег"
    python story-gen/preview.py "реклама крема для рук WB"
    python story-gen/preview.py "сделай меня в стиле аниме"
    python story-gen/preview.py "закат над горами в стиле oil painting"
    python story-gen/preview.py "идея" --scene 2
"""

import sys, os, json, argparse, base64, subprocess, tempfile
from urllib import request

if sys.stdout.encoding and sys.stdout.encoding.lower() not in ('utf-8', 'utf8'):
    sys.stdout.reconfigure(encoding='utf-8')

API_URL     = os.environ.get("OPENAI_BASE_URL", "https://llm.lambda.coredump.ru/v1")
API_KEY     = os.environ.get("OPENAI_API_KEY", "")
MODEL       = os.environ.get("STORY_MODEL", "qwen3.5-122b")
IMAGE_MODEL = os.environ.get("IMAGE_MODEL", "gpt-image-1")

# ---------------------------------------------------------------------------
# Step 1: classify the request
# ---------------------------------------------------------------------------

CLASSIFY_PROMPT = """You are a media production assistant. Classify the user request into one of four types:

- "wb_product": user wants to create a commercial video ad for a Wildberries product — they provide a WB URL, product name, description, photos, or any combination of product assets. Key signals: mentions WB/Wildberries, product URL (wildberries.ru), product specs, wants to sell something on a marketplace.
- "video": user wants a generic video — reel, TikTok, story, scenario — NOT specifically a WB product ad
- "trend": user wants a portrait transformation based on a popular visual trend (anime, medieval, doll, soviet era, studio photo, aged self, flowers in hair, etc.)
- "image": user wants a single image of anything that is NOT a portrait trend, NOT a video, NOT a WB product

Return ONLY valid JSON, no markdown:
{
  "type": "wb_product|video|trend|image",
  "reason": "one sentence explanation"
}
"""

def classify(user_input: str) -> str:
    """Returns 'video', 'trend', or 'image'."""
    payload = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": CLASSIFY_PROMPT},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.1
    }).encode()
    req = request.Request(
        f"{API_URL}/chat/completions",
        data=payload,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
    )
    with request.urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read())
    content = data["choices"][0]["message"]["content"].strip()
    if content.startswith("```"):
        content = content.split("```")[1]
        if content.startswith("json"):
            content = content[4:]
    result = json.loads(content.strip())
    return result["type"], result.get("reason", "")


# ---------------------------------------------------------------------------
# Step 2a: video — full scenario
# ---------------------------------------------------------------------------

def generate_scenario(user_input: str, fmt="auto", platform="auto") -> dict:
    scripts_dir = os.path.join(os.path.dirname(__file__), "scripts")
    sys.path.insert(0, scripts_dir)
    import generate as gen

    assets = gen.analyze_assets(user_input)
    if fmt == "auto" and assets.get("recommended_format"):
        fmt = assets["recommended_format"]
    if platform == "auto" and assets.get("recommended_platform"):
        platform = assets["recommended_platform"]

    scenario = gen.generate(user_input, fmt, platform, assets=assets)
    scenario["asset_analysis"] = assets
    return scenario


# ---------------------------------------------------------------------------
# Step 2b: trend — portrait prompt expansion (file-based + LLM fallback)
# ---------------------------------------------------------------------------

TRENDS_DIR = os.path.join(os.path.dirname(__file__), "scripts", "trends")

TREND_FALLBACK_PROMPT = """You are a professional prompt engineer for AI portrait image generation.

The user gives you a short phrase describing a visual style or trend.
Expand it into a detailed professional portrait prompt.

Structure:
Transform this photo into a human portrait. Use the uploaded photo for the face — preserve ALL features exactly: face shape, eyes, nose, lips, eyebrows, hair color, hairstyle.

Clothing: [detailed clothing]
Location: [detailed background/setting]
Pose & Action: [body position, gesture, eye direction]
Lighting: [lighting setup and atmosphere]
Mood: [emotional tone, keywords]
Technical: [lens, aperture, ISO, art style, quality tags]

Rules:
- Write entirely in English
- Be very specific and detailed
- Return ONLY the prompt text, no explanations
"""

def _load_trends() -> list[dict]:
    trends = []
    if not os.path.isdir(TRENDS_DIR):
        return trends
    for fname in os.listdir(TRENDS_DIR):
        if not fname.endswith(".txt"):
            continue
        with open(os.path.join(TRENDS_DIR, fname), encoding="utf-8") as f:
            content = f.read().strip()
        lines = content.splitlines()
        keywords = []
        prompt_lines = []
        for i, line in enumerate(lines):
            if line.startswith("keywords:"):
                keywords = [k.strip().lower() for k in line[len("keywords:"):].split(",") if k.strip()]
            else:
                prompt_lines = lines[i:]
                break
        trends.append({
            "name": fname.replace(".txt", ""),
            "keywords": keywords,
            "prompt": "\n".join(prompt_lines).strip()
        })
    return trends


def expand_trend(user_input: str) -> str:
    query = user_input.lower()
    for trend in _load_trends():
        if any(kw in query for kw in trend["keywords"]):
            print(f"  Matched trend: {trend['name']}", file=sys.stderr)
            return trend["prompt"]

    print("  No trend matched — using LLM to generate prompt", file=sys.stderr)
    payload = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": TREND_FALLBACK_PROMPT},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.7
    }).encode()
    req = request.Request(
        f"{API_URL}/chat/completions",
        data=payload,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
    )
    with request.urlopen(req, timeout=60) as resp:
        data = json.loads(resp.read())
    return data["choices"][0]["message"]["content"].strip()


# ---------------------------------------------------------------------------
# Step 2c: image — direct prompt expansion
# ---------------------------------------------------------------------------

IMAGE_PROMPT = """You are a professional prompt engineer for AI image generation.

The user gives you an idea for an image. Expand it into a detailed, high-quality generation prompt.

Rules:
- Write entirely in English
- Include: subject, style, lighting, colors, mood, composition, technical quality tags
- Be specific and detailed
- Return ONLY the prompt text, no explanations
"""

def expand_image_prompt(user_input: str) -> str:
    payload = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": IMAGE_PROMPT},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.7
    }).encode()
    req = request.Request(
        f"{API_URL}/chat/completions",
        data=payload,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
    )
    with request.urlopen(req, timeout=60) as resp:
        data = json.loads(resp.read())
    return data["choices"][0]["message"]["content"].strip()


# ---------------------------------------------------------------------------
# Image generation + display
# ---------------------------------------------------------------------------

def read_photo(photo_path: str) -> tuple[bytes, str, str]:
    """Read photo file, return (bytes, mime_type, filename)."""
    ext = os.path.splitext(photo_path)[1].lower()
    # .jfif is JPEG with different extension — API needs image/jpeg
    mime_map = {
        ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
        ".jfif": "image/jpeg", ".jpe": "image/jpeg",
        ".png": "image/png", ".webp": "image/webp",
    }
    mime = mime_map.get(ext, "image/jpeg")
    # Always send as .jpg filename so API doesn't reject on extension
    filename = os.path.splitext(os.path.basename(photo_path))[0] + ".jpg"
    with open(photo_path, "rb") as f:
        return f.read(), mime, filename


def generate_image(prompt: str, size: str, photo_path: str = None) -> bytes:
    """Generate image. If photo_path given — edit the photo using the prompt."""
    if photo_path:
        print("  Generating image with your photo...", file=sys.stderr)
        import uuid
        boundary = uuid.uuid4().hex

        def field(name, value):
            return (
                f"--{boundary}\r\n"
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'
                f"{value}\r\n"
            ).encode()

        photo_bytes, mime, filename = read_photo(photo_path)

        body = (
            f"--{boundary}\r\n"
            f'Content-Disposition: form-data; name="image"; filename="{filename}"\r\n'
            f"Content-Type: {mime}\r\n\r\n"
        ).encode() + photo_bytes + b"\r\n"
        body += field("prompt", prompt)
        body += field("model", IMAGE_MODEL)
        body += field("size", size)
        body += field("n", "1")
        body += f"--{boundary}--\r\n".encode()

        req = request.Request(
            f"{API_URL}/images/edits",
            data=body,
            headers={
                "Authorization": f"Bearer {API_KEY}",
                "Content-Type": f"multipart/form-data; boundary={boundary}"
            }
        )
    else:
        print("  Generating image...", file=sys.stderr)
        body = json.dumps({
            "model": IMAGE_MODEL,
            "prompt": prompt,
            "n": 1,
            "size": size
        }).encode()
        req = request.Request(
            f"{API_URL}/images/generations",
            data=body,
            headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
        )

    timeout = 240 if photo_path else 120
    with request.urlopen(req, timeout=timeout) as resp:
        data = json.loads(resp.read())
    b64 = data["data"][0].get("b64_json")
    if not b64:
        raise ValueError("API returned no image data")
    return base64.b64decode(b64)


def open_image(path: str):
    if sys.platform == "win32":
        os.startfile(path)
    elif sys.platform == "darwin":
        subprocess.run(["open", path])
    else:
        subprocess.run(["xdg-open", path])


def save_and_open(image_bytes: bytes, label: str) -> str:
    out_dir = os.path.join(tempfile.gettempdir(), "story-gen-previews")
    os.makedirs(out_dir, exist_ok=True)
    safe = "".join(c if c.isalnum() or c in "-_ " else "_" for c in label)[:50].strip()
    path = os.path.join(out_dir, f"{safe}.png")
    with open(path, "wb") as f:
        f.write(image_bytes)
    print(f"  Saved: {path}", file=sys.stderr)
    open_image(path)
    return path


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Universal media preview")
    parser.add_argument("input", help="Any idea — video, trend, or image")
    parser.add_argument("--format",   default="auto",
                        choices=["auto","wb_ad","reels","viral","long","postcard","educational"])
    parser.add_argument("--platform", default="auto",
                        choices=["auto","tiktok","instagram","wb","youtube","vk"])
    parser.add_argument("--scene",    type=int, default=None,
                        help="For video: preview a specific scene number instead of the grid")
    parser.add_argument("--size",     default="1024x1024",
                        choices=["1024x1024", "1792x1024", "1024x1792"])
    parser.add_argument("--photo",    default=None,
                        help="Path to your photo (for portrait trends — face will be preserved)")
    parser.add_argument("--assets",   nargs="*", default=None,
                        help="Additional assets for WB product: photo paths, video paths, WB URL, text strings")
    parser.add_argument("--mode",     default="auto",
                        choices=["auto", "card", "video"],
                        help="WB product mode: card=generate infographic images, video=generate video scenario")
    parser.add_argument("--cards",    type=int, default=1,
                        help="Number of WB cards to generate (default: 1)")
    args = parser.parse_args()

    if not API_KEY:
        print("Error: OPENAI_API_KEY not set", file=sys.stderr)
        sys.exit(1)

    # Classify
    print(f"\nAnalyzing request...", file=sys.stderr)
    request_type, reason = classify(args.input)
    print(f"  Type: {request_type} — {reason}", file=sys.stderr)

    # --- WB PRODUCT → card infographics or video scenario ---
    if request_type == "wb_product":
        scripts_dir = os.path.join(os.path.dirname(__file__), "scripts")
        sys.path.insert(0, scripts_dir)
        import generate as gen

        # Build assets dict from CLI args
        assets = {}
        if args.photo:
            if os.path.exists(args.photo):
                assets["photo_1"] = {"type": "photo", "path": args.photo}
            else:
                assets["photo_1"] = {"type": "photo", "url": args.photo}
        if args.assets:
            for i, a in enumerate(args.assets, start=(2 if args.photo else 1)):
                ext = os.path.splitext(a)[1].lower()
                if ext in (".mp4", ".mov", ".avi", ".webm"):
                    assets[f"video_{i}"] = {"type": "video", "path": a}
                elif ext in (".jpg", ".jpeg", ".png", ".webp", ".jfif"):
                    assets[f"photo_{i}"] = {"type": "photo", "path": a}
                elif a.startswith("http"):
                    if "wildberries.ru" in a or "wb.ru" in a:
                        assets["wb_url"] = {"type": "url", "url": a}
                    else:
                        assets[f"url_{i}"] = {"type": "url", "url": a}
                else:
                    assets[f"text_{i}"] = {"type": "text", "content": a}
        assets = assets if assets else None

        # Determine mode: card or video
        # Auto-detect: if user mentions "карточка", "инфографика", "фото" → card; "видео", "ролик" → video
        wb_mode = args.mode
        if wb_mode == "auto":
            inp_lower = args.input.lower()
            if any(w in inp_lower for w in ["карточк", "инфографик", "фото", "картинк", "изображен", "card"]):
                wb_mode = "card"
            elif any(w in inp_lower for w in ["видео", "ролик", "video", "reels", "тикток"]):
                wb_mode = "video"
            else:
                wb_mode = "card"  # default for WB product is card

        if wb_mode == "card":
            print("  Building WB product card series...", file=sys.stderr)
            plan = gen.generate_wb_product_cards(args.input, assets, count=args.cards)
            product_name = plan.get("product_name", "product")
            cards = plan.get("cards", [])
            print(f"  Product: {product_name}", file=sys.stderr)
            print(f"  Cards: {len(cards)}", file=sys.stderr)
            print(f"  Color scheme: {plan.get('color_scheme', '')}", file=sys.stderr)
            print(json.dumps(plan, ensure_ascii=False, indent=2))

            # Generate images for each card
            print(f"\n  Generating {len(cards)} card images...", file=sys.stderr)
            for card in cards:
                cid = card["id"]
                prompt = card["image_prompt"]
                headline = card.get("headline_ru", "")
                print(f"  Card {cid}: {card['card_type']} — {headline[:40]}", file=sys.stderr)
                try:
                    image_bytes = generate_image(prompt, "1024x1024")  # square, closest supported
                    label = f"wb_card_{cid}_{card['card_type']}_{product_name[:20]}"
                    save_and_open(image_bytes, label)
                except Exception as e:
                    print(f"  Card {cid} generation failed: {e}", file=sys.stderr)

        else:
            print("  Building WB product video scenario...", file=sys.stderr)
            scenario = gen.generate_wb_product_video(
                args.input,
                platform=args.platform if args.platform != "auto" else "wb",
                lang="ru",
                duration=None,
                assets=assets
            )
            print(f"  Title: {scenario['title']}", file=sys.stderr)
            print(f"  Scenes: {len(scenario['scenes'])}", file=sys.stderr)
            print(f"  Duration: {scenario.get('duration_sec')}s", file=sys.stderr)
            print(json.dumps(scenario, ensure_ascii=False, indent=2))

    # --- VIDEO → только сценарий, без генерации картинок ---
    elif request_type == "video":
        print("  Building video scenario...", file=sys.stderr)
        scripts_dir = os.path.join(os.path.dirname(__file__), "scripts")
        sys.path.insert(0, scripts_dir)
        import generate as gen
        scenario = gen.generate_video(args.input, args.platform, audience="", lang="auto")
        print(f"  Title: {scenario['title']}", file=sys.stderr)
        print(f"  Scenes: {len(scenario['scenes'])}", file=sys.stderr)
        print(f"  Duration: {scenario.get('duration_sec')}s", file=sys.stderr)
        print(json.dumps(scenario, ensure_ascii=False, indent=2))

    # --- TREND ---
    elif request_type == "trend":
        if not args.photo:
            # No photo — trend transform is impossible, fall back to plain image generation
            print("  Trend detected but no --photo provided — generating image directly", file=sys.stderr)
            prompt = expand_image_prompt(args.input)
            print(f"\nPrompt:\n{prompt}\n", file=sys.stderr)
            image_bytes = generate_image(prompt, args.size)
            save_and_open(image_bytes, f"image_{args.input[:30]}")
        elif not os.path.exists(args.photo):
            print(f"Error: photo not found: {args.photo}", file=sys.stderr)
            sys.exit(1)
        else:
            print("  Expanding trend prompt...", file=sys.stderr)
            prompt = expand_trend(args.input)
            print(f"\nPrompt:\n{prompt}\n", file=sys.stderr)
            image_bytes = generate_image(prompt, args.size, photo_path=args.photo)
            save_and_open(image_bytes, f"trend_{args.input[:30]}")

    # --- IMAGE ---
    else:
        print("  Expanding image prompt...", file=sys.stderr)
        prompt = expand_image_prompt(args.input)
        print(f"\nPrompt:\n{prompt}\n", file=sys.stderr)
        image_bytes = generate_image(prompt, args.size, photo_path=args.photo)
        save_and_open(image_bytes, f"image_{args.input[:30]}")


if __name__ == "__main__":
    main()
