src/generate_traces.py

"""
Synthetic search agent trace generator.

Generates thousands of diverse, realistic search agent traces using
actual code chunks from the corpus. Each trace follows the gold-standard
format: reasoning → search → results → analysis → evidence → finish.

Trace types generated:
  1. Simple lookup (1 search, 1 result)
  2. Lookup with noise (1 search, 2-3 results, some irrelevant)
  3. Multi-step search (2 searches, progressive refinement)
  4. Query decomposition (complex query → subqueries)
  5. Not found (search returns nothing)
  6. Type/struct inspection (field enumeration)
  7. Function behavior analysis (what does it do?)
  8. Usage pattern (how is X used?)

Query templates are varied to prevent memorization:
  - "What does {name} do?"
  - "Where is {name} defined?"
  - "How does {name} work?"
  - "What parameters does {name} accept?"
  - "What fields does {name} have?"
  - "How is {name} used in the codebase?"
  - etc.

Output: data/sft_traces.jsonl
"""

import json
import os
import random
import re

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")
TRACES_PATH = os.path.join(PROJECT_DIR, "data", "sft_traces.jsonl")

SYSTEM_PROMPT = (
    "You are a code search agent. Given a query from a reasoning model, "
    "decompose it into subqueries, search the codebase, inspect results, "
    "and return curated evidence. Use <|search|> to issue searches, "
    "<|reasoning|> to analyze, and <|evidence|> to return findings. "
    "Be concise. Extract only the relevant facts. End with <|finish|>."
)

random.seed(42)

# ─── Query templates ─────────────────────────────────────────────────────────
QUERY_TEMPLATES = {
    "function": [
        "What does {name} do?",
        "Where is {name} defined?",
        "How does {name} work?",
        "What parameters does {name} accept?",
        "What does {name} return?",
        "Find the implementation of {name}.",
        "Explain what {name} does step by step.",
        "What is the signature of {name}?",
    ],
    "struct": [
        "What fields does {name} have?",
        "What is the structure of {name}?",
        "Where is {name} defined?",
        "What data does {name} contain?",
        "Describe the {name} type.",
    ],
    "class": [
        "What does the {name} class do?",
        "What methods does {name} have?",
        "Where is the {name} class defined?",
        "How is {name} structured?",
        "What is the inheritance of {name}?",
    ],
    "macro": [
        "What does {name} do?",
        "Where is {name} defined?",
        "What is {name}?",
    ],
    "enum": [
        "What variants does {name} have?",
        "Where is {name} defined?",
        "What values can {name} take?",
    ],
    "block": [
        "What does {name} do?",
        "Where is {name} defined?",
        "Find code related to {name}.",
        "How is {name} used?",
    ],
}

# ─── Reasoning templates ─────────────────────────────────────────────────────
INITIAL_REASONING = [
    "I need to find {desc}. Let me search for {search_term}.",
    "Looking for {desc}. I'll search using the term '{search_term}'.",
    "The query asks about {desc}. Let me search the codebase.",
    "I should find {desc}. Searching for '{search_term}'.",
    "This requires finding {desc}. Let me issue a search.",
]

ANALYSIS_REASONING_FOUND = [
    "Found the relevant code. {analysis}",
    "This result contains what I need. {analysis}",
    "I found {desc}. {analysis}",
    "This is the right code. {analysis}",
]

ANALYSIS_REASONING_NOISE = [
    "The first result is relevant. {analysis} The other results are not directly related.",
    "Result 1 matches the query. {analysis} The remaining results appear to be unrelated code.",
    "I found the target in the first result. {analysis} The other results don't match.",
    "The relevant code is in the first result. {analysis} Ignoring the noise results.",
]

ANALYSIS_REASONING_REFINE = [
    "Found the definition, but I need to see how it's used. Let me search for usages.",
    "I have the implementation. Now let me find where it's called.",
    "Got the definition. Let me also search for related patterns.",
    "Found it. Let me do one more search to get complete context.",
]

NOT_FOUND_REASONING = [
    "No results found for '{search_term}'. Let me try a broader search.",
    "The search returned nothing. Let me try different terms.",
    "No matches. The codebase may not contain this. Let me verify with another search.",
]

