src/generate_traces_v2.py

"""
Diverse search-agent trace generator v2.

Generates 20-50k unique, high-quality search agent traces using real code
chunks from the deduplicated corpus. Each trace follows the Toast 1-style
agentic search loop: reasoning -> search -> results -> analysis -> evidence -> finish.

Improvements over generate_traces.py:
  - Combinatorial query phrasings (hundreds of unique templates)
  - Smarter evidence extraction (analyzes actual code structure)
  - More trace types: simple, noise, multi-step, decomposition,
    not-found, struct inspection, behavior analysis, cross-file,
    usage tracking, error-handling patterns
  - Trace-level dedup (no two traces share the same query+search combo)
  - Varied reasoning text (not random.choice from 5 options)
  - Realistic evidence packages that extract key facts, not dump code

Output: data/traces/sft_traces_v2.jsonl
"""

import hashlib
import json
import os
import random
import re

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(PROJECT_DIR, "data")
CHUNKS_PATH = os.path.join(DATA_DIR, "chunks", "chunks.jsonl")
TRACES_DIR = os.path.join(DATA_DIR, "traces")
TRACES_PATH = os.path.join(TRACES_DIR, "sft_traces_v2.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(2026)

# ─── Query phrasing components (combinatorial) ───────────────────────────────
QUERY_VERBS = ["What does", "How does", "Where is", "Explain", "Find",
               "Describe", "What is", "How is", "What are", "Show me",
               "Locate", "Identify", "Detail"]
QUERY_ACTIONS = {
    "function": ["do", "work", "behave", "function", "operate", "process",
                 "handle", "compute", "return", "accept"],
    "struct": ["contain", "have", "store", "represent", "look like",
               "be structured", "be defined"],
    "class": ["do", "work", "be structured", "be organized", "be designed",
              "function", "be implemented"],
    "enum": ["have", "contain", "represent", "be defined", "include"],
    "macro": ["do", "define", "expand to", "be used for"],
    "trait": ["define", "require", "provide", "specify"],
    "interface": ["define", "require", "specify", "declare"],
    "impl": ["implement", "provide", "define"],
    "module": ["contain", "provide", "export", "define"],
    "typedef": ["define", "represent", "be"],
    "block": ["do", "work", "be used for"],
}
QUERY_SUFFIXES = ["?", "in the codebase?", "in this project?",
                  "? Explain briefly.", "? Give me the details.",
                  "", "? I need specifics.", "? Summarize the key points."]

# ─── Reasoning components ────────────────────────────────────────────────────
REASONING_OPENERS = [
    "I need to find {desc}.", "Looking for {desc}.",
    "The query is about {desc}.", "This requires locating {desc}.",
    "Let me search for {desc}.", "I should find {desc} first.",
    "Starting by searching for {desc}.", "The user wants to know about {desc}.",
    "I'll look up {desc}.", "First, I need to find {desc}.",
]
REASONING_PLANS = [
    "Let me search using '{term}'.",
    "I'll query the index for '{term}'.",
    "Searching for '{term}' should find it.",
    "Let me issue a search for '{term}'.",
    "I'll search with the term '{term}'.",
    "A search for '{term}' should surface the relevant code.",
    "Let me look for '{term}' in the codebase.",
]
ANALYSIS_FOUND = [
    "Found the relevant code.", "This result contains what I need.",
    "This is the right code.", "I found the target.",
    "This matches the query.", "The search returned the relevant implementation.",
    "This is exactly what I was looking for.",
]
ANALYSIS_NOISE = [
    "The first result is relevant; the others are not directly related.",
    "Result 1 matches. The remaining results appear to be unrelated code.",
    "I found the target in the first result. Ignoring the noise.",
    "The relevant code is in the first result; the rest don't match.",
    "Only the first result is pertinent to this query.",
]
ANALYSIS_REFINE = [
    "Found the definition. Now let me find where it's used.",
    "I have the implementation. Let me search for usages.",
    "Got the definition. Let me get more context with another search.",
    "Found it. Let me do one more search for completeness.",
    "Now I need to see how this is called. Let me search for that.",
]
NOT_FOUND_REASONING = [
    "No results found for '{term}'. Let me try a broader search.",
    "The search returned nothing for '{term}'. Let me try different terms.",
    "No matches for '{term}'. Let me verify with another query.",
    "Empty results for '{term}'. Let me rephrase the search.",
]


def generate_query(chunk: dict) -> str:
    """Generate a diverse query for a chunk."""
    name = chunk["name"]
    typ = chunk["type"]
    if name == "unknown" or name == "block":
        # Use a key identifier from the code
        ids = re.findall(r"\b(ngx_\w+|def\s+\w+|fn\s+\w+|struct\s+\w+|class\s+\w+)\b", chunk["code"])
        if ids:
            name = ids[0].replace("def ", "").replace("fn ", "").replace("struct ", "").replace("class ", "")
        else:
            words = re.findall(r"\b[a-z_]{4,}\b", chunk["code"])
            name = words[0] if words else "this code"

    verb = random.choice(QUERY_VERBS)
    actions = QUERY_ACTIONS.get(typ, QUERY_ACTIONS["block"])
    action = random.choice(actions)
    suffix = random.choice(QUERY_SUFFIXES)

    if verb in ("Where is", "Locate", "Find", "Show me", "Identify"):
        return f"{verb} {name} {action}{suffix}".replace(f" {action}", "")
    return f"{verb} {name} {action}{suffix}"


def generate_description(chunk: dict) -> str:
    name = chunk["name"]
    typ = chunk["type"]
    lang = chunk["language"]
    type_words = {
        "function": f"the {name} function", "struct": f"the {name} structure",
        "class": f"the {name} class", "enum": f"the {name} enum",
        "macro": f"the {name} macro", "trait": f"the {name} trait",
        "interface": f"the {name} interface", "impl": f"the {name} impl block",
        "module": f"the {name} module", "typedef": f"the {name} typedef",
        "block": f"code related to {name}",
    }
    return type_words.get(typ, f"the {name} definition")


def generate_search_term(chunk: dict) -> str:
    name = chunk["name"]
    if name != "unknown" and name != "block":
        return name
    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 ", "")
    words = re.findall(r"\b[a-z_]{4,}\b", code)
    return words[0] if words else "code"


# ─── Evidence extraction (smarter) ───────────────────────────────────────────

def extract_function_info(code: str, name: str) -> dict:
    info = {"name": name, "params": [], "returns": "", "key_lines": [], "calls": []}
    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()][:6]
    ret_match = re.search(r"->\s*([^{]+)", code)
    if ret_match:
        info["returns"] = ret_match.group(1).strip()
    calls = re.findall(r"\b(\w+)\s*\(", code)
    info["calls"] = [c for c in calls if c != name and c not in ("if", "for", "while", "return", "print")][:5]
    for line in code.split("\n"):
        s = line.strip()
        if not s or s.startswith("//") or s.startswith("#"):
            continue
        if any(kw in s for kw in ["return ", "if ", "for ", "while ", "raise ", "throw ", "Err(", "Ok(", "assert", "break", "continue"]):
            info["key_lines"].append(s)
    info["key_lines"] = info["key_lines"][:4]
    return info


