#!/usr/bin/env python3
"""
Generate AI-powered PDP model descriptions using Claude Haiku 4.5.
Runs once per brand+model combo; resumable. Writes to model_content.pdp_body
and variant pools (opener_variants, closer_variants).

Usage:
  python scripts/generate-model-descriptions.py              # generate missing
  python scripts/generate-model-descriptions.py --limit 5    # test run
  python scripts/generate-model-descriptions.py --force      # regen even if exists
  python scripts/generate-model-descriptions.py --brand nike # filter per brand slug
"""
import os
import sys
import json
import time
import argparse
from pathlib import Path

import psycopg2
import psycopg2.extras
from dotenv import load_dotenv
from anthropic import Anthropic, APIError

load_dotenv()

MODEL = "claude-haiku-4-5"
MAX_TOKENS = 1200
DB_URL = os.environ["DATABASE_URL"]

BRAND_HERITAGE_PATH = Path(__file__).parent / "brand-heritage.json"

SLUG_EXPR = (
    "LOWER(REGEXP_REPLACE(REGEXP_REPLACE(TRIM(p.model), '[^a-zA-Z0-9]+', '-', 'g'), "
    "'(^-|-$)', '', 'g'))"
)


def load_brand_heritage():
    with open(BRAND_HERITAGE_PATH, "r", encoding="utf-8") as f:
        return json.load(f)


def brand_slug(brand_name: str) -> str:
    import re
    s = re.sub(r"[^a-zA-Z0-9]+", "-", brand_name.strip()).strip("-").lower()
    return s


def fetch_models(conn, force: bool = False, brand: str | None = None, limit: int | None = None):
    """Return list of dicts with brand_slug, brand_name, model_name, model_slug, variant_count."""
    brand_filter = ""
    params: list = []
    if brand:
        brand_filter = "AND LOWER(REGEXP_REPLACE(REGEXP_REPLACE(p.brand, '[^a-zA-Z0-9]+', '-', 'g'), '(^-|-$)', '', 'g')) = %s"
        params.append(brand.lower())

    sql = f"""
        SELECT
          LOWER(REGEXP_REPLACE(REGEXP_REPLACE(p.brand, '[^a-zA-Z0-9]+', '-', 'g'), '(^-|-$)', '', 'g')) AS brand_slug,
          p.brand AS brand_name,
          p.model AS model_name,
          {SLUG_EXPR} AS model_slug,
          COUNT(DISTINCT p.id) AS variant_count
        FROM products p
        WHERE p.is_sneaker = true
          AND p.model IS NOT NULL AND TRIM(p.model) <> ''
          AND p.model ~ '[A-Za-z]{{2,}}'
          {brand_filter}
        GROUP BY p.brand, p.model
        HAVING COUNT(DISTINCT p.id) >= 2
        ORDER BY variant_count DESC
    """
    with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
        cur.execute(sql, params)
        all_rows = cur.fetchall()

    if not force:
        # Filter out rows that already have pdp_body
        with conn.cursor() as cur:
            cur.execute(
                "SELECT brand_slug, model_slug FROM model_content WHERE pdp_body IS NOT NULL"
            )
            done = {(r[0], r[1]) for r in cur.fetchall()}
        all_rows = [r for r in all_rows if (r["brand_slug"], r["model_slug"]) not in done]

    if limit:
        all_rows = all_rows[:limit]
    return all_rows


def build_prompt(brand_name: str, brand_heritage: str, model_name: str) -> str:
    heritage_block = brand_heritage or f"{brand_name} is een bekend sneakermerk."
    return f"""Je bent copywriter voor SneakerPicks, een Nederlandse sneaker-prijsvergelijker.
Schrijf in natuurlijk Nederlands, informeel maar professioneel. Geen cliches, geen uitroeptekens.

Merk: {brand_name}
Merk-achtergrond: {heritage_block}
Model: {model_name}

Genereer EXACT dit JSON-object, geen extra tekst:
{{
  "pdp_body": "<een paragraaf van 80 tot 120 woorden over dit specifieke model. Noem silhouet, geschiedenis indien relevant, waarom mensen het kopen, technologie of bijzondere materialen. NIET over prijs, kleur of shops schrijven.>",
  "opener_variants": [
    "<openingszin, 15-25 woorden. Variant 1>",
    "<variant 2 met duidelijk andere structuur en woordkeuze>",
    "<variant 3 met duidelijk andere structuur en woordkeuze>"
  ],
  "closer_variants": [
    "<afsluitende CTA richting prijsvergelijken, 15-25 woorden. Variant 1>",
    "<variant 2>",
    "<variant 3>"
  ],
  "style_tags": ["<max 3 tags zoals retro, skate, running, heritage, technical>"]
}}

Belangrijk:
- Noem NOOIT concrete prijzen, maten, kleuren of specifieke shops.
- Alleen Nederlands. Geen Engelse zinnen.
- Zorg dat elke variant duidelijk anders voelt.
- Output uitsluitend geldig JSON, niets eromheen.
"""


def extract_json(text: str) -> dict | None:
    """Pull the first {...} block out of the response text."""
    text = text.strip()
    if text.startswith("```"):
        # strip code fences
        text = text.strip("`")
        if text.startswith("json"):
            text = text[4:]
    start = text.find("{")
    end = text.rfind("}")
    if start < 0 or end <= start:
        return None
    try:
        return json.loads(text[start : end + 1])
    except json.JSONDecodeError:
        return None


