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

Extracts structured fields from comments for a selected cluster in a CSV:
- workaround
- proposed_solution
- tradeoffs

It uses:
- `canonical` as the known pain point (e.g., "device tree complexity")
- `comment` as the user input
- `cluster` to filter a specific cluster (passed via CLI)

Usage:
  python extract_cluster_resolutions.py input.csv --cluster 24 --output output.csv

Requires OpenAI API key if --use_openai is enabled:
  export OPENAI_API_KEY=your_key
"""

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

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

PROMPT_TEMPLATE = """
You are analyzing user comments related to the pain point:
"{pain_point}"

From the comment below, extract only the following three fields and return a valid JSON object:

- workaround: What the user actually did to bypass or cope with the pain point.
- proposed_solution: What the user suggests or wishes existed to solve it better.
- tradeoffs: Any downsides or compromises the user mentions about their approach.

Guidelines:
- Use short, clear sentences or bullet fragments.
- If a field is not mentioned in the comment, leave it as an empty string "".
- Do not invent information. Only extract what is actually present in the comment.
- Return your response as raw JSON only.
- Do not wrap it in Markdown or code fences like ```json.
- Return only the JSON object and nothing else.

Comment:
\"\"\"{comment}\"\"\"
"""

def extract_fields(comment, pain_point, model="gpt-4o"):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    prompt = PROMPT_TEMPLATE.format(pain_point=pain_point, comment=comment)
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt.strip()}],
            temperature=0.2,
        )
        content = response.choices[0].message.content
        try:
            data = json.loads(content)
        except json.JSONDecodeError:
            print("Invalid JSON from model:\n", content)
            raise

        return {
            "workaround": data.get("workaround", ""),
            "proposed_solution": data.get("proposed_solution", ""),
            "tradeoffs": data.get("tradeoffs", "")
        }
    except Exception as e:
        print(f"Error: {e}")
        return {"workaround": "", "proposed_solution": "", "tradeoffs": ""}

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input_csv", help="Input CSV file")
    parser.add_argument("--cluster", type=int, required=True, help="Cluster number to process")
    parser.add_argument("--output", default="cluster_extracted.csv", help="Output CSV file")
    parser.add_argument("--model", default="gpt-4o", help="OpenAI model (default: gpt-4o)")
    args = parser.parse_args()

    df = pd.read_csv(args.input_csv)

    if not {"canonical", "cluster", "comment"}.issubset(df.columns):
        raise ValueError("CSV must contain 'canonical', 'cluster', and 'comment' columns.")

    subset = df[df["cluster"] == args.cluster].copy()
    if subset.empty:
        print(f"No rows found for cluster {args.cluster}")
        return

    pain_point = subset["canonical"].iloc[0]
    print(f"Processing cluster {args.cluster} for pain point: {pain_point}")

    results = []
    for i, row in subset.iterrows():
        extracted = extract_fields(row["comment"], pain_point, model=args.model)
        results.append(extracted)

    enriched = pd.concat([subset.reset_index(drop=True), pd.DataFrame(results)], axis=1)
    enriched.to_csv(args.output, index=False)
    print(f"Done. Output written to: {args.output}")

if __name__ == "__main__":
    main()