# ─── Code analysis functions ─────────────────────────────────────────────────

def extract_function_info(code: str, name: str) -> dict:
    """Extract key information from a function."""
    info = {"name": name, "params": [], "returns": "", "key_lines": []}

    # Extract parameters from signature
    sig_match = re.search(r"(?:def|fn|func|function)\s+" + re.escape(name) + r"\s*\(([^)]*)\)", code)
    if sig_match:
        params_raw = sig_match.group(1).strip()
        if params_raw:
            info["params"] = [p.strip() for p in params_raw.split(",") if p.strip()]

    # Extract return type (Rust/TS)
    ret_match = re.search(r"->\s*([^{]+)", code)
    if ret_match:
        info["returns"] = ret_match.group(1).strip()

    # Extract key lines (non-trivial lines)
    lines = code.split("\n")
    key_lines = []
    for line in lines:
        stripped = line.strip()
        if not stripped or stripped.startswith("//") or stripped.startswith("#"):
            continue
        if any(kw in stripped for kw in ["return ", "if ", "for ", "while ", "raise ", "throw ", "Err(", "Ok(", "assert"]):
            key_lines.append(stripped)
    info["key_lines"] = key_lines[:5]

    return info


def extract_struct_info(code: str, name: str) -> dict:
    """Extract field information from a struct."""
    info = {"name": name, "fields": []}

    # Find field declarations (type name; or name: type)
    field_patterns = [
        re.compile(r"^\s+(\w+)\s+(\w+);", re.M),           # C: type name;
        re.compile(r"^\s+(\w+):\s+([^,;]+)", re.M),         # Rust/TS: name: type
        re.compile(r"^\s+(self\.)?(\w+)\s*[:=]", re.M),      # Python: self.name =
    ]

    for pat in field_patterns:
        matches = pat.findall(code)
        for m in matches:
            if isinstance(m, tuple):
                field_name = m[-1] if m[-1] else m[0]
            else:
                field_name = m
            if field_name and field_name not in ("self", "pub", "fn", "def", "class"):
                info["fields"].append(field_name)

    return info


def generate_evidence(chunk: dict, analysis_type: str = "found") -> str:
    """Generate curated evidence from a code chunk."""
    code = chunk["code"]
    name = chunk["name"]
    typ = chunk["type"]
    lang = chunk["language"]
    filepath = chunk["filepath"]

    if typ in ("function",):
        info = extract_function_info(code, name)
        parts = [f"`{name}` is a {lang} function:"]

        # Signature
        if info["params"]:
            params_str = ", ".join(info["params"][:6])
            if len(info["params"]) > 6:
                params_str += ", ..."
            parts.append(f"- Parameters: {params_str}")

        if info["returns"]:
            parts.append(f"- Returns: {info['returns']}")

        # Key behavior
        if info["key_lines"]:
            parts.append("- Key behavior:")
            for line in info["key_lines"][:3]:
                parts.append(f"  - `{line}`")

        parts.append(f"Source: {filepath}")
        return "\n".join(parts)

    elif typ in ("struct",):
        info = extract_struct_info(code, name)
        parts = [f"`{name}` is a {lang} structure with fields:"]
        for field in info["fields"][:10]:
            parts.append(f"- `{field}`")
        if len(info["fields"]) > 10:
            parts.append(f"- ... ({len(info['fields'])} total fields)")
        parts.append(f"Source: {filepath}")
        return "\n".join(parts)

    elif typ in ("class",):
        parts = [f"`{name}` is a {lang} class:"]
        # Find methods
        methods = re.findall(r"(?:def |fn |func |function )\s*(\w+)", code)
        if methods:
            parts.append(f"- Methods: {', '.join(methods[:8])}")
        # Find inheritance
        base_match = re.search(r"class\s+\w+\s*(?:\(|:\s*|extends\s+|<\s*)(\w+)", code)
        if base_match:
            parts.append(f"- Inherits from: {base_match.group(1)}")
        parts.append(f"Source: {filepath}")
        return "\n".join(parts)

    elif typ in ("macro",):
        parts = [f"`{name}` is a {lang} macro/preprocessor definition:"]
        # First few lines of the macro
        lines = code.strip().split("\n")[:5]
        for line in lines:
            parts.append(f"  `{line}`")
        parts.append(f"Source: {filepath}")
        return "\n".join(parts)

    elif typ in ("enum",):
        parts = [f"`{name}` is a {lang} enum with variants:"]
        variants = re.findall(r"^\s+(\w+)[,\s]*$", code, re.M)
        for v in variants[:10]:
            parts.append(f"- `{v}`")
        parts.append(f"Source: {filepath}")
        return "\n".join(parts)

    else:
        # Generic block
        parts = [f"Found code related to `{name}` ({lang}):"]
        lines = code.strip().split("\n")[:8]
        for line in lines:
            parts.append(f"  `{line}`")
        parts.append(f"Source: {filepath}")
        return "\n".join(parts)


