import os
import subprocess
import json
import sqlite3
import pytest
import shutil
import yaml
import time
import fcntl
import concurrent.futures
import threading
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_e2e_scenario1_standard_lifecycle(mam_sandbox, mock_herdr, mock_agents):
    """
    Scenario 1: Standard Agent Session Lifecycle
    Spawn session, check registry status, query status via status.sh, and stop session.
    """
    tmp_path = mam_sandbox
    create_script = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    status_script = tmp_path / "skills" / "multi-agent-mux-status" / "scripts" / "status.sh"
    stop_script = tmp_path / "skills" / "multi-agent-mux-stop" / "scripts" / "stop_session.sh"
    
    session_name = "e2e-sess1-creator-claude"
    
    # 1. Spawn session using create_session.sh
    cmd_create = [
        "bash", str(create_script),
        "--workspace", str(tmp_path),
        "--agent", "claude",
        "--role", "Creator",
        "--session", session_name
    ]
    res_create = subprocess.run(cmd_create, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_create.returncode == 0, f"Stderr: {res_create.stderr}"
    
    # Verify mock herdr registers it
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    assert session_name in state["agents"]
    assert state["agents"][session_name]["status"] == "running"
    
    # 2. Check status via status.sh --json
    cmd_status = ["bash", str(status_script), "--json"]
    res_status = subprocess.run(cmd_status, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_status.returncode == 0
    
    status_data = json.loads(res_status.stdout)
    sessions_detail = status_data["sessions_detail"]
    assert any(s["name"] == session_name and s["status"] == "running" for s in sessions_detail)
    
    # 3. Stop it via stop_session.sh
    cmd_stop = ["bash", str(stop_script), "--session", session_name]
    res_stop = subprocess.run(cmd_stop, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_stop.returncode == 0
    
    # Verify status in YAML/DB becomes stopped
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    session_entry = [s for s in reg.get("herdr_sessions", []) if s["name"] == session_name][0]
    assert session_entry["status"] == "stopped"


def test_e2e_scenario2_disconnect_resume(mam_sandbox, mock_herdr, mock_agents):
    """
    Scenario 2: Session Disconnect and Resume
    Spawn session, simulate process death in herdr state, run resume_session.sh, and assert resume UUID.
    """
    tmp_path = mam_sandbox
    create_script = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    resume_script = tmp_path / "skills" / "multi-agent-mux-resume" / "scripts" / "resume_session.sh"
    
    session_name = "e2e-sess2-creator-claude"
    
    # 1. Spawn session
    cmd_create = [
        "bash", str(create_script),
        "--workspace", str(tmp_path),
        "--agent", "claude",
        "--role", "Creator",
        "--session", session_name
    ]
    res_create = subprocess.run(cmd_create, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_create.returncode == 0
    
    # Get own UUID from YAML
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    orig_session = reg["herdr_sessions"][0]
    # Simulate first message creating own ID (materialized own ID)
    own_uuid = "e2e-own-uuid-111"
    
    key = str(tmp_path).replace('/', '-').replace('_', '-')
    claude_dir = Path(os.environ.get("CLAUDE_PROJECT_DIR", Path.home() / ".claude" / "projects")) / key
    claude_dir.mkdir(parents=True, exist_ok=True)
    (claude_dir / f"{own_uuid}.jsonl").write_text(json.dumps({"sessionId": own_uuid}) + "\n")
    
    mutation = f"""
for s in d.get('herdr_sessions', []):
    if s.get('name') == '{session_name}':
        s['status'] = 'stopped'
        s['claude_session_id_own'] = '{own_uuid}'
"""
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # 2. Simulate process death in mock herdr state
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    if session_name in state["agents"]:
        del state["agents"][session_name]
    with open(mock_herdr, 'w') as f:
        json.dump(state, f, indent=2)
        
    # Clear herdr calls to isolate assertions
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    state["calls"] = []
    with open(mock_herdr, 'w') as f:
        json.dump(state, f, indent=2)
        
    # 3. Run resume_session.sh
    cmd_resume = [
        "bash", str(resume_script),
        "--workspace", str(tmp_path),
        "--agent", "claude",
        "--session", session_name
    ]
    res_resume = subprocess.run(cmd_resume, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_resume.returncode == 0, f"Stderr: {res_resume.stderr}"
    
    # 4. Assert that it resumes using the correct UUID in the start command arguments
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
        
    calls = state.get("calls", [])
    resume_call = None
    for call in calls:
        if "agent" in call and "start" in call and session_name in call:
            resume_call = call
            break
            
    assert resume_call is not None, f"Could not find resume agent start call in calls: {calls}"
    assert any(own_uuid in arg for arg in resume_call), f"Expected UUID {own_uuid} to be in resume command: {resume_call}"


def test_e2e_scenario3_drift_auto_reconciliation(mam_sandbox, mock_herdr):
    """
    Scenario 3: Drift Detection and Auto-Reconciliation
    Setup drift states (running in herdr but not in YAML registry, and running in YAML registry but terminated in herdr).
    Verify reconcile.sh automatically reconciles both.
    """
    tmp_path = mam_sandbox
    reconcile_script = tmp_path / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
    
    # Drift 1: Running in herdr but not in YAML
    drift_herdr_only = "drift-herdr-only-creator-claude"
    lock_path = str(mock_herdr) + ".lock"
    with open(lock_path, 'w') as lock_f:
        fcntl.flock(lock_f, fcntl.LOCK_EX)
        with open(mock_herdr, 'r') as f:
            state = json.load(f)
        state["agents"][drift_herdr_only] = {
            "status": "running",
            "agent": "claude",
            "cwd": str(tmp_path),
            "pid": 5555,
            "pane_id": "w1:p1",
            "command": "claude",
            "buffer": "Anthropic Claude Ready"
        }
        with open(mock_herdr, 'w') as f:
            json.dump(state, f, indent=2)
        fcntl.flock(lock_f, fcntl.LOCK_UN)
        
    # Drift 2: Running in YAML registry but terminated in herdr
    drift_yaml_only = "drift-yaml-only-creator-claude"
    mutation = f"""
d['herdr_sessions'] = [{{
    'name': '{drift_yaml_only}',
    'status': 'running',
    'role': 'Creator',
    'pane': {{
        'cwd': 'WS_PLACEHOLDER',
        'pid': 6666,
        'cmd': 'claude',
        'cmd_full': 'claude'
    }}
}}]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # Run reconcile.sh --once
    cmd_reconcile = ["bash", str(reconcile_script), "--once"]
    run_env = dict(os.environ)
    run_env["LOCAL_BIN"] = str(tmp_path / "bin")
    run_env["PATH"] = str(tmp_path / "bin") + ":" + os.environ.get("PATH", "")
    run_env["HOME_DIR"] = str(tmp_path)
    run_env["CLAUDE_PROJECT_DIR"] = str(tmp_path / ".claude" / "projects")
    res_recon = subprocess.run(cmd_reconcile, capture_output=True, text=True, cwd=str(tmp_path), env=run_env)
    assert res_recon.returncode == 0, f"Stderr: {res_recon.stderr}"
    
    # Verify YAML/DB states
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    sessions = reg.get("herdr_sessions", [])
    
    # drift-herdr-only-creator-claude should have been auto-registered as running
    sess_herdr_only = [s for s in sessions if s["name"] == drift_herdr_only]
    assert len(sess_herdr_only) == 1
    assert sess_herdr_only[0]["status"] == "running"
    
    # drift-yaml-only-creator-claude should have been auto-terminated
    sess_yaml_only = [s for s in sessions if s["name"] == drift_yaml_only]
    assert len(sess_yaml_only) == 1
    assert sess_yaml_only[0]["status"] == "terminated"


def test_e2e_scenario4_parallel_flock_locking(mam_sandbox, mock_herdr, mock_agents):
    """
    Scenario 4: Parallel Session Operations with flock Locking
    Run multiple concurrent session creation scripts to verify SQLite locking prevents database corruption.
    """
    tmp_path = mam_sandbox
    create_script = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    
    num_sessions = 6
    
    def run_create(i):
        session_name = f"parallel-sess-{i}-creator-claude"
        cmd = [
            "bash", str(create_script),
            "--workspace", str(tmp_path),
            "--agent", "claude",
            "--role", "Creator",
            "--session", session_name
        ]
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=str(tmp_path))
        return res
        
    # Execute in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_sessions) as executor:
        futures = [executor.submit(run_create, i) for i in range(num_sessions)]
        results = [f.result() for f in futures]
        
    # Verify that all succeeded
    for i, res in enumerate(results):
        assert res.returncode == 0, f"Session {i} failed. Stdout: {res.stdout}\nStderr: {res.stderr}"
        
    # Verify all 6 sessions are present in YAML registry
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f) or {}
    sessions = reg.get("herdr_sessions", [])
    
    registered_names = {s["name"] for s in sessions if isinstance(s, dict) and s.get("name")}
    assert len(registered_names) == num_sessions
    for i in range(num_sessions):
        assert f"parallel-sess-{i}-creator-claude" in registered_names


def test_e2e_scenario5_multi_agent_review_loop(mam_sandbox, mock_herdr, mock_agents):
    """
    Scenario 5: Multi-Agent Review Loop
    Run run_loop.sh under simulated conditions where reviewer output is mocked.
    Verify loop terminates correctly with expected PASS and NOT PASS verdicts.
    """
    tmp_path = mam_sandbox
    loop_script = tmp_path / ".agents" / "skills" / "multi-agent-mux-loop" / "scripts" / "run_loop.sh"
    
    # 1. Seed sessions in registry for the worker, reviewer, and planner
    worker_name = "test-worker-creator-claude"
    reviewer_name = "test-reviewer-creator-claude"
    planner_name = "test-planner-creator-claude"
    
    mutation = f"""
d['herdr_sessions'] = [
    {{
        'name': '{worker_name}',
        'status': 'running',
        'role': 'worker',
        'claude_session_id_own': '11111111-1111-4111-a111-111111111111',
        'pane': {{'cwd': 'WS_PLACEHOLDER'}}
    }},
    {{
        'name': '{reviewer_name}',
        'status': 'running',
        'role': 'reviewer',
        'claude_session_id_own': '22222222-2222-4222-a222-222222222222',
        'pane': {{'cwd': 'WS_PLACEHOLDER'}}
    }},
    {{
        'name': '{planner_name}',
        'status': 'running',
        'role': 'planner',
        'claude_session_id_own': '33333333-3333-4333-a333-333333333333',
        'pane': {{'cwd': 'WS_PLACEHOLDER'}}
    }}
]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # Seed mock herdr state with these running sessions to satisfy has-session checks
    lock_path = str(mock_herdr) + ".lock"
    with open(lock_path, 'w') as lock_f:
        fcntl.flock(lock_f, fcntl.LOCK_EX)
        with open(mock_herdr, 'r') as f:
            herdr_state = json.load(f)
        pane_ids = {
            worker_name: "w1:p1",
            reviewer_name: "w1:p2",
            planner_name: "w1:p3"
        }
        for name in [worker_name, reviewer_name, planner_name]:
            herdr_state["agents"][name] = {
                "status": "running",
                "agent": "claude",
                "cwd": str(tmp_path),
                "pid": 9999,
                "pane_id": pane_ids[name],
                "command": "claude",
                "buffer": "Anthropic Claude Ready"
            }
        with open(mock_herdr, 'w') as f:
            json.dump(herdr_state, f, indent=2)
        fcntl.flock(lock_f, fcntl.LOCK_UN)
        
    # Initialize git repo in sandbox for e2e loop review
    subprocess.run(["git", "init"], cwd=str(tmp_path), capture_output=True)
    subprocess.run(["git", "config", "user.email", "test@example.com"], cwd=str(tmp_path), capture_output=True)
    subprocess.run(["git", "config", "user.name", "Test User"], cwd=str(tmp_path), capture_output=True)
    (tmp_path / "README.md").write_text("# Test Repo\n")
    subprocess.run(["git", "add", "."], cwd=str(tmp_path), capture_output=True)
    subprocess.run(["git", "commit", "-m", "initial commit"], cwd=str(tmp_path), capture_output=True)

    # Define mock reviewer and planner outputs
    # Let's mock a scenario:
    # Critique: worker challenge
    # Refinement: refined plan
    # Review: First try NOT PASS, Second try PASS
    job_responses = {
        "Planner": "Refined Plan:\n1. Implement X\n2. Verify X",
        "critique": "Creator Critique: Plan has 1 edge case.",
        "Worker": "Creator Output: Code updated.",
        "Reviewer": "Reviewer verdict:\n\n[VERDICT: PASS]" # will override inside simulator to test iteration logic
    }
    
    # We will run a background simulator thread to resolve jobs
    stop_event = threading.Event()
    
    def simulate_delegate_jobs():
        jobs_dir = tmp_path / ".mam" / "jobs"
        iteration = 1
        
        while not stop_event.is_set():
            if not jobs_dir.exists():
                time.sleep(0.1)
                continue
                
            for job_file in jobs_dir.glob("*.json"):
                try:
                    with open(job_file, 'r+') as f:
                        job = json.load(f)
                        if job.get("status") == "pending":
                            job_id = job["job_id"]
                            role = job.get("role", "Worker")
                            prompt = job.get("prompt", "")
                            
                            # Determine response text
                            response_text = ""
                            if role == "Planner":
                                response_text = job_responses["Planner"]
                            elif "Challenge" in prompt or "Critique" in prompt:
                                response_text = job_responses["critique"]
                            elif role == "Worker":
                                response_text = job_responses["Worker"]
                            elif role == "Reviewer":
                                # For Reviewer, fail the first time, pass the second time
                                if iteration == 1:
                                    response_text = "Review report:\nSome lint issues found.\n\n[VERDICT: NOT PASS]"
                                    iteration += 1
                                else:
                                    response_text = "Review report:\nAll clean.\n\n[VERDICT: PASS]"
                                    
                            # Write final report
                            job_work_dir = jobs_dir / job_id
                            job_work_dir.mkdir(parents=True, exist_ok=True)
                            (job_work_dir / "report-final.md").write_text(response_text)
                            
                            # Complete job
                            job["status"] = "completed"
                            f.seek(0)
                            json.dump(job, f, indent=2)
                            f.truncate()

                            # Publish completed event over MQTT using publish_event.py
                            pub_script = tmp_path / ".agents" / "skills" / "multi-agent-mux-delegate-job" / "scripts" / "publish_event.py"
                            cmd_pub = [
                                sys.executable, str(pub_script),
                                "--registry-dir", str(jobs_dir),
                                "--job", job_id,
                                "--event", "completed",
                                "--detail", f"{role} finished work"
                            ]
                            subprocess.run(cmd_pub, capture_output=True, text=True)
                except Exception:
                    pass
            time.sleep(0.1)
            
    sim_thread = threading.Thread(target=simulate_delegate_jobs)
    sim_thread.daemon = True
    sim_thread.start()
    
    try:
        # Run run_loop.sh
        cmd_loop = [
            "bash", str(loop_script),
            "--creator", worker_name,
            "--reviewer", reviewer_name,
            "--plan",
            "--plan-talk", "1",
            "--max-loop", "3",
            "--task", "Implement feature X and verify"
        ]
        import sys
        run_env = dict(os.environ)
        run_env["DELEGATE_JOB_PYTHON"] = sys.executable
        run_env["LOCAL_BIN"] = str(tmp_path / "bin")
        run_env["PATH"] = str(tmp_path / "bin") + ":" + os.environ.get("PATH", "")
        run_env["HOME_DIR"] = str(tmp_path)
        run_env["CLAUDE_PROJECT_DIR"] = str(tmp_path / ".claude" / "projects")
        res_loop = subprocess.run(cmd_loop, capture_output=True, text=True, cwd=str(tmp_path), env=run_env)
        assert res_loop.returncode == 0, f"Loop failed. Stdout: {res_loop.stdout}\nStderr: {res_loop.stderr}"
        assert "Reviewer 'test-reviewer-creator-claude': NOT PASS" in res_loop.stdout or "Reviewer 'test-reviewer-creator-claude': NOT PASS" in res_loop.stderr or "NOT PASS" in res_loop.stdout
        assert "Reviewer 'test-reviewer-creator-claude': PASS" in res_loop.stdout
        assert "Mux loop finished with 100% PASS verdicts." in res_loop.stdout
        
    finally:
        stop_event.set()
        sim_thread.join(timeout=1.0)
