import os
import subprocess
import json
import sqlite3
import pytest
import shutil
import yaml
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_stop_session_unvalidated_agent(mam_sandbox, mock_herdr, mock_agents):
    """
    Test case 1: Unvalidated agent argument in stop_session.sh
    Verify that calling stop_session.sh with an invalid agent name returns code 2
    and exits with a clear error message to stderr.
    """
    tmp_path = mam_sandbox
    stop_script = tmp_path / "skills" / "multi-agent-mux-stop" / "scripts" / "stop_session.sh"
    
    # 1. Register a session first
    session_name = "invalid-agent-session-creator-claude"
    mutation = """
d['herdr_sessions'] = [{
    'name': 'invalid-agent-session-creator-claude',
    'status': 'running',
    'role': 'Creator',
    'pane': {
        'cwd': 'WS_PLACEHOLDER',
        'pid': 8888,
        'cmd': 'claude',
        'cmd_full': 'claude'
    }
}]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # 2. Try stopping with --agent invalid_agent and --purge-conversation
    cmd_stop = [
        "bash", str(stop_script),
        "--session", session_name,
        "--agent", "invalid_agent",
        "--purge-conversation",
        "--yes"
    ]
    res_stop = subprocess.run(cmd_stop, capture_output=True, text=True, cwd=str(tmp_path))
    
    # It returns code 2
    assert res_stop.returncode == 2
    
    # And prints invalid agent type error to stderr
    assert "ERROR: invalid agent type 'invalid_agent'" in res_stop.stderr
    
    # The registry entry is NOT removed
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    assert len(reg.get("herdr_sessions", [])) == 1


def test_special_character_workspace_slug_interpretation(mam_sandbox):
    """
    Test case 2: Special character workspace slug interpretation
    Verify that workspace paths with only special characters result in derived session names
    that do NOT start with dashes, preventing option flags injection.
    """
    tmp_path = mam_sandbox
    lib_path = tmp_path / ".agents" / "skills" / "lib.sh"
    
    # Workspace consisting of purely special characters
    special_ws = tmp_path / "@#$*" / "@#$*"
    special_ws.mkdir(parents=True, exist_ok=True)
    
    cmd = f"source {lib_path} && derive_session_name '{special_ws}' 'claude'"
    res = subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 0
    derived_name = res.stdout.strip()
    
    # Confirm that the derived name does NOT start with a dash and is default "ws"
    assert not derived_name.startswith("-")
    assert derived_name == "ws-creator-claude"


def test_monitor_reconcile_race_on_purge(mam_sandbox, mock_herdr):
    """
    Test case 3: Concurrency/Race condition during session purge
    Verify that if a session is being purged, the presence of the purging lock file
    prevents concurrent monitor checks from auto-registering it back.
    """
    tmp_path = mam_sandbox
    reconcile_script = tmp_path / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
    
    # 1. Register a running session in herdr and YAML
    session_name = "purged-race-creator-claude"
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    state["agents"][session_name] = {
        "status": "running",
        "agent": "claude",
        "cwd": str(tmp_path),
        "pid": 9999,
        "pane_id": "w1:p1",
        "command": "claude"
    }
    with open(mock_herdr, 'w') as f:
        json.dump(state, f, indent=2)
        
    mutation = f"""
d['herdr_sessions'] = [{{
    'name': '{session_name}',
    'status': 'running',
    'role': 'Creator',
    'pane': {{
        'cwd': 'WS_PLACEHOLDER',
        'pid': 9999,
        'cmd': 'claude',
        'cmd_full': 'claude'
    }}
}}]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # 2. Simulate a purge: stop_session.sh deletes the registry row, but herdr session takes a moment to die.
    # We remove it from the YAML registry, and create the purging lock file.
    mutation_purge = f"""
d['herdr_sessions'] = [s for s in d.get('herdr_sessions', []) if s.get('name') != '{session_name}']
"""
    res_purge_mut = run_mutation(tmp_path, mutation_purge)
    assert res_purge_mut.returncode == 0
    
    # Create the purging lock file
    purging_file = tmp_path / ".mam" / f"purging-{session_name}"
    purging_file.parent.mkdir(parents=True, exist_ok=True)
    purging_file.touch()
    
    # Confirm YAML registry has no sessions
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    assert len(reg.get("herdr_sessions", [])) == 0
    
    # 3. Run reconcile.sh once. It should NOT auto-register it back because the purging lock file exists!
    cmd_reconcile = ["bash", str(reconcile_script), "--once", "--emit-diff"]
    res_rec = subprocess.run(cmd_reconcile, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_rec.returncode == 0
    
    # Verify that the session has NOT been auto-registered back in YAML
    with open(yaml_path, 'r') as f:
        reg_after = yaml.safe_load(f)
    sessions_after = reg_after.get("herdr_sessions", [])
    assert len(sessions_after) == 0


def test_mqtt_hmac_and_seq_validation(mam_sandbox):
    """
    Test case 4: MQTT HMAC and sequence number validation
    Verify that mqtt_common correctly verifies HMAC signatures and drops messages with
    invalid signatures or out-of-order sequence numbers.
    """
    # Import mqtt_common from sandboxed folder
    sys.path.insert(0, str(mam_sandbox / ".agents" / "skills" / "multi-agent-mux-delegate-job" / "scripts"))
    import mqtt_common
    
    # 1. verify_hmac behaviour with no auth_token (PoC mode)
    payload_poc = {"data": {"val": 123}}
    assert mqtt_common.verify_hmac(payload_poc, None) is True
    
    # 2. verify_hmac with auth_token and valid HMAC
    auth_token = "mysecrettoken"
    payload = {
        "job_id": "job123",
        "event": "started",
        "seq": 1,
        "data": {
            "val": 123
        }
    }
    
    # Compute valid signature
    msg = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
    import hmac
    import hashlib
    sig = hmac.new(auth_token.encode(), msg, hashlib.sha256).hexdigest()
    
    payload_signed = dict(payload)
    payload_signed["data"] = dict(payload["data"])
    payload_signed["data"]["hmac_sig"] = sig
    
    assert mqtt_common.verify_hmac(payload_signed, auth_token) is True
    
    # 3. verify_hmac with invalid signature
    payload_signed["data"]["hmac_sig"] = "invalidsignature"
    assert mqtt_common.verify_hmac(payload_signed, auth_token) is False