#!/usr/bin/env python3
"""
Retrain safety threshold using chat-context features.

Sends all 80 labeled prompts through the actual chat API (with system prompt,
tools, full context), captures the SAE features from the residuals, and
rebuilds the threshold classifier.

Usage:
    python scripts/retrain_threshold_chat_context.py
"""

import json
import time
import sys
from collections import defaultdict
from pathlib import Path

import requests

ROOT = Path("/data/share169/creditscope")
BENIGN_PATH = ROOT / "benign_results_expanded.json"
ADVERSARIAL_PATH = ROOT / "adversarial_results_expanded.json"
OUTPUT_PATH = ROOT / "circuit_tracer" / "data" / "checkpoints" / "safety_threshold.json"
API_BASE = "http://127.0.0.1:8081"

# ── Load prompts ──────────────────────────────────────────────────────────

benign_prompts = [e["question"].strip() for e in json.loads(BENIGN_PATH.read_text())]
adversarial_prompts = [e["question"].strip() for e in json.loads(ADVERSARIAL_PATH.read_text())]

print(f"Benign: {len(benign_prompts)}, Adversarial: {len(adversarial_prompts)}")
print(f"Total: {len(benign_prompts) + len(adversarial_prompts)} prompts to process")

# ── Login ─────────────────────────────────────────────────────────────────

session = requests.Session()
resp = session.post(f"{API_BASE}/api/auth/login", json={
    "email": "sarelw1@gmail.com", "password": "WEF4"
})
resp.raise_for_status()
print("Logged in")

# ── Send all prompts through chat API ─────────────────────────────────────

all_results: list[dict] = []

for category, prompts in [("benign", benign_prompts), ("adversarial", adversarial_prompts)]:
    for i, prompt in enumerate(prompts):
        started = time.time()
        try:
            resp = session.post(f"{API_BASE}/api/chat", json={
                "message": prompt,
                "cot_config": {"mode": "off", "budget": "none", "visibility": "hidden"},
            }, timeout=120)
            resp.raise_for_status()
            data = resp.json()
            ca = data.get("circuit_analysis", {})
            duration = time.time() - started

            features_by_key: dict[str, float] = {}
            for layer_str, features in ca.get("layer_features", {}).items():
                for f in features:
                    key = f"L{f['layer']}_F{f['feature_idx']}"
                    act = abs(f["activation"])
                    if key not in features_by_key or act > features_by_key[key]:
                        features_by_key[key] = act

            all_results.append({
                "prompt": prompt,
                "category": category,
                "features": features_by_key,
                "total_features": ca.get("total_active_features", 0),
                "num_layers": ca.get("num_layers_analyzed", 0),
            })

            print(f"  [{category[:3].upper()}] {i+1}/{len(prompts)} ({duration:.0f}s) "
                  f"feats={len(features_by_key)} | {prompt[:60]}")

        except Exception as e:
            print(f"  [{category[:3].upper()}] {i+1}/{len(prompts)} ERROR: {e} | {prompt[:60]}")
            all_results.append({
                "prompt": prompt,
                "category": category,
                "features": {},
                "total_features": 0,
                "num_layers": 0,
            })

print(f"\nCompleted {len(all_results)} prompts")

# ── Save raw results ──────────────────────────────────────────────────────

raw_path = ROOT / "circuit_tracer" / "data" / "checkpoints" / "chat_context_features.json"
with open(raw_path, "w") as f:
    json.dump(all_results, f, indent=2)
print(f"Raw features saved to {raw_path}")

# ── Compute threshold ─────────────────────────────────────────────────────

benign_results = [r for r in all_results if r["category"] == "benign" and r["total_features"] > 0]
adv_results = [r for r in all_results if r["category"] == "adversarial" and r["total_features"] > 0]

print(f"\nUsable: {len(benign_results)} benign, {len(adv_results)} adversarial")

# Per-feature stats
all_keys = set()
for r in all_results:
    all_keys.update(r["features"].keys())

feature_stats = []
for key in sorted(all_keys):
    b_vals = [r["features"].get(key, 0) for r in benign_results if key in r["features"]]
    a_vals = [r["features"].get(key, 0) for r in adv_results if key in r["features"]]

    b_mean = sum(b_vals) / len(b_vals) if b_vals else 0
    a_mean = sum(a_vals) / len(a_vals) if a_vals else 0

    parts = key.split("_")
    layer = int(parts[0][1:])
    feat_idx = int(parts[1][1:])

    feature_stats.append({
        "key": key, "layer": layer, "feature_idx": feat_idx,
        "benign_count": len(b_vals), "adv_count": len(a_vals),
        "benign_mean": round(b_mean, 6), "adv_mean": round(a_mean, 6),
        "diff_mean": round(a_mean - b_mean, 6),
        "ratio": round(a_mean / b_mean, 4) if b_mean > 0 else (999 if a_mean > 0 else 0),
    })

# Select discriminative features
discriminative = set()
safety_features = set()

# Adversarial-only features
for f in feature_stats:
    if f["adv_count"] >= 3 and f["benign_count"] == 0 and f["adv_mean"] > 0.01:
        discriminative.add(f["key"])

# Features significantly higher in adversarial
for f in feature_stats:
    if f["adv_count"] >= 5 and f["diff_mean"] > 0.02 and f["adv_mean"] > 0.01:
        discriminative.add(f["key"])

# Features significantly higher in adversarial by ratio
for f in feature_stats:
    if f["adv_count"] >= 5 and f["ratio"] > 1.15 and f["adv_mean"] > 0.01:
        discriminative.add(f["key"])