def generate_search_term(chunk: dict) -> str:
    """Generate a realistic search term for a chunk."""
    name = chunk["name"]
    if name != "unknown" and name != "block":
        return name

    # For unnamed blocks, extract a key identifier
    code = chunk["code"]
    identifiers = re.findall(r"\b(ngx_\w+|def\s+\w+|fn\s+\w+|struct\s+\w+|class\s+\w+)\b", code)
    if identifiers:
        return identifiers[0].replace("def ", "").replace("fn ", "").replace("struct ", "").replace("class ", "")

    # Fallback: first meaningful word
    words = re.findall(r"\b[a-z_]{4,}\b", code)
    return words[0] if words else "code"


def generate_description(chunk: dict) -> str:
    """Generate a natural language description of what we're looking for."""
    name = chunk["name"]
    typ = chunk["type"]
    lang = chunk["language"]

    if typ == "function":
        return f"the {name} function"
    elif typ == "struct":
        return f"the {name} structure"
    elif typ == "class":
        return f"the {name} class"
    elif typ == "macro":
        return f"the {name} macro"
    elif typ == "enum":
        return f"the {name} enum"
    else:
        return f"code related to {name}"


# ─── Trace generators ────────────────────────────────────────────────────────

def make_trace_simple(chunk: dict) -> dict:
    """Type 1: Simple lookup — 1 search, 1 result."""
    name = chunk["name"]
    typ = chunk["type"]
    templates = QUERY_TEMPLATES.get(typ, QUERY_TEMPLATES["block"])
    query = random.choice(templates).format(name=name)

    search_term = generate_search_term(chunk)
    desc = generate_description(chunk)

    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term)
    evidence = generate_evidence(chunk)

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found the relevant code. This is {desc}.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
        ],
    }


def make_trace_with_noise(chunk: dict, noise_chunks: list) -> dict:
    """Type 2: Lookup with noise — 1 search, 2-3 results (some irrelevant)."""
    name = chunk["name"]
    typ = chunk["type"]
    templates = QUERY_TEMPLATES.get(typ, QUERY_TEMPLATES["block"])
    query = random.choice(templates).format(name=name)

    search_term = generate_search_term(chunk)
    desc = generate_description(chunk)
    evidence = generate_evidence(chunk)

    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term)

    # Select 1-2 noise chunks
    n_noise = random.choice([1, 2])
    noise = random.sample(noise_chunks, min(n_noise, len(noise_chunks)))

    analysis = f"This contains {desc}."
    analysis_template = random.choice(ANALYSIS_REASONING_NOISE).format(analysis=analysis)

    trace = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
        {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term}<|end|>"},
        {"role": "result", "content": chunk["code"]},
    ]
    for nc in noise:
        trace.append({"role": "result", "content": nc["code"]})

    trace.append({"role": "assistant", "content": f"<|reasoning|>{analysis_template}<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"})

    return {"query": query, "trace": trace}


