#!/usr/bin/env python3
import os
import sys
import pandas as pd
from openai import OpenAI

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

def label_cluster(words, model="gpt-4o-mini"):
    prompt = f"""
You are given a cluster of words extracted from Reddit comments about doomscrolling.
Summarize the theme of this cluster as a short label (max 3 words).
Do not just repeat the words, infer the category.

Words: {", ".join(words[:20])}
Answer with only the label.
"""
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful assistant that labels word clusters."},
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )
    return resp.choices[0].message.content.strip()

def main(input_file, top_n):
    df = pd.read_csv(input_file)

    # Count cluster sizes
    cluster_sizes = df.groupby("cluster_id").size().reset_index(name="count")
    top_clusters = cluster_sizes.sort_values("count", ascending=False).head(top_n)

    results = []
    for cid in top_clusters["cluster_id"]:
        words = df[df["cluster_id"] == cid]["word"].tolist()
        label = label_cluster(words)
        results.append({
            "cluster_id": cid,
            "size": len(words),
            "label": label,
            "top_words": ", ".join(words[:15])
        })
        print(f"Cluster {cid} → {label}")

    out_file = f"{input_file.split('.csv')[0]}_labels.csv"
    pd.DataFrame(results).to_csv(out_file, index=False)
    print(f"\nSaved labels → {out_file}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python cluster_labeler.py <clusters.csv> <top_n>")
        sys.exit(1)

    input_file = sys.argv[1]
    top_n = int(sys.argv[2])
    main(input_file, top_n)