#!/usr/bin/env python3
"""
aggregate_delighter_clusters_with_labels.py

Reads a clustered delighters CSV and produces an aggregated CSV with:
- cluster ID
- cluster size
- representative delighter
- all canonical items in the cluster
- auto-generated theme_label + theme_description via OpenAI API

Filters:
- Only clusters with size > 3
- Only the 20–30 largest clusters

Usage:
  export OPENAI_API_KEY=your_api_key
  python label_delighter_clusters.py input.csv --output aggregated.csv
"""

import argparse
import os
import pandas as pd
import textwrap
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def build_prompt(cluster_id, canonicals):
    canon_text = "\n".join(canonicals)
    prompt = textwrap.dedent(f"""
    I have a cluster of *user delighters* (cluster ID: {cluster_id}).

    Your task:
    1. Suggest a short theme label (3–6 words).
    2. Write a 2–3 sentence description that explains the main idea of the cluster,
       highlighting what users appreciate, value, or enjoy about the product.
    3. Return ONLY valid JSON, with no explanations, no commentary, no code fences.

    JSON format to use:
    {{
      "theme_label": "...",
      "theme_description": "..."
    }}

    Here is the cluster (each line is a canonical delighter):
    {canon_text}
    """)
    return prompt.strip()

def get_label_and_description(prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,
    )
    text = response.choices[0].message.content
    try:
        import json
        parsed = json.loads(text)
        return parsed.get("theme_label", ""), parsed.get("theme_description", "")
    except Exception:
        return "", text

def aggregate_clusters(df):
    grouped = df.groupby("cluster")

    rows = []
    for idx, (cluster_id, group) in enumerate(grouped, start=1):
        cluster_size = group["cluster_size"].iloc[0] if "cluster_size" in group.columns else len(group)
        if cluster_size <= 3:
            continue

        representative = group["representative"].iloc[0] if "representative" in group.columns else group["canonical"].iloc[0]
        canonicals = group["canonical"].tolist()

        prompt = build_prompt(cluster_id, canonicals)
        print(f"→ Processing cluster {cluster_id} (size={cluster_size}, rep='{representative[:40]}...')")

        theme_label, theme_description = get_label_and_description(prompt)

        rows.append({
            "cluster": cluster_id,
            "cluster_size": cluster_size,
            "representative": representative,
            "canonicals": " || ".join(canonicals),
            "theme_label": theme_label,
            "theme_description": theme_description
        })
    return pd.DataFrame(rows)

def main():
    parser = argparse.ArgumentParser(description="Aggregate clustered delighters into theme-level summaries")
    parser.add_argument("input_csv", help="Path to input clustered CSV")
    parser.add_argument("--output", default=None, help="Path to output aggregated CSV (default: input_delighters_aggregated.csv)")
    args = parser.parse_args()

    outp = args.output or args.input_csv.replace(".csv", "_delighters_aggregated.csv")

    df = pd.read_csv(args.input_csv)

    # compute cluster sizes (in case not in file or unreliable)
    sizes = df.groupby("cluster")["canonical"].count().reset_index(name="true_size")
    df = df.merge(sizes, on="cluster", how="left")
    df["cluster_size"] = df["cluster_size"].fillna(df["true_size"])

    # filter clusters >3 and pick top 20–30 largest
    top_clusters = (
        df.groupby("cluster")["cluster_size"].first()
        .sort_values(ascending=False)
        .loc[lambda x: x > 3]
        .head(30)
        .index
    )
    df = df[df["cluster"].isin(top_clusters)]

    agg = aggregate_clusters(df)
    agg.to_csv(outp, index=False)
    print(f"\n✅ Wrote: {outp}")

if __name__ == "__main__":
    main()