def extract_struct_info(code: str, name: str) -> dict:
    fields = []
    for pat in [re.compile(r"^\s+(\w+)\s+(\w+);", re.M),
                re.compile(r"^\s+(\w+):\s+([^,;]+)", re.M),
                re.compile(r"^\s+(self\.)?(\w+)\s*[:=]", re.M),
                re.compile(r"^\s+(\w+)\s+\w+", re.M)]:
        for m in pat.findall(code):
            fn = m[-1] if isinstance(m, tuple) else m
            if fn and fn not in ("self", "pub", "fn", "def", "class", "struct", "const", "var"):
                fields.append(fn)
    return {"name": name, "fields": fields[:12]}


def generate_evidence(chunk: dict) -> str:
    code = chunk["code"]
    name = chunk["name"]
    typ = chunk["type"]
    lang = chunk["language"]
    filepath = chunk["filepath"]
    repo = chunk.get("repo", "")

    if typ == "function":
        info = extract_function_info(code, name)
        parts = [f"`{name}` is a {lang} function."]
        if info["params"]:
            params_str = ", ".join(info["params"])
            parts.append(f"- Parameters: {params_str}")
        if info["returns"]:
            parts.append(f"- Returns: {info['returns']}")
        if info["key_lines"]:
            parts.append("- Key behavior:")
            for line in info["key_lines"][:3]:
                parts.append(f"  - `{line}`")
        if info["calls"]:
            parts.append(f"- Calls: {', '.join(info['calls'])}")
        parts.append(f"- Source: {filepath}")
        return "\n".join(parts)

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

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

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

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

    elif typ in ("trait", "interface"):
        methods = re.findall(r"(?:fn |def |function )\s*(\w+)", code)
        parts = [f"`{name}` is a {lang} {typ}."]
        if methods:
            parts.append(f"- Required methods: {', '.join(methods[:8])}")
        parts.append(f"- Source: {filepath}")
        return "\n".join(parts)

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