# Safety features (higher in benign)
for f in feature_stats:
    if f["benign_count"] >= 5 and f["diff_mean"] < -0.02 and f["benign_mean"] > 0.01:
        safety_features.add(f["key"])

print(f"\nDiscriminative features: {len(discriminative)}")
print(f"Safety features: {len(safety_features)}")

# Compute per-trace risk scores
def compute_score(result):
    adv_signal = sum(result["features"].get(k, 0) for k in discriminative)
    safety_signal = sum(result["features"].get(k, 0) for k in safety_features)
    adv_hits = sum(1 for k in discriminative if k in result["features"])
    safety_hits = sum(1 for k in safety_features if k in result["features"])
    return {
        "risk_score": round(adv_signal - safety_signal * 0.5, 4),
        "adv_signal": round(adv_signal, 4),
        "adv_hits": adv_hits,
        "safety_signal": round(safety_signal, 4),
        "safety_hits": safety_hits,
    }

benign_scores = []
adv_scores = []

print(f"\n--- BENIGN ---")
for r in sorted(benign_results, key=lambda x: x["prompt"]):
    s = compute_score(r)
    benign_scores.append(s["risk_score"])
    print(f"  [{s['risk_score']:+8.2f}] adv={s['adv_signal']:6.2f}({s['adv_hits']}) safety={s['safety_signal']:6.2f}({s['safety_hits']}) | {r['prompt'][:65]}")

print(f"\n--- ADVERSARIAL ---")
for r in sorted(adv_results, key=lambda x: x["prompt"]):
    s = compute_score(r)
    adv_scores.append(s["risk_score"])
    print(f"  [{s['risk_score']:+8.2f}] adv={s['adv_signal']:6.2f}({s['adv_hits']}) safety={s['safety_signal']:6.2f}({s['safety_hits']}) | {r['prompt'][:65]}")

# Find optimal thresholds
print(f"\nBenign scores: min={min(benign_scores):+.2f} avg={sum(benign_scores)/len(benign_scores):+.2f} max={max(benign_scores):+.2f}")
print(f"Adversarial:   min={min(adv_scores):+.2f} avg={sum(adv_scores)/len(adv_scores):+.2f} max={max(adv_scores):+.2f}")

all_score_vals = sorted(set(benign_scores + adv_scores))
best_f1 = None
full_recall = None

for i in range(len(all_score_vals) - 1):
    threshold = (all_score_vals[i] + all_score_vals[i + 1]) / 2
    tp = sum(1 for s in adv_scores if s >= threshold)
    fn = sum(1 for s in adv_scores if s < threshold)
    fp = sum(1 for s in benign_scores if s >= threshold)
    tn = sum(1 for s in benign_scores if s < threshold)
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

    if best_f1 is None or f1 > best_f1["f1"]:
        best_f1 = {"threshold": round(threshold, 4), "tp": tp, "fn": fn, "fp": fp, "tn": tn,
                    "precision": round(precision, 4), "recall": round(recall, 4), "f1": round(f1, 4),
                    "accuracy": round((tp + tn) / (tp + fn + fp + tn), 4)}

    if recall == 1.0:
        if full_recall is None or fp < full_recall["fp"]:
            full_recall = {"threshold": round(threshold, 4), "tp": tp, "fn": fn, "fp": fp, "tn": tn,
                           "precision": round(precision, 4), "recall": 1.0,
                           "f1": round(2 * precision / (precision + 1), 4),
                           "accuracy": round((tp + tn) / (tp + fn + fp + tn), 4)}

print(f"\n--- BEST F1 ---")
if best_f1:
    print(f"  Threshold: {best_f1['threshold']}, P={best_f1['precision']:.1%} R={best_f1['recall']:.1%} F1={best_f1['f1']:.1%} Acc={best_f1['accuracy']:.1%}")
    print(f"  TP={best_f1['tp']} FN={best_f1['fn']} FP={best_f1['fp']} TN={best_f1['tn']}")

print(f"\n--- 100% RECALL ---")
if full_recall:
    print(f"  Threshold: {full_recall['threshold']}, P={full_recall['precision']:.1%} R={full_recall['recall']:.1%} Acc={full_recall['accuracy']:.1%}")
    print(f"  TP={full_recall['tp']} FN={full_recall['fn']} FP={full_recall['fp']} TN={full_recall['tn']}")

# Save threshold
disc_list = [f for f in feature_stats if f["key"] in discriminative]
safety_list = [f for f in feature_stats if f["key"] in safety_features]

threshold_config = {
    "best_f1": best_f1,
    "full_recall": full_recall,
    "discriminative_features": sorted(disc_list, key=lambda f: f["diff_mean"], reverse=True),
    "safety_features": sorted(safety_list, key=lambda f: f["diff_mean"]),
    "scoring": {
        "method": "adversarial_signal - 0.5 * safety_signal",
        "context": "chat (system prompt + tools + user message)",
    },
    "dataset": {
        "benign_traces": len(benign_results),
        "adversarial_traces": len(adv_results),
        "total_features": len(all_keys),
    },
}

OUTPUT_PATH.write_text(json.dumps(threshold_config, indent=2))
print(f"\nSaved to {OUTPUT_PATH}")

# Misclassifications
if full_recall:
    thr = full_recall["threshold"]
    print(f"\n=== FALSE POSITIVES at 100% recall (threshold={thr}) ===")
    for r in benign_results:
        s = compute_score(r)
        if s["risk_score"] >= thr:
            print(f"  [{s['risk_score']:+.2f}] {r['prompt'][:80]}")