# jtbd_openai_extraction_batches.py
import pandas as pd
import json
import os
import argparse

from llm_client import call_llm_chat_json

def main(input_file, output_file, batch_size=10):
    # Load all comments, drop NaN and duplicates
    df = pd.read_csv(input_file)
    comments = df["comment"].dropna().drop_duplicates().tolist()

    # Resume option if output file exists
    if os.path.exists(output_file):
        existing = pd.read_csv(output_file)
        processed = set(existing["comment"].tolist())
        results = existing.to_dict("records")
        print(f"▶ Resuming, {len(processed)} comments already processed")
    else:
        processed = set()
        results = []

    # prompt_template = """
    # You are a Jobs-to-be-Done (JTBD) extractor.

    # For the given comment:
    # - If it contains a clear JTBD signal, return this JSON:
    #   {{
    #     "jtbd": "YES",
    #     "situation": "...",
    #     "struggle": "...",
    #     "outcome": "..."
    #   }}

    # - If it does NOT contain a JTBD (no clear situation/struggle/outcome), return this JSON:
    #   {{ "jtbd": "NO" }}

    # Rules:
    # - Always output valid JSON.
    # - Do not include any text outside the JSON.
    # - If unsure, prefer "jtbd": "NO".

    # Comment: "{comment}"
    # """
    system_message = """You are a Jobs-to-be-Done (JTBD) extractor.

For the given comment, identify whether it expresses a clear JTBD — a "job" someone is trying to get done in a specific situation, along with what they want to achieve and why.

If the comment contains a clear JTBD, return **only** this JSON:
{
  "jtbd": "YES",
  "job_statement": "When I [situation], I want to [motivation/struggle], so I can [desired outcome].",
  "situation": "...",
  "struggle": "...",
  "outcome": "..."
}

If the comment does NOT contain a clear JTBD signal (no situation, motivation, or outcome), return this JSON:
{ "jtbd": "NO" }

Guidelines:
- Always output **valid JSON only**.
- Use the user's own language when possible, but paraphrase slightly for clarity.
- Be concise: each field should be 1–2 short sentences.
- If unsure whether the comment includes a complete job, prefer { "jtbd": "NO" }.
"""

    seen_in_run = set(processed)  # track duplicates inside this run as well

    for start in range(0, len(comments), batch_size):
        end = start + batch_size
        batch = comments[start:end]

        for c in batch:
            if c in seen_in_run:
                continue  # skip duplicates (already processed in past or in this run)

            result = call_llm_chat_json(
                {"temperature": 0.0, "max_tokens": 512},
                system_message=system_message,
                user_message=f'Comment: "{c}"',
            )
            if result is None:
                result = {"jtbd": "NO"}

            row = {"comment": c, **result}
            results.append(row)

            seen_in_run.add(c)  # mark as processed

        # Save after each batch
        out_df = pd.DataFrame(results)
        out_df.to_csv(output_file, index=False)

        print(f"✅ Processed {start+1}–{end} / {len(comments)} comments. Saved to {output_file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="JTBD extractor for comments")
    parser.add_argument("input_file", help="Path to input CSV containing comments")
    parser.add_argument("output_file", help="Path to output CSV for results")
    parser.add_argument("--batch-size", type=int, default=10, help="Number of comments per batch")
    args = parser.parse_args()

    main(args.input_file, args.output_file, batch_size=args.batch_size)