# ─── Trace builders ──────────────────────────────────────────────────────────

def make_reasoning(desc: str, term: str) -> str:
    opener = random.choice(REASONING_OPENERS).format(desc=desc)
    plan = random.choice(REASONING_PLANS).format(term=term)
    return f"{opener} {plan}"


def trace_key(query: str, searches: list[str]) -> str:
    """Dedup key for a trace."""
    return hashlib.md5((query + "|" + "|".join(searches)).encode()).hexdigest()


def make_trace_simple(chunk: dict) -> dict:
    query = generate_query(chunk)
    term = generate_search_term(chunk)
    desc = generate_description(chunk)
    evidence = generate_evidence(chunk)
    reasoning1 = make_reasoning(desc, term)
    analysis = random.choice(ANALYSIS_FOUND)
    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{term}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>{analysis} {desc.replace('the ', 'This is ', 1)}.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
        ],
    }


def make_trace_with_noise(chunk: dict, noise_chunks: list) -> dict:
    query = generate_query(chunk)
    term = generate_search_term(chunk)
    desc = generate_description(chunk)
    evidence = generate_evidence(chunk)
    reasoning1 = make_reasoning(desc, term)
    n_noise = random.choice([1, 2])
    noise = random.sample(noise_chunks, min(n_noise, len(noise_chunks)))
    analysis = random.choice(ANALYSIS_NOISE)
    trace = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
        {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|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}<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"})
    return {"query": query, "trace": trace}


