import pandas as pd
import json
import csv
import time
import os
import requests

from tqdm import tqdm


MODEL_NAME = "gpt-4.1-mini"
TEXT_COLUMN = "summary"
BATCH_SIZE = 50
PART = 1
N_PARTS = 4
OUTPUT_FILE = "community_notes_labels_part1.csv"
API_KEY = ""



print("Loading data part 1/4...")

df_note = pd.read_csv('./dataset_summary_noteId.csv')

# Load already done from ALL files
already_done = set()

# Load from original file
if os.path.exists('community_notes_labels.csv'):
    existing = pd.read_csv('community_notes_labels.csv')
    already_done.update(existing['noteId'].tolist())
    print("Found {} notes in main file".format(len(already_done)))

# Load from this part file
if os.path.exists(OUTPUT_FILE):
    existing = pd.read_csv(OUTPUT_FILE)
    already_done.update(existing['noteId'].tolist())
    print("Found {} notes in part file".format(len(existing)))

# Split dataset in 4 parts
total = len(df_note)
part_size = total // N_PARTS
start_idx = (PART - 1) * part_size
end_idx = total if PART == N_PARTS else PART * part_size

df_part = df_note.iloc[start_idx:end_idx].reset_index(drop=True)
print("Part 1: rows {}-{} ({} notes)".format(start_idx, end_idx, len(df_part)))

# Filter out already done
df_part = df_part[~df_part['noteId'].isin(already_done)].reset_index(drop=True)
print("Remaining to classify: {}".format(len(df_part)))

category = pd.read_csv('./category_counts_final.csv')
CATEGORIES = list(category['category'].unique())



SYSTEM_PROMPT = """
You are an expert qualitative researcher.

Your task is to classify Community Notes.

You MUST assign:

- one primary label (label1)
- optionally one secondary label (label2)

Allowed categories:

{}

Rules:

- label1 is mandatory
- label2 is optional
- use label2 only when the note clearly belongs to two categories
- if only one category applies, set label2 to null
- DO NOT invent categories
- labels must exactly match one of the allowed categories
- return a confidence score:
    - high
    - medium
    - low

Return ONLY valid JSON.

Format:

[
  {{
    "temp_id": 0,
    "label1": "Health misinformation",
    "label2": "Missing context",
    "confidence": "high"
  }}
]
""".format(json.dumps(CATEGORIES, indent=2))


# CREATE OUTPUT FILE IF NOT EXISTS


if not os.path.exists(OUTPUT_FILE):
    with open(OUTPUT_FILE, mode="w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["noteId", "summary", "label1", "label2", "confidence"])



def call_openai(system_prompt, user_prompt):
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer {}".format(API_KEY)
    }
    payload = {
        "model": MODEL_NAME,
        "temperature": 0,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    }
    response = requests.post(
        "https://api.openai.com/v1/chat/completions",
        headers=headers,
        json=payload,
        timeout=60
    )
    response.raise_for_status()
    return response.json()["choices"][0]["message"]["content"].strip()




def classify_batch(batch_df):
    notes = []
    temp_mapping = {}
    for i, (_, row) in enumerate(batch_df.iterrows()):
        note_id = row["noteId"]
        summary = str(row[TEXT_COLUMN])
        temp_mapping[i] = {"noteId": note_id, "summary": summary}
        notes.append({"temp_id": i, "summary": summary})

    user_prompt = """
Classify the following Community Notes.

Remember:
- use ONLY categories from the allowed list
- assign exactly one primary label (label1)
- assign a secondary label (label2) only if genuinely needed
- otherwise set label2 to null

Notes:

{}
""".format(json.dumps(notes, ensure_ascii=False, indent=2))

    content = call_openai(SYSTEM_PROMPT, user_prompt)

    if content.startswith("```"):
        content = content.split("```")[1]
        if content.startswith("json"):
            content = content[4:]
        content = content.strip()

    try:
        parsed = json.loads(content)
    except Exception:
        print("\nFAILED RESPONSE:\n")
        print(content)
        raise

    if not isinstance(parsed, list):
        raise ValueError("Expected list, got {}".format(type(parsed)))

    final_results = []
    for item in parsed:
        if not isinstance(item, dict):
            continue
        temp_id = item.get("temp_id")
        if temp_id not in temp_mapping:
            continue
        label1 = item.get("label1")
        label2 = item.get("label2")
        confidence = item.get("confidence")
        if label1 not in CATEGORIES:
            print("WARNING: Invalid label1 -> {}".format(label1))
        if label2 is not None and label2 not in CATEGORIES:
            print("WARNING: Invalid label2 -> {}".format(label2))
        final_results.append({
            "noteId": temp_mapping[temp_id]["noteId"],
            "summary": temp_mapping[temp_id]["summary"],
            "label1": label1,
            "label2": label2,
            "confidence": confidence
        })
    return final_results



start_time = time.time()
total_batches = (len(df_part) + BATCH_SIZE - 1) // BATCH_SIZE

for batch_num, start in enumerate(tqdm(range(0, len(df_part), BATCH_SIZE))):
    end = start + BATCH_SIZE
    batch_df = df_part.iloc[start:end]
    max_retries = 5

    for attempt in range(max_retries):
        try:
            results = classify_batch(batch_df)
            with open(OUTPUT_FILE, mode="a", newline="", encoding="utf-8") as f:
                writer = csv.writer(f)
                for r in results:
                    writer.writerow([
                        r["noteId"],
                        r["summary"],
                        r["label1"],
                        r["label2"],
                        r["confidence"]
                    ])
            break
        except Exception as e:
            print("\nRetry {} failed for batch {}-{}".format(attempt+1, start, end))
            print(e)
            time.sleep(2 ** attempt)

    if batch_num > 0 and batch_num % 100 == 0:
        elapsed = time.time() - start_time
        rate = batch_num / elapsed
        remaining = total_batches - batch_num
        eta_hours = (remaining / rate) / 3600
        print("\nPart 1 | Batch {}/{} | ETA: {:.1f}h".format(
            batch_num, total_batches, eta_hours))

print("\nPart 1 DONE")
print("Saved to: {}".format(OUTPUT_FILE))