def make_trace_multi_step(chunk: dict, related_chunk: dict) -> dict:
    """Type 3: Multi-step search — 2 searches, progressive refinement."""
    name = chunk["name"]
    typ = chunk["type"]
    query = f"What does {name} do and where is it used?"

    search_term1 = generate_search_term(chunk)
    search_term2 = generate_search_term(related_chunk)
    desc = generate_description(chunk)
    evidence1 = generate_evidence(chunk)
    evidence2 = generate_evidence(related_chunk)

    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term1)
    reasoning2 = random.choice(ANALYSIS_REASONING_REFINE)

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term1}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found {desc}. {reasoning2}<|end|>\n<|search|>{search_term2}<|end|>"},
            {"role": "result", "content": related_chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found both the definition and a usage. I now have complete context.<|end|>\n<|evidence|>{evidence1}\n\nRelated usage:\n{evidence2}<|end|>\n<|finish|>"},
        ],
    }


def make_trace_not_found() -> dict:
    """Type 5: Not found — search returns nothing."""
    fake_names = ["database_pool", "auth_middleware", "cache_manager", "session_store",
                  "websocket_handler", "graphql_resolver", "redis_client", "email_service"]
    name = random.choice(fake_names)
    query = f"Where is the {name} implemented?"

    search1 = name
    search2 = name.replace("_", " ")

    reasoning1 = random.choice(INITIAL_REASONING).format(desc=f"the {name}", search_term=search1)
    reasoning2 = random.choice(NOT_FOUND_REASONING).format(search_term=search1)

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search1}<|end|>"},
            {"role": "result", "content": ""},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning2}<|end|>\n<|search|>{search2}<|end|>"},
            {"role": "result", "content": ""},
            {"role": "assistant", "content": f"<|reasoning|>No results found. The codebase does not contain {name}.<|end|>\n<|evidence|>No implementation of `{name}` was found in the codebase. Searched for: '{search1}', '{search2}'.<|end|>\n<|finish|>"},
        ],
    }


def make_trace_decomposition(chunk: dict, related_chunk: dict) -> dict:
    """Type 4: Query decomposition — complex query broken into subqueries."""
    name = chunk["name"]
    related_name = related_chunk["name"]

    query = f"How does {name} work and how does it relate to {related_name}?"

    search1 = generate_search_term(chunk)
    search2 = generate_search_term(related_chunk)
    evidence1 = generate_evidence(chunk)
    evidence2 = generate_evidence(related_chunk)

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>This query has two parts: (1) how {name} works, (2) its relationship to {related_name}. Let me search for {name} first.<|end|>\n<|search|>{search1}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found {name}. Now let me search for {related_name} to understand the relationship.<|end|>\n<|search|>{search2}<|end|>"},
            {"role": "result", "content": related_chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found both components. I can now explain how they relate.<|end|>\n<|evidence|>{evidence1}\n\nRelated: {evidence2}<|end|>\n<|finish|>"},
        ],
    }


def make_trace_struct_inspection(chunk: dict) -> dict:
    """Type 6: Struct inspection — enumerate fields."""
    if chunk["type"] not in ("struct",):
        return make_trace_simple(chunk)

    name = chunk["name"]
    query = f"What fields does {name} have?"
    search_term = name
    evidence = generate_evidence(chunk)

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>Looking for the {name} struct definition to enumerate its fields.<|end|>\n<|search|>{search_term}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found the struct definition. I can extract all fields from it.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
        ],
    }


def make_trace_behavior_analysis(chunk: dict) -> dict:
    """Type 7: Function behavior analysis — what does it do?"""
    if chunk["type"] not in ("function",):
        return make_trace_simple(chunk)

    name = chunk["name"]
    query = f"What does {name} do? Explain its behavior."
    search_term = name
    evidence = generate_evidence(chunk)

    info = extract_function_info(chunk["code"], name)
    params = ", ".join(info["params"][:4]) if info["params"] else "no parameters"

    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>I need to find and analyze the {name} function. Let me search for its definition.<|end|>\n<|search|>{search_term}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found the function. It takes {params}. Let me analyze its behavior from the code.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
        ],
    }


# ─── Main generation ─────────────────────────────────────────────────────────