def call_ai(client: Anthropic, prompt: str, retries: int = 2) -> dict | None:
    for attempt in range(retries + 1):
        try:
            response = client.messages.create(
                model=MODEL,
                max_tokens=MAX_TOKENS,
                messages=[{"role": "user", "content": prompt}],
            )
            text = "".join(
                block.text for block in response.content if getattr(block, "type", "") == "text"
            )
            parsed = extract_json(text)
            if parsed:
                usage = getattr(response, "usage", None)
                if usage:
                    parsed["_usage"] = {
                        "input_tokens": usage.input_tokens,
                        "output_tokens": usage.output_tokens,
                    }
                return parsed
            print(f"    JSON parse failed on attempt {attempt + 1}, retrying...", file=sys.stderr)
        except APIError as e:
            print(f"    API error: {e}. Retrying in {2 ** attempt}s...", file=sys.stderr)
            time.sleep(2 ** attempt)
    return None


def validate(payload: dict) -> bool:
    if not isinstance(payload.get("pdp_body"), str) or len(payload["pdp_body"]) < 60:
        return False
    for key in ("opener_variants", "closer_variants"):
        arr = payload.get(key)
        if not isinstance(arr, list) or len(arr) < 2:
            return False
        if not all(isinstance(s, str) and len(s) > 5 for s in arr):
            return False
    tags = payload.get("style_tags")
    if tags is not None and not (isinstance(tags, list) and all(isinstance(t, str) for t in tags)):
        return False
    return True


def upsert(conn, brand_slug: str, model_slug: str, model_name: str, payload: dict):
    sql = """
        INSERT INTO model_content (brand_slug, model_slug, model_name, pdp_body,
                                   opener_variants, closer_variants, style_tags,
                                   ai_model, ai_generated_at, updated_at)
        VALUES (%s, %s, %s, %s, %s::jsonb, %s::jsonb, %s::jsonb, %s, NOW(), NOW())
        ON CONFLICT (brand_slug, model_slug) DO UPDATE SET
          pdp_body = EXCLUDED.pdp_body,
          opener_variants = EXCLUDED.opener_variants,
          closer_variants = EXCLUDED.closer_variants,
          style_tags = EXCLUDED.style_tags,
          ai_model = EXCLUDED.ai_model,
          ai_generated_at = NOW(),
          updated_at = NOW()
    """
    with conn.cursor() as cur:
        cur.execute(
            sql,
            (
                brand_slug,
                model_slug,
                model_name,
                payload["pdp_body"].strip(),
                json.dumps(payload["opener_variants"]),
                json.dumps(payload["closer_variants"]),
                json.dumps(payload.get("style_tags") or []),
                MODEL,
            ),
        )
    conn.commit()


def cost_estimate(usage: dict | None) -> float:
    if not usage:
        return 0.0
    # Claude Haiku 4.5 pricing (confirm at runtime if pricing changes)
    in_cost = usage.get("input_tokens", 0) / 1_000_000 * 1.0
    out_cost = usage.get("output_tokens", 0) / 1_000_000 * 5.0
    return in_cost + out_cost


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=None, help="Max number of models to process")
    parser.add_argument("--force", action="store_true", help="Regenerate even if pdp_body exists")
    parser.add_argument("--brand", type=str, default=None, help="Filter by brand slug (e.g. nike)")
    parser.add_argument("--dry-run", action="store_true", help="Print but don't call the API")
    args = parser.parse_args()

    brand_heritage = load_brand_heritage()
    client = Anthropic()  # reads ANTHROPIC_API_KEY

    conn = psycopg2.connect(DB_URL)
    try:
        models = fetch_models(conn, force=args.force, brand=args.brand, limit=args.limit)
        print(f"Found {len(models)} models to process")

        total_cost = 0.0
        ok = 0
        fail = 0

        for idx, m in enumerate(models, 1):
            bslug = m["brand_slug"]
            mslug = m["model_slug"]
            bname = m["brand_name"]
            mname = m["model_name"]
            heritage = brand_heritage.get(bslug, "")

            print(f"[{idx}/{len(models)}] {bname} / {mname} ({m['variant_count']} varianten)")

            if args.dry_run:
                prompt = build_prompt(bname, heritage, mname)
                print(prompt[:300] + "...")
                continue

            prompt = build_prompt(bname, heritage, mname)
            payload = call_ai(client, prompt)

            if not payload or not validate(payload):
                print(f"    FAIL: invalid or missing payload")
                fail += 1
                continue

            usage = payload.pop("_usage", None)
            cost = cost_estimate(usage)
            total_cost += cost

            try:
                upsert(conn, bslug, mslug, mname, payload)
                ok += 1
                print(f"    OK  (est. ${cost:.4f}, running total ${total_cost:.4f})")
            except Exception as e:
                print(f"    DB error: {e}")
                fail += 1

            time.sleep(0.3)

        print(f"\nDone. ok={ok} fail={fail} total_est_cost=${total_cost:.2f}")
    finally:
        conn.close()


if __name__ == "__main__":
    main()