def make_trace_multi_step(chunk: dict, related: dict) -> dict:
    name = chunk["name"]
    rname = related["name"]
    query = f"What does {name} do and where is it used?"
    term1 = generate_search_term(chunk)
    term2 = generate_search_term(related)
    desc = generate_description(chunk)
    evidence1 = generate_evidence(chunk)
    evidence2 = generate_evidence(related)
    reasoning1 = make_reasoning(desc, term1)
    reasoning2 = random.choice(ANALYSIS_REFINE)
    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{term1}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found {desc}. {reasoning2}<|end|>\n<|search|>{term2}<|end|>"},
            {"role": "result", "content": related["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_decomposition(chunk: dict, related: dict) -> dict:
    name = chunk["name"]
    rname = related["name"]
    query = f"How does {name} work and how does it relate to {rname}?"
    term1 = generate_search_term(chunk)
    term2 = generate_search_term(related)
    evidence1 = generate_evidence(chunk)
    evidence2 = generate_evidence(related)
    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 {rname}. Let me search for {name} first.<|end|>\n<|search|>{term1}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found {name}. Now let me search for {rname} to understand the relationship.<|end|>\n<|search|>{term2}<|end|>"},
            {"role": "result", "content": related["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_not_found() -> dict:
    fake_names = ["database_pool", "auth_middleware", "cache_manager", "session_store",
                  "websocket_handler", "graphql_resolver", "redis_client", "email_service",
                  "rate_limiter", "circuit_breaker", "event_bus", "message_queue",
                  "load_balancer", "api_gateway", "service_mesh", "config_center"]
    name = random.choice(fake_names)
    query = f"Where is the {name} implemented?"
    search1 = name
    search2 = name.replace("_", " ")
    reasoning1 = make_reasoning(f"the {name}", search1)
    reasoning2 = random.choice(NOT_FOUND_REASONING).format(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_struct_inspection(chunk: dict) -> dict:
    if chunk["type"] != "struct":
        return make_trace_simple(chunk)
    name = chunk["name"]
    query = f"What fields does {name} have?"
    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|>{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(chunk: dict) -> dict:
    if chunk["type"] != "function":
        return make_trace_simple(chunk)
    name = chunk["name"]
    query = f"What does {name} do? Explain its behavior."
    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|>{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|>"},
        ],
    }


def make_trace_cross_file(chunk: dict, related: dict) -> dict:
    """Cross-file: definition in one file, usage in another."""
    name = chunk["name"]
    rname = related["name"]
    query = f"How is {name} used across the codebase?"
    term1 = generate_search_term(chunk)
    term2 = generate_search_term(related)
    evidence1 = generate_evidence(chunk)
    evidence2 = generate_evidence(related)
    return {
        "query": query,
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": query},
            {"role": "assistant", "content": f"<|reasoning|>I need to find the definition of {name} and then trace its usages across files.<|end|>\n<|search|>{term1}<|end|>"},
            {"role": "result", "content": chunk["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found the definition in {chunk['filepath']}. Now let me search for where it's called or referenced.<|end|>\n<|search|>{term2}<|end|>"},
            {"role": "result", "content": related["code"]},
            {"role": "assistant", "content": f"<|reasoning|>Found a usage in {related['filepath']}. I now have both the definition and a call site.<|end|>\n<|evidence|>Definition ({chunk['filepath']}):\n{evidence1}\n\nUsage ({related['filepath']}):\n{evidence2}<|end|>\n<|finish|>"},
        ],
    }


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

def main():
    os.makedirs(TRACES_DIR, exist_ok=True)

    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")

    good_chunks = [c for c in chunks if c["name"] != "unknown" and c["name"] != "block" and 50 < len(c["code"]) < 4000]
    print(f"  Named chunks (good): {len(good_chunks):,}")

    by_lang = {}
    for c in good_chunks:
        by_lang.setdefault(c["language"], []).append(c)

    traces = []
    seen_keys = set()

    def add_trace(trace):
        searches = [m["content"] for m in trace["trace"] if m["role"] == "assistant" and "<|search|>" in m["content"]]
        # Extract search terms
        search_terms = re.findall(r"<\|search\|>(.*?)<\|end\|>", "".join(m["content"] for m in trace["trace"] if m["role"] == "assistant"))
        key = trace_key(trace["query"], search_terms)
        if key in seen_keys:
            return False
        seen_keys.add(key)
        traces.append(trace)
        return True

    def get_related(chunk):
        lang_chunks = by_lang.get(chunk["language"], good_chunks)
        pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
        return random.choice(pool) if pool else chunk

    # Type 1: Simple lookup (25%)
    n_simple = min(8000, len(good_chunks))
    for chunk in random.sample(good_chunks, n_simple):
        add_trace(make_trace_simple(chunk))
    print(f"  Simple lookup: {len(traces)}")

    # Type 2: With noise (20%)
    n_noise = min(6000, len(good_chunks))
    for chunk in random.sample(good_chunks, n_noise):
        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:
            add_trace(make_trace_with_noise(chunk, noise_pool))
    print(f"  With noise: {len(traces)}")

    # Type 3: Multi-step (15%)
    n_multi = min(4000, len(good_chunks) // 2)
    for chunk in random.sample(good_chunks, n_multi):
        add_trace(make_trace_multi_step(chunk, get_related(chunk)))
    print(f"  Multi-step: {len(traces)}")

    # Type 4: Decomposition (10%)
    n_decomp = min(3000, len(good_chunks) // 3)
    for chunk in random.sample(good_chunks, n_decomp):
        add_trace(make_trace_decomposition(chunk, get_related(chunk)))
    print(f"  Decomposition: {len(traces)}")

    # Type 5: Not found (5%)
    for _ in range(1500):
        add_trace(make_trace_not_found())
    print(f"  Not found: {len(traces)}")

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

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

    # Type 8: Cross-file (5%)
    n_cross = min(1500, len(good_chunks) // 4)
    for chunk in random.sample(good_chunks, n_cross):
        add_trace(make_trace_cross_file(chunk, get_related(chunk)))
    print(f"  Cross-file: {len(traces)}")

    random.shuffle(traces)

    # Write
    print(f"\nTotal unique traces: {len(traces):,}")
    total_size = 0
    with open(TRACES_PATH, "w", encoding="utf-8") as f:
        for trace in traces:
            line = json.dumps(trace, ensure_ascii=False)
            f.write(line + "\n")
            total_size += len(line)
    print(f"Traces written to {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"]:
        content = msg["content"]
        if len(content) > 200:
            content = content[:200] + "..."
        print(f"\n[{msg['role']}]")
        print(content)

    # Stats
    stats = {
        "total_traces": len(traces),
        "file_size_mb": total_size / 1e6,
        "unique_queries": len(set(t["query"] for t in traces)),
    }
    stats_path = os.path.join(TRACES_DIR, "trace_stats.json")
    with open(stats_path, "w") as f:
        json.dump(stats, f, indent=2)
    print(f"\nStats: {stats}")


if __name__ == "__main__":
    main()