#!/usr/bin/env python3
"""
build_4b_v43_patch2.py — Second surgical patch targeting 8 specific BFCL failures.

86.9% → 100% target. Exact failures addressed:
  1. knowledge_search vs session_search_memory (3 failures)
  2. session_task_route for "local or cloud?" (1 failure)
  3. session_delete_memory (hallucinated) → session_forget_memory (1 failure)
  4. session_init (hallucinated) → session_load_context (1 failure)
  5. knowledge_forget vs session_forget_memory (1 failure)
  6. session_save_experience vs session_save_ledger on followup (1 failure)
  7. CS abstain even when "retry/backoff" appears in prompt (1 failure)
"""
import json, random
from pathlib import Path

random.seed(2028)

SYS_PROMPT = (
    "You are Synalux, a memory-augmented coding and clinical reasoning assistant. "
    "You have access to Prism Memory tools (session_save_ledger, session_load_context, "
    "session_search_memory, session_save_handoff, session_forget_memory, session_health_check, "
    "session_compact_ledger, session_export_memory, session_task_route, session_save_experience, "
    "session_synthesize_edges, session_backfill_links, knowledge_search, knowledge_forget, "
    "knowledge_upvote, knowledge_downvote, knowledge_set_retention, session_save_image, session_view_image) "
    "and 13 multimodal tool modules (image_gen, office, web_scraper, browser, tts, ocr, git, "
    "terminal, deps_scanner, hipaa, data_graph, templates, pdf_parser). "
    "Think step-by-step before answering. When the user references past work, prior decisions, "
    "or stored context, use the appropriate Prism Memory tool. "
    "TOOL DISTINCTION: Use knowledge_search to query the persistent knowledge base (accumulated "
    "documentation, best practices, reusable insights). Use session_search_memory to find past "
    "session work, project history, prior conversations, and what we worked on before. "
    "Use session_task_route when asked whether the local or cloud agent should handle a task. "
    "Format tool calls inside <tool_call>...</tool_call> JSON blocks with fields 'name' and 'arguments'. "
    "If no tool is needed, answer directly in plain text. "
    "ABSTAIN for general programming questions, CS concepts (algorithms, data structures, "
    "networking, design patterns, frameworks), greetings, and capability questions — even if "
    "the question mentions words like 'retry', 'session', 'memory', or 'knowledge' in a CS context."
)

def ex(user, tool_name, args):
    args_json = json.dumps(args, ensure_ascii=False)
    return {"text": (
        f"<|im_start|>system\n{SYS_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        f"<|im_start|>assistant\n"
        f"<tool_call>\n"
        f'{{\"name\": \"{tool_name}\", \"arguments\": {args_json}}}\n'
        f"</tool_call>\n<|im_end|>"
    )}

def ex_abstain(user, reply):
    return {"text": (
        f"<|im_start|>system\n{SYS_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        f"<|im_start|>assistant\n{reply}<|im_end|>"
    )}

def ex_multiturn(user, tool1, args1, tool_resp, tool2, args2):
    a1 = json.dumps(args1, ensure_ascii=False)
    first = f'<tool_call>\n{{"name": "{tool1}", "arguments": {a1}}}\n</tool_call>'
    if tool2 == "NO_TOOL":
        second = args2.get("reply", "Done.")
    else:
        a2 = json.dumps(args2, ensure_ascii=False)
        second = f'<tool_call>\n{{"name": "{tool2}", "arguments": {a2}}}\n</tool_call>'
    return {"text": (
        f"<|im_start|>system\n{SYS_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        f"<|im_start|>assistant\n{first}<|im_end|>\n"
        f"<|im_start|>tool\n{tool_resp}<|im_end|>\n"
        f"<|im_start|>assistant\n{second}<|im_end|>"
    )}

rows = []

# =============================================================================
# FIX 1: knowledge_search vs session_search_memory (40 examples each)
# =============================================================================

