#!/usr/bin/env python3
"""
trend.py — portrait trend prompt enhancer.

Checks if the user's request matches one of the known trends (stored in trends/ folder).
If matched — returns the curated prompt directly.
If not matched — asks LLM to generate a prompt based on the request.

Usage:
    python story-gen/scripts/trend.py "аниме"
    python story-gen/scripts/trend.py "я в средневековье" --generate
    python story-gen/scripts/trend.py "vintage photo booth" --generate --size 1024x1792
"""

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")

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

# ---------------------------------------------------------------------------
# Trend matching
# ---------------------------------------------------------------------------

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


def match_trend(user_input: str, trends: list[dict]) -> dict | None:
    """Return matched trend dict or None if no match."""
    query = user_input.lower()
    for trend in trends:
        for kw in trend["keywords"]:
            if kw in query:
                return trend
    return None


# ---------------------------------------------------------------------------
# LLM fallback for unknown trends
# ---------------------------------------------------------------------------

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

The user gives you a short casual phrase describing a visual style or trend.
Your job: expand it into a detailed, professional portrait prompt.

Always structure the output as follows:

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 description]

Location: [detailed background/setting description]

Pose & Action: [body position, gesture, eye direction]

Lighting: [lighting setup and atmosphere]

Mood: [emotional tone, keywords]

Technical: [camera lens, aperture, ISO, art style, quality tags]

Rules:
- Write entirely in English
- Be very specific and detailed in every section
- The prompt must be ready to paste directly into an image generation model
- Return ONLY the prompt text, no explanations, no markdown
"""


def expand_trend_llm(user_input: str) -> str:
    payload = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": SYSTEM_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()


def expand_trend(user_input: str) -> tuple[str, str]:
    """Returns (prompt, source) where source is trend name or 'llm'."""
    trends = load_trends()
    matched = match_trend(user_input, trends)
    if matched:
        print(f"  Matched trend: {matched['name']}", file=sys.stderr)
        return matched["prompt"], matched["name"]
    print("  No trend matched — using LLM to generate prompt", file=sys.stderr)
    return expand_trend_llm(user_input), "llm"


# ---------------------------------------------------------------------------
# Image generation
# ---------------------------------------------------------------------------

def generate_image(prompt: str, size: str) -> bytes:
    print("  Generating image...", file=sys.stderr)
    payload = json.dumps({
        "model": IMAGE_MODEL,
        "prompt": prompt,
        "n": 1,
        "size": size
    }).encode()
    req = request.Request(
        f"{API_URL}/images/generations",
        data=payload,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
    )
    with request.urlopen(req, timeout=120) 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 main():
    parser = argparse.ArgumentParser(description="Portrait trend prompt enhancer")
    parser.add_argument("input", help="Short trend phrase in any language")
    parser.add_argument("--generate", "-g", action="store_true",
                        help="Generate image after expanding the prompt")
    parser.add_argument("--size", default="1024x1024",
                        choices=["1024x1024", "1792x1024", "1024x1792"])
    args = parser.parse_args()

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

    print(f"Processing trend request: {args.input!r}", file=sys.stderr)
    prompt, _ = expand_trend(args.input)

    print(prompt)

    if args.generate:
        image_bytes = generate_image(prompt, args.size)

        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 args.input)[:40]
        path = os.path.join(out_dir, f"trend_{safe}.png")
        with open(path, "wb") as f:
            f.write(image_bytes)
        print(f"\nSaved: {path}", file=sys.stderr)
        open_image(path)


if __name__ == "__main__":
    main()
