import os
import subprocess
import json
import pytest
import sys
from pathlib import Path

# Helper to run mutation on agent-sessions.yaml using atomic_dump_yaml in bash
def run_mutation(mam_sandbox, mutation_str, env=None):
    lib_path = mam_sandbox / ".agents" / "skills" / "lib.sh"
    yaml_path = mam_sandbox / ".mam" / "agent-sessions.yaml"
    cmd_str = f"source {lib_path} && atomic_dump_yaml {yaml_path}"
    run_env = dict(os.environ)
    if env:
        run_env.update(env)
    res = subprocess.run(["bash", "-c", cmd_str], input=mutation_str, capture_output=True, text=True, env=run_env)
    return res


def test_unbound_key_handling(mam_sandbox, mock_herdr):
    """
    Verify unbound key handling: herdr has-session/kill-session with trailing parameters
    handles it gracefully (returns status 1 with error) instead of crashing with unbound var or shift errors.
    """
    lib_path = mam_sandbox / ".agents" / "skills" / "lib.sh"
    subprocess.run(["bash", "-c", f"source {lib_path}"], cwd=str(mam_sandbox), env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox)))
    
    shim_path = mam_sandbox / ".mam" / "shim" / "herdr"
    assert shim_path.exists(), "Shim herdr was not created"
    
    # Run has-session with trailing -t
    res = subprocess.run([str(shim_path), "has-session", "-t"], capture_output=True, text=True, env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr)))
    assert res.returncode == 1
    assert "Error: -t requires a value" in res.stderr
    
    # Run kill-session with trailing -t
    res = subprocess.run([str(shim_path), "kill-session", "-t"], capture_output=True, text=True, env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr)))
    assert res.returncode == 1
    assert "Error: -t requires a value" in res.stderr


def test_new_session_fallback_shell(mam_sandbox, mock_herdr):
    """
    Verify new-session: launches shell when no command is provided.
    """
    lib_path = mam_sandbox / ".agents" / "skills" / "lib.sh"
    subprocess.run(["bash", "-c", f"source {lib_path}"], cwd=str(mam_sandbox), env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox)))
    shim_path = mam_sandbox / ".mam" / "shim" / "herdr"
    
    test_shell = "/bin/custom_sh"
    run_env = dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr), SHELL=test_shell)
    
    res = subprocess.run([str(shim_path), "new-session", "-s", "fallback-session"], capture_output=True, text=True, env=run_env)
    assert res.returncode == 0, f"Stderr: {res.stderr}"
    
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    calls = state.get("calls", [])
    
    start_call = None
    for call in calls:
        if "agent" in call and "start" in call and "fallback-session" in call:
            start_call = call
            break
    assert start_call is not None, f"No agent start call found in: {calls}"
    assert start_call[-1] == test_shell, f"Expected last arg to be {test_shell}, got {start_call[-1]}"


def test_ls_key_error(mam_sandbox):
    """
    Verify ls key error: parses agent list without 'name' correctly using fallback keys.
    """
    py_code = """
import sys, json
try:
    data = json.load(sys.stdin)
    res = data.get('result', data)
    for a in res.get('agents', []):
        try:
            name = a.get('name') or a.get('agent') or 'unknown'
            print(f"{name}|0")
        except Exception:
            pass
except Exception:
    pass
"""
    input_json = json.dumps({
        "result": {
            "agents": [
                {"agent": "claude", "status": "running"},
                {"name": "agent-with-name", "agent": "agy"}
            ]
        }
    })
    
    res = subprocess.run([sys.executable, "-c", py_code], input=input_json, capture_output=True, text=True)
    assert res.returncode == 0
    lines = res.stdout.strip().split('\n')
    assert "claude|0" in lines
    assert "agent-with-name|0" in lines


def test_reconcile_relocatability(mam_sandbox, mock_herdr):
    """
    Verify relocatability: reconcile.sh runs when relocated or inside different paths.
    """
    reconcile_script = mam_sandbox / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
    run_dir = mam_sandbox / "some_other_dir"
    run_dir.mkdir()
    
    res = subprocess.run(["bash", str(reconcile_script), "--once", "--emit-diff", "--dry-run"], cwd=str(run_dir), capture_output=True, text=True, env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr)))
    assert res.returncode == 0, f"Failed when run from different directory. Stderr: {res.stderr}"


def test_workspace_server_mapping(mam_sandbox, mock_herdr):
    """
    Verify workspace/server mapping: status output workspace is correct.
    """
    status_script = mam_sandbox / "skills" / "multi-agent-mux-status" / "scripts" / "status.sh"
    session_name = "test-mapping-sess-creator-claude"
    
    mutation = f"""
d['herdr_sessions'] = [{{
    'name': '{session_name}',
    'status': 'running',
    'role': 'Creator',
    'herdr_server': 'my-custom-workspace',
    'pane': {{
        'cwd': 'WS_PLACEHOLDER',
        'pid': 7777,
        'cmd': 'claude',
        'cmd_full': 'claude'
    }}
}}]
""".replace("WS_PLACEHOLDER", str(mam_sandbox))
    res_mut = run_mutation(mam_sandbox, mutation)
    assert res_mut.returncode == 0
    
    res = subprocess.run(["bash", str(status_script), "--json"], capture_output=True, text=True, cwd=str(mam_sandbox), env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr)))
    assert res.returncode == 0, f"Stderr: {res.stderr}"
    
    data = json.loads(res.stdout)
    sess_detail = data["sessions_detail"]
    target_sess = [s for s in sess_detail if s["name"] == session_name][0]
    
    assert target_sess["server"] == "my-custom-workspace"


def test_variable_splicing_injection_safety(mam_sandbox, mock_herdr):
    """
    Verify variable splicing injection: no vulnerabilities or syntax errors remain.
    """
    lib_path = mam_sandbox / ".agents" / "skills" / "lib.sh"
    subprocess.run(["bash", "-c", f"source {lib_path}"], cwd=str(mam_sandbox), env=dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox)))
    shim_path = mam_sandbox / ".mam" / "shim" / "herdr"
    
    adversarial_name = 'my"server; import os; os.system("echo INJECTED")'
    run_env = dict(os.environ, HOME=str(mam_sandbox), WORKSPACE_ROOT=str(mam_sandbox), REAL_HERDR=str(mock_herdr), HERDR_SERVER_NAME=adversarial_name)
    
    res = subprocess.run([str(shim_path), "new-session", "-s", "injection-test-session"], capture_output=True, text=True, env=run_env)
    assert res.returncode == 0, f"Failed with adversarial HERDR_SERVER_NAME. Stderr: {res.stderr}"
    
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    assert any("injection-test-session" in call for call in state["calls"]), "Expected injection-test-session call to succeed"


def test_export_masking_exit_code_preservation():
    """
    Verify export masking: exit codes are preserved when assigning and exporting.
    """
    # 1. Export masked version exits 0 (silent fail)
    cmd_masked = ["bash", "-c", "set -e; export TEST_VAR=$(false); echo 'survived'"]
    res_masked = subprocess.run(cmd_masked, capture_output=True, text=True)
    assert res_masked.returncode == 0
    assert res_masked.stdout.strip() == "survived"
    
    # 2. Fixed split version exits non-zero (preserves failure exit code)
    cmd_fixed = ["bash", "-c", "set -e; TEST_VAR=$(false); export TEST_VAR; echo 'survived'"]
    res_fixed = subprocess.run(cmd_fixed, capture_output=True, text=True)
    assert res_fixed.returncode != 0
    assert res_fixed.stdout.strip() != "survived"