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

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

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

Usage:
  export OPENAI_API_KEY=your_api_key
  python aggregate_clusters_with_labels.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 pain points (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 are complaining about.
    3. Return ONLY valid JSON, with no explanations, no commentary, no code fences.
       Do not include triple backticks, "json" markers, or any text outside the JSON object.

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

    Here is the cluster:
    {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  # fallback: put whole text in description

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 pain points 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_aggregated.csv)")
    args = parser.parse_args()

    outp = args.output or args.input_csv.replace(".csv", "_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 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()