# knowledge_search: the persistent knowledge base, accumulated docs, best practices, reusable insights
ks_prompts = [
    "Search our accumulated documentation for {topic}.",
    "Look up {topic} in the knowledge base.",
    "Find {topic} in our knowledge base.",
    "Search knowledge for {topic}.",
    "Query the knowledge base for {topic}.",
    "What does our knowledge base say about {topic}?",
    "Check the accumulated knowledge for {topic}.",
    "Find {topic} in our documentation knowledge.",
    "Search persisted knowledge for {topic}.",
    "Pull up knowledge base entries about {topic}.",
    "Look for {topic} in the knowledge repository.",
    "Find reusable insights about {topic}.",
    "Knowledge base search: {topic}.",
    "Find best practices for {topic} in our knowledge base.",
    "Search the knowledge store for {topic}.",
]
ks_topics = [
    "WebSocket best practices", "retry strategies", "caching patterns",
    "auth flow", "rate limiting", "database indexing", "circuit breaker pattern",
    "API versioning", "error handling strategies", "deployment checklists",
    "code review guidelines", "security best practices", "logging conventions",
    "microservice communication", "data validation patterns",
]
for i in range(40):
    topic = ks_topics[i % len(ks_topics)]
    user = ks_prompts[i % len(ks_prompts)].format(topic=topic)
    rows.append(ex(user, "knowledge_search", {"query": topic}))

# session_search_memory: past sessions, what we worked on, project history, prior decisions
ssm_prompts = [
    "What did we work on last time for {proj}?",
    "Search my session history for {topic}.",
    "Find what we discussed about {topic} in past sessions.",
    "Look up our prior work on {topic}.",
    "What have we worked on related to {topic}?",
    "Find previous decisions about {topic} in my memory.",
    "Search session memory for {topic}.",
    "What did we decide about {topic} last time?",
    "Look through our past sessions for {topic}.",
    "Find recent session work on {topic}.",
]
ssm_topics = [
    "the auth module", "the deploy pipeline", "the payment service", "database migrations",
    "the API gateway", "the caching layer", "the websocket handler", "performance optimization",
]
projs = ["portal", "analytics", "billing", "auth-service", "dashboard"]
for i in range(40):
    topic = ssm_topics[i % len(ssm_topics)]
    proj = projs[i % len(projs)]
    user = ssm_prompts[i % len(ssm_prompts)].format(topic=topic, proj=proj)
    rows.append(ex(user, "session_search_memory", {"query": topic}))

print(f"After FIX 1 (knowledge vs session search): {len(rows)} rows")

# =============================================================================
# FIX 2: session_task_route (30 examples)
# =============================================================================
task_route_prompts = [
    "Should the local agent handle this {task}? If cloud, just tell me.",
    "Route this {task} — local or cloud?",
    "Should I run this {task} locally or use the cloud model?",
    "Task routing for {task}: local agent or cloud?",
    "Is this {task} suitable for the local agent?",
    "Which agent should handle this {task}: local or host?",
    "Route: should local handle this {task}?",
    "Local or cloud for {task}?",
    "Task route check: can local model do this {task}?",
    "Should I use the local model for {task} or route to cloud?",
]
tasks = [
    "TypeScript refactor", "Python debugging", "code review",
    "SQL query optimization", "React component", "security audit",
    "performance profiling", "architecture design", "bug fix",
    "unit test generation",
]
for i in range(30):
    task = tasks[i % len(tasks)]
    user = task_route_prompts[i % len(task_route_prompts)].format(task=task)
    rows.append(ex(user, "session_task_route", {"task": task}))

print(f"After FIX 2 (session_task_route): {len(rows)} rows")

# =============================================================================
# FIX 3: session_forget_memory (not session_delete_memory — doesn't exist)
# =============================================================================
forget_prompts = [
    "Delete memory entry '{mem_id}' — it's outdated.",
    "Remove memory entry {mem_id} from my session memory.",
    "Forget memory entry ID {mem_id}.",
    "Delete specific memory {mem_id}.",
    "Clear out memory entry {mem_id} — it's wrong.",
    "Remove the memory with id {mem_id}.",
    "Erase memory entry {mem_id}.",
    "Drop memory {mem_id} from session.",
]
mem_ids = ["mem-42", "mem-007", "mem-123", "entry-99", "session-mem-5"]
for i in range(20):
    mid = mem_ids[i % len(mem_ids)]
    user = forget_prompts[i % len(forget_prompts)].format(mem_id=mid)
    rows.append(ex(user, "session_forget_memory", {"memory_id": mid}))