def main():
    print(f"Loading chunks from {CHUNKS_PATH}...")
    with open(CHUNKS_PATH, "r", encoding="utf-8") as f:
        chunks = [json.loads(line) for line in f]
    print(f"  Loaded {len(chunks):,} chunks")

    # Filter to chunks with real names and reasonable size
    good_chunks = [c for c in chunks if c["name"] != "unknown" and c["name"] != "block" and 50 < len(c["code"]) < 3000]
    block_chunks = [c for c in chunks if c["name"] == "block" or c["name"] == "unknown"]
    print(f"  Named chunks (good): {len(good_chunks):,}")
    print(f"  Block chunks: {len(block_chunks):,}")

    # Group by language for related chunk pairing
    by_lang = {}
    for c in good_chunks:
        by_lang.setdefault(c["language"], []).append(c)

    traces = []

    # ─── Generate traces ─────────────────────────────────────────────────────
    print("\nGenerating traces...")

    # Type 1: Simple lookup (30% of traces)
    n_simple = min(2000, len(good_chunks))
    sampled = random.sample(good_chunks, n_simple)
    for chunk in sampled:
        traces.append(make_trace_simple(chunk))
    print(f"  Simple lookup: {n_simple}")

    # Type 2: With noise (20%)
    n_noise = min(1500, len(good_chunks))
    sampled = random.sample(good_chunks, n_noise)
    for chunk in sampled:
        # Pick noise chunks from same language
        lang_chunks = by_lang.get(chunk["language"], good_chunks)
        noise_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
        if len(noise_pool) >= 2:
            traces.append(make_trace_with_noise(chunk, noise_pool))
    print(f"  With noise: {n_noise}")

    # Type 3: Multi-step (15%)
    n_multi = min(1000, len(good_chunks) // 2)
    sampled = random.sample(good_chunks, n_multi)
    for chunk in sampled:
        lang_chunks = by_lang.get(chunk["language"], good_chunks)
        related_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
        if related_pool:
            related = random.choice(related_pool)
            traces.append(make_trace_multi_step(chunk, related))
    print(f"  Multi-step: {n_multi}")

    # Type 4: Query decomposition (10%)
    n_decomp = min(700, len(good_chunks) // 3)
    sampled = random.sample(good_chunks, n_decomp)
    for chunk in sampled:
        lang_chunks = by_lang.get(chunk["language"], good_chunks)
        related_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
        if related_pool:
            related = random.choice(related_pool)
            traces.append(make_trace_decomposition(chunk, related))
    print(f"  Query decomposition: {n_decomp}")

    # Type 5: Not found (5%)
    n_notfound = 300
    for _ in range(n_notfound):
        traces.append(make_trace_not_found())
    print(f"  Not found: {n_notfound}")

    # Type 6: Struct inspection (10%)
    struct_chunks = [c for c in good_chunks if c["type"] == "struct"]
    for chunk in struct_chunks[:500]:
        traces.append(make_trace_struct_inspection(chunk))
    print(f"  Struct inspection: {min(len(struct_chunks), 500)}")

    # Type 7: Behavior analysis (10%)
    func_chunks = [c for c in good_chunks if c["type"] == "function"]
    for chunk in func_chunks[:500]:
        traces.append(make_trace_behavior_analysis(chunk))
    print(f"  Behavior analysis: {min(len(func_chunks), 500)}")

    # Shuffle
    random.shuffle(traces)

    # Write
    print(f"\nTotal traces: {len(traces):,}")
    with open(TRACES_PATH, "w", encoding="utf-8") as f:
        for trace in traces:
            f.write(json.dumps(trace) + "\n")
    print(f"Traces written to {TRACES_PATH}")

    # Stats
    total_size = os.path.getsize(TRACES_PATH)
    print(f"  File size: {total_size / 1e6:.2f} MB")

    # Show a sample
    print("\n" + "=" * 60)
    print("SAMPLE TRACE")
    print("=" * 60)
    sample = traces[0]
    print(f"\nQuery: {sample['query']}")
    for msg in sample["trace"]:
        role = msg["role"]
        content = msg["content"]
        if len(content) > 200:
            content = content[:200] + "..."
        print(f"\n[{role}]")
        print(content)


if __name__ == "__main__":
    main()