#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
QUALITY_FILE="${QUALITY_FILE:-$ROOT_DIR/training_data/lora/mcp_call_quality.jsonl}"
GOOD_FILE="${GOOD_FILE:-$ROOT_DIR/training_data/lora/mcp_calls_good.jsonl}"
BAD_FILE="${BAD_FILE:-$ROOT_DIR/training_data/lora/mcp_calls_bad.jsonl}"

python3 - "$QUALITY_FILE" "$GOOD_FILE" "$BAD_FILE" <<'PY'
import json
import re
import sys
from pathlib import Path
from collections import Counter

quality=Path(sys.argv[1])
good=Path(sys.argv[2])
bad=Path(sys.argv[3])

for p in (quality,good,bad):
    if not p.exists():
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text('',encoding='utf-8')

schema_re = re.compile(
    r"invalid_action|invalid[_\s-]?json|json schema|schema|missing_required_args|"
    r"schema_boundary_violation|task_id_missing|taskitems_missing_or_invalid|"
    r"task_entry_invalid|no_sections_found|invalid[_\s-]?arg|field.+missing",
    re.IGNORECASE
)
timeout_re = re.compile(
    r"timeout|timed out|max_turns_reached|no_final_action|deadline exceeded|executor timeout",
    re.IGNORECASE
)
transport_re = re.compile(
    r"curl:\s*\(7\)|failed to connect|connection refused|no route to host|network unreachable|dns",
    re.IGNORECASE
)
policy_re = re.compile(
    r"permission denied|forbidden|unauthorized|blocked_by_policy|policy|scope|out[-_\s]?of[-_\s]?scope",
    re.IGNORECASE
)

def infer_label(rec):
    label = rec.get("label")
    if label in ("good", "bad"):
        return label
    src = rec.get("source")
    if src == "pipeline_run":
        return "good"
    if src in ("adapter_loop", "mcp_tool_response", "stage_labeled", "synthetic_counterfactual"):
        return "bad"
    payload = rec.get("payload", {})
    if isinstance(payload, dict):
        if payload.get("stage") == "tool_execution" and payload.get("result") == "success":
            return "good"
    return None

def infer_tool_call_success(rec, label):
    payload = rec.get("payload", {})
    if isinstance(payload, dict):
        if payload.get("stage") == "tool_execution":
            if payload.get("result") == "success":
                return True
            if payload.get("result") in ("error", "failed", "failure"):
                return False
    if label == "good":
        return True
    if label == "bad":
        return False
    return None

def classify_failure(rec, tool_call_success):
    if tool_call_success is True:
        return None
    blob = " ".join(
        str(x) for x in [
            rec.get("malformed_kind", ""),
            rec.get("error", ""),
            rec.get("failure_reason", ""),
            rec.get("failure_detail", ""),
            json.dumps(rec.get("payload", {}), ensure_ascii=False)
        ]
    )
    if transport_re.search(blob):
        return "transport"
    if timeout_re.search(blob):
        return "timeout"
    if policy_re.search(blob):
        return "policy"
    if schema_re.search(blob):
        return "schema"
    return "logic"

def classify_recovery_pattern(rec, tool_call_success):
    if tool_call_success is True:
        return "none"
    mk = " ".join(
        str(x) for x in [
            rec.get("malformed_kind", ""),
            rec.get("error", ""),
            rec.get("failure_reason", ""),
            rec.get("failure_detail", ""),
            json.dumps(rec.get("payload", {}), ensure_ascii=False)
        ]
    ).lower()
    if "retry" in mk:
        return "retry"
    if any(x in mk for x in [
        "missing_required_args", "schema_boundary_violation", "task_id_missing",
        "taskitems_missing_or_invalid", "task_entry_invalid", "no_sections_found",
        "invalid_action_format", "invalid_action"
    ]):
        return "reformulate_args"
    if any(x in mk for x in [
        "wrong_tool_selection", "recovery_path_misroute", "dependency_order_violation",
        "overfetch_context"
    ]):
        return "choose_different_tool"
    if any(x in mk for x in [
        "premature_finalize", "non_terminal_loop", "max_turns_reached", "no_final_action", "blocked"
    ]):
        return "abort"
    return "unknown"

def classify_failure_terminality(recovery_pattern):
    if recovery_pattern == "none":
        return None
    if recovery_pattern in ("retry", "reformulate_args", "choose_different_tool"):
        return "recoverable"
    if recovery_pattern == "abort":
        return "terminal"
    return "unknown"

records=[]
for ln in quality.read_text(encoding='utf-8',errors='ignore').splitlines():
    ln=ln.strip()
    if not ln:
        continue
    try:
        rec=json.loads(ln)
    except Exception:
        continue

    label = infer_label(rec)
    if label is None:
        continue
    rec["label"] = label
    if label == "good":
        rec["malformed_kind"] = None

    tool_call_success = infer_tool_call_success(rec, label)
    rec["tool_call_success"] = tool_call_success
    rec["failure_class"] = classify_failure(rec, tool_call_success)
    rec["recovery_pattern"] = classify_recovery_pattern(rec, tool_call_success)
    rec["failure_terminality"] = classify_failure_terminality(rec["recovery_pattern"])
    rec["quality_schema_version"] = 2
    records.append(rec)

# Deduplicate by serialized record.
seen=set(); dedup=[]
for r in records:
    s=json.dumps(r,sort_keys=True,separators=(',',':'))
    if s in seen:
        continue
    seen.add(s)
    dedup.append(r)

quality.write_text(''.join(json.dumps(r,separators=(',',':'))+'\n' for r in dedup),encoding='utf-8')

good_rows=[r for r in dedup if r.get('label')=='good']
bad_rows=[r for r in dedup if r.get('label')=='bad']

good.write_text(''.join(json.dumps(r,separators=(',',':'))+'\n' for r in good_rows),encoding='utf-8')
bad.write_text(''.join(json.dumps(r,separators=(',',':'))+'\n' for r in bad_rows),encoding='utf-8')

fc = Counter((r.get("failure_class") or "none") for r in dedup)
rp = Counter((r.get("recovery_pattern") or "none") for r in dedup)
print(f'reclassified total={len(dedup)} good={len(good_rows)} bad={len(bad_rows)}')
print("failure_class_counts", dict(fc))
print("recovery_pattern_counts", dict(rp))
PY