print(f"After FIX 3 (session_forget_memory not delete): {len(rows)} rows")

# =============================================================================
# FIX 4: session_load_context for "initialize/start/begin/setup" context
# =============================================================================
init_prompts = [
    "Initialize the session context for project {proj} at the {level} level.",
    "Start up the session context for {proj}.",
    "Begin session with context for {proj}.",
    "Set up context for {proj} project.",
    "Init session for {proj} at {level} level.",
    "Please initialize session context for {proj}.",
    "Start loading context for {proj}.",
    "Open up the context for project {proj}.",
    "Boot up context for {proj}.",
    "Set context for {proj} ({level}).",
]
levels = ["standard", "deep", "shallow", "full"]
for i in range(20):
    proj = projs[i % len(projs)]
    level = levels[i % len(levels)]
    user = init_prompts[i % len(init_prompts)].format(proj=proj, level=level)
    rows.append(ex(user, "session_load_context", {"project": proj, "level": level}))

print(f"After FIX 4 (session_load_context for 'initialize'): {len(rows)} rows")

# =============================================================================
# FIX 5: knowledge_forget vs session_forget_memory
# knowledge_forget = clear knowledge base entries (by category/project)
# session_forget_memory = clear a specific session memory entry (by ID)
# =============================================================================
kf_prompts = [
    "Clear out all old knowledge entries in the '{cat}' category for {proj}.",
    "Remove all {cat} knowledge entries for the {proj} project.",
    "Forget all knowledge about {cat} in {proj}.",
    "Delete {proj} knowledge entries tagged {cat}.",
    "Purge {cat} knowledge for {proj}.",
    "Clear the {cat} knowledge base entries for {proj}.",
    "Remove all {cat}-category knowledge from {proj}.",
    "Delete outdated knowledge in {cat} for {proj}.",
]
cats = ["testing", "deprecated", "v1", "staging", "draft", "archived"]
for i in range(20):
    cat = cats[i % len(cats)]
    proj = projs[i % len(projs)]
    user = kf_prompts[i % len(kf_prompts)].format(cat=cat, proj=proj)
    rows.append(ex(user, "knowledge_forget", {"project": proj, "category": cat}))

print(f"After FIX 5 (knowledge_forget): {len(rows)} rows")

# =============================================================================
# FIX 6: session_save_experience vs session_save_ledger on followup
# session_save_experience = record a correction/insight/learning (event_type matters)
# session_save_ledger = log a session summary/progress
# =============================================================================

# Multi-turn: load context, then log what we EXPERIENCED (correction/insight)
load_experience_chains = [
    ("Load context for {proj} and then log that we tried {what} but should have used {better} instead.",
     "session_load_context", lambda p, w, b: {"project": p},
     '{{"project": "{proj}", "last_summary": "Working on {proj}"}}',
     "session_save_experience", lambda p, w, b: {"project": p, "event_type": "correction",
                                                   "content": f"Tried {w} but should have used {b}"}),
    ("Get {proj} context, then record the correction: used {what} when {better} was better.",
     "session_load_context", lambda p, w, b: {"project": p},
     '{{"project": "{proj}"}}',
     "session_save_experience", lambda p, w, b: {"project": p, "event_type": "correction",
                                                   "content": f"Used {what} instead of {better}"}),
]
whats = ["batch inserts", "polling", "mutex locks", "REST calls", "eager loading"]
betters = ["streaming writes", "webhooks", "read-write locks", "GraphQL", "lazy loading"]

for i in range(15):
    ch = load_experience_chains[i % len(load_experience_chains)]
    proj = projs[i % len(projs)]
    what = whats[i % len(whats)]
    better = betters[i % len(betters)]
    user = ch[0].format(proj=proj, what=what, better=better)
    t1 = ch[2](proj, what, better)
    resp = ch[3].format(proj=proj)
    t2 = ch[5](proj, what, better)
    rows.append(ex_multiturn(user, ch[1], t1, resp, ch[4], t2))

