#!/usr/bin/env python3
import csv
import hashlib
import os
import argparse
import math
import json

def gen_id(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, help="Path to the original input CSV (all comments)")
    ap.add_argument("--stop-batch", type=int, required=True, help="Batch number you stopped at (e.g. 169)")
    ap.add_argument("--batch-tokens", type=int, default=3000, help="Same batch size token budget you used in extract_delighters.py")
    args = ap.parse_args()

    # Sidecar file based on input filename
    base_name = os.path.splitext(os.path.basename(args.input))[0]
    sidecar_path = f"{base_name}.processed_ids"

    # Load all comments
    comments = []
    with open(args.input, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            txt = str(row.get("comment", "")).strip()
            cid = gen_id(txt)
            comments.append({
                "comment_id": cid,
                "text": txt,
                "post_url": row.get("post_url", "").strip()
            })

    # Token counter
    try:
        import tiktoken
        _ENC = tiktoken.get_encoding("cl100k_base")
        def count_tokens(text: str) -> int:
            return len(_ENC.encode(text))
    except Exception:
        def count_tokens(text: str) -> int:
            return max(1, math.ceil(len(text) / 4))

    # Batch packing (same as main script)
    def pack_batches(comments, max_total_tokens, expected_output_tokens_per_comment=60):
        batches, current = [], []
        current_input_tokens = 0
        overhead = 1000
        for item in comments:
            s = json.dumps(item, ensure_ascii=False)
            t = count_tokens(s)
            est_output = (len(current) + 1) * expected_output_tokens_per_comment
            if t + current_input_tokens + overhead + est_output > max_total_tokens:
                if current:
                    batches.append(current)
                current = [item]
                current_input_tokens = t
            else:
                current.append(item)
                current_input_tokens += t
        if current:
            batches.append(current)
        return batches

    batches = pack_batches(comments, args.batch_tokens)
    print(f"Total {len(comments)} comments → {len(batches)} batches")

    # Select all comment_ids up to stop-batch
    processed_ids = []
    for i, batch in enumerate(batches, 1):
        if i > args.stop_batch:
            break
        processed_ids.extend([c["comment_id"] for c in batch])

    # Write sidecar
    with open(sidecar_path, "w", encoding="utf-8") as f:
        for cid in processed_ids:
            f.write(cid + "\n")

    print(f"✅ Prefilled {len(processed_ids)} IDs into {sidecar_path}")
    print(f"➡️ Resume will now start at batch {args.stop_batch+1}/{len(batches)}")

if __name__ == "__main__":
    main()