04_survey.py

"""
04_survey.py
------------
Generate a survey draft by sending topics/<topic>/all_cards.md through the
survey prompts as a multi-turn conversation with Gemini via Vertex AI.

Each prompt's response is saved individually so the run is resume-safe.
When resuming, prior turns are replayed as chat history so the model
retains the full conversation context.

Usage:
    python 04_survey.py --topic <name>
    python 04_survey.py --topic <name> --topic-name "AI-based Vulnerability Detection"
"""

import argparse
import os
import re
import sys
import time
from pathlib import Path

from dotenv import load_dotenv
import vertexai
from vertexai.generative_models import Content, GenerativeModel, Part

load_dotenv()

# ── Args ──────────────────────────────────────────────────────────────────────

parser = argparse.ArgumentParser()
parser.add_argument("--topic", required=True,
                    help="Topic name under topics/ (e.g. vul)")
parser.add_argument("--topic-name", default=None,
                    help="Human-readable topic name, e.g. 'AI-based Vulnerability Detection'")
args = parser.parse_args()

# ── Config ────────────────────────────────────────────────────────────────────

PROJECT_ID = os.getenv("GCP_PROJECT_ID")
LOCATION   = os.getenv("GCP_LOCATION", "us-central1")
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")

BASE         = Path(f"./topics/{args.topic}")
CARDS_FILE   = BASE / "all_cards.md"
PROMPTS_FILE = BASE / "prompts" / "survey.md"
PARTS_DIR    = BASE / "survey_parts"
DRAFT_FILE   = BASE / "survey_draft.md"

RATE_LIMIT_DELAY = 5  # seconds between requests

# ── Validation ────────────────────────────────────────────────────────────────

if not PROJECT_ID:
    sys.exit("Error: GCP_PROJECT_ID not found. Check your .env file.")
if not BASE.exists():
    sys.exit(f"Error: topic directory not found: {BASE}")
if not CARDS_FILE.exists():
    sys.exit(f"Error: {CARDS_FILE} not found. Run 03_merge.py first.")
if not PROMPTS_FILE.exists():
    sys.exit(f"Error: {PROMPTS_FILE} not found.")

vertexai.init(project=PROJECT_ID, location=LOCATION)
model = GenerativeModel(MODEL_NAME)

PARTS_DIR.mkdir(exist_ok=True)

# ── Load ──────────────────────────────────────────────────────────────────────

cards_text   = CARDS_FILE.read_text(encoding="utf-8")
prompts_text = PROMPTS_FILE.read_text(encoding="utf-8")

m = re.search(r"^Total papers:\s*(\d+)", cards_text, re.MULTILINE)
n_papers   = int(m.group(1)) if m else 0
topic_name = args.topic_name or args.topic

# ── Parse survey.md ───────────────────────────────────────────────────────────

def parse_prompts(text: str) -> list[tuple[str, str]]:
    """Return (title, body) for each PROMPT N block in survey.md."""
    pattern = r"={40,}\n(PROMPT \d+ — .+?)\n={40,}\n(.*?)(?=\n={40,}|\Z)"
    return [(t.strip(), b.strip())
            for t, b in re.findall(pattern, text, re.DOTALL)]


def fill_placeholders(text: str) -> str:
    return text.replace("[TOPIC]", topic_name).replace("[N]", str(n_papers))


prompts = parse_prompts(prompts_text)
if not prompts:
    sys.exit(f"Error: no PROMPT blocks found in {PROMPTS_FILE}")

# Warn about unfilled template placeholders
remaining = re.findall(r"\[[A-Z][A-Z _]+\]",
                        prompts_text.replace("[TOPIC]", "").replace("[N]", ""))
if remaining:
    print(f"Warning: unfilled placeholders in survey prompts: {', '.join(sorted(set(remaining)))}")
    print(f"  Edit {PROMPTS_FILE} to fill them in for best results.\n")

# ── Build messages ────────────────────────────────────────────────────────────

def build_prompt0(body: str) -> str:
    """Adapt PROMPT 0 for API use: strip upload instruction, inline the cards."""
    # Remove the "(Upload all_cards.md, then send this message)" line + blank line
    body = re.sub(r"^\(Upload.*?\n\n", "", body, flags=re.DOTALL)
    body = fill_placeholders(body)
    body = body.replace(
        "I have attached a file\ncontaining",
        "Below is the complete data for",
    ).replace(
        "I have attached a file containing",
        "Below is the complete data for",
    )
    return (
        body
        + "\n\n--- BEGIN LITERATURE CARDS ---\n\n"
        + cards_text
        + "\n\n--- END LITERATURE CARDS ---"
    )


def user_message(i: int, body: str) -> str:
    return build_prompt0(body) if i == 0 else fill_placeholders(body)

# ── File helpers ──────────────────────────────────────────────────────────────

def part_path(i: int, title: str) -> Path:
    slug = re.sub(r"[^a-z0-9]+", "_", title.lower()).strip("_")
    return PARTS_DIR / f"{i:02d}_{slug}.md"


def write_draft(responses: list[tuple[str, str]]) -> None:
    parts = [f"# {title}\n\n{text}" for title, text in responses]
    DRAFT_FILE.write_text("\n\n---\n\n".join(parts), encoding="utf-8")
    print(f"\nDraft saved → {DRAFT_FILE}")

# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    print(f"Topic      : {args.topic} ({topic_name})")
    print(f"Project    : {PROJECT_ID}")
    print(f"Model      : {MODEL_NAME}")
    print(f"Papers     : {n_papers}")
    print(f"Prompts    : {len(prompts)}")
    print(f"Parts dir  : {PARTS_DIR}\n")

    # Scan which prompts are already done; collect (title, user_msg, response)
    # for replaying into chat history on resume.
    done: list[tuple[str, str, str]] = []
    start_idx = 0
    for i, (title, body) in enumerate(prompts):
        out = part_path(i, title)
        if out.exists():
            msg = user_message(i, body)
            resp = out.read_text(encoding="utf-8")
            done.append((title, msg, resp))
            start_idx = i + 1
            print(f"[{i+1}/{len(prompts)}] SKIP  {title}  (already done)")
        else:
            break

    if start_idx == len(prompts):
        print("\nAll prompts done. Rebuilding draft...")
        write_draft([(t, r) for t, _, r in done])
        return

    # Replay completed turns as history so the model has full context on resume
    history = []
    for _, msg, resp in done:
        history.append(Content(role="user",  parts=[Part.from_text(msg)]))
        history.append(Content(role="model", parts=[Part.from_text(resp)]))

    chat = model.start_chat(history=history)
    all_responses = [(t, r) for t, _, r in done]

    for i in range(start_idx, len(prompts)):
        title, body = prompts[i]
        out  = part_path(i, title)
        msg  = user_message(i, body)

        print(f"[{i+1}/{len(prompts)}] {title} ...", end="", flush=True)
        try:
            response = chat.send_message(msg)
            text = response.text
            out.write_text(text, encoding="utf-8")
            all_responses.append((title, text))
            print("  ✓")
        except Exception as e:
            print(f"  ✗  ERROR: {e}")
            out.with_suffix(".error.txt").write_text(str(e), encoding="utf-8")
            all_responses.append((title, f"[ERROR: {e}]"))

        if i < len(prompts) - 1:
            time.sleep(RATE_LIMIT_DELAY)

    write_draft(all_responses)


if __name__ == "__main__":
    main()