# Distinguish ledger (progress/session log) from experience (insight/correction)
ledger_vs_exp = [
    ("Load {proj} context, then save a session ledger entry about today's progress.",
     "session_load_context", lambda p: {"project": p},
     '{{"project": "{proj}", "last_summary": "Progress on {proj}"}}',
     "session_save_ledger", lambda p: {"project": p, "summary": "Session progress logged"}),
    ("Get {proj} context, then record that we successfully completed the migration.",
     "session_load_context", lambda p: {"project": p},
     '{{"project": "{proj}"}}',
     "session_save_ledger", lambda p: {"project": p, "summary": "Migration completed successfully"}),
]
for i in range(10):
    ch = ledger_vs_exp[i % len(ledger_vs_exp)]
    proj = projs[i % len(projs)]
    user = ch[0].format(proj=proj)
    t1 = ch[2](proj)
    resp = ch[3].format(proj=proj)
    t2 = ch[5](proj)
    rows.append(ex_multiturn(user, ch[1], t1, resp, ch[4], t2))

print(f"After FIX 6 (save_experience vs save_ledger): {len(rows)} rows")

# =============================================================================
# FIX 7: CS abstain even with "retry/backoff/session/memory/knowledge" in CS context
# =============================================================================
cs_with_prism_keywords = [
    ("How do I implement exponential backoff with jitter for API retries?",
     "Exponential backoff with jitter: `wait = random.uniform(0, base * 2^attempt)`. Cap at max_delay. This prevents thundering herd. Pure CS — no Prism tool needed."),
    ("Explain retry semantics in distributed systems.",
     "Retry semantics in distributed systems cover idempotency (safe to retry), at-least-once vs exactly-once delivery, and exponential backoff to avoid overload. This is a CS concept, not a Prism memory operation."),
    ("How does session management work in Django?",
     "Django uses a session middleware that stores a session ID in a cookie and keeps session data server-side (database, cache, or file). Access via `request.session`. Not a Prism memory operation."),
    ("What is memory management in C++?",
     "C++ memory management involves heap allocation (new/delete), stack allocation, RAII, and smart pointers (unique_ptr, shared_ptr). Not related to Prism session memory."),
    ("Explain knowledge graphs and how they work.",
     "Knowledge graphs represent entities and relationships as nodes and edges. They power semantic search and reasoning engines. This is a CS/ML concept — not a Prism knowledge_search operation."),
    ("How do I implement retry logic with exponential backoff in Python?",
     "Use the `tenacity` library or implement manually: `time.sleep(min(max_delay, base * 2**attempt) * random.uniform(0.5, 1.5))`. This is a general Python pattern, not a Prism memory task."),
    ("What is session affinity in load balancers?",
     "Session affinity (sticky sessions) routes all requests from a client to the same backend server, typically using a cookie. This is a networking/infrastructure concept, not a Prism operation."),
    ("How does Redis handle memory eviction?",
     "Redis uses LRU (least recently used), LFU (least frequently used), or TTL-based eviction policies, configured via `maxmemory-policy`. This is a Redis/CS concept, not a Prism memory tool call."),
    ("Explain the actor model for concurrent programming.",
     "The actor model treats concurrent computation as actors that communicate via message passing, with no shared state. Erlang, Akka, and Pony use this model. Pure CS concept."),
    ("What is the difference between a stack and a queue?",
     "Stack: LIFO (last in, first out) — push/pop. Queue: FIFO (first in, first out) — enqueue/dequeue. These are fundamental data structures."),
]
for user, reply in cs_with_prism_keywords:
    rows.append(ex_abstain(user, reply))

print(f"After FIX 7 (CS abstain with prism-like keywords): {len(rows)} rows")

# =============================================================================
# SHUFFLE AND WRITE
# =============================================================================
random.shuffle(rows)
out = Path("/tmp/4b_v43_patch2.jsonl")
out.write_text("\n".join(json.dumps(r, ensure_ascii=False) for r in rows) + "\n")
print(f"\n✅ Wrote {len(rows)} patch2 rows to {out}")