from __future__ import annotations

import os
import sys
import time
import logging
import threading
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

# Add src to sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from api_client import ApiClient, ApiError
from config import Config
from heartbeat import send_once
from command_runner import (
    process_commands,
    load_command_states,
    update_command_state,
    recover_interrupted_commands,
    CommandWorkerManager,
    HANDLERS,
)
from logger import setup_logging


def _fresh_config(tmp: str) -> Config:
    os.environ["WOLFPANEL_DEV_HOME"] = tmp
    os.environ["WOLFPANEL_API_MOCK"] = "1"
    import importlib
    import config as config_module
    importlib.reload(config_module)
    return config_module.load_config()


def _registered_config(tmp: str) -> Config:
    cfg = _fresh_config(tmp)
    from identity import ensure_dirs, store_credentials
    ensure_dirs(cfg)
    store_credentials(cfg, "server_123", "wp_agent_dev_token")
    return _fresh_config(tmp)


def test_command_state_transitions(tmp_path):
    cfg = _registered_config(str(tmp_path))
    
    # Enqueue a mock command and verify state transitions
    api_client = MagicMock()
    api_client.get_pending_commands.return_value = [
        {"id": 101, "command_type": "test.noop", "payload": {}}
    ]
    
    # Register a noop handler
    mock_noop = MagicMock(return_value={"success": True})
    HANDLERS["test.noop"] = mock_noop
    
    # Enable async execution for tests
    os.environ["WOLFPANEL_TEST_ASYNC"] = "1"
    os.environ["WOLFPANEL_MAX_CONCURRENT_COMMANDS"] = "1"
    
    with patch("command_runner.load_agent_token", return_value="wp_agent_dev_token"):
        process_commands(cfg, api_client)
        
        # Give worker a moment to start and run the command
        time.sleep(0.5)
        
        # Verify it was called
        mock_noop.assert_called_once()
        
        # Verify local state transitions to success
        states = load_command_states(cfg.var_dir)
        assert states["101"]["status"] == "success"
        
        # Verify backend was called with running then succeeded
        api_client.post_command_result.assert_any_call(101, {"status": "running"}, "wp_agent_dev_token")
        api_client.post_command_result.assert_any_call(101, {"status": "succeeded", "result": {"success": True}}, "wp_agent_dev_token")


def test_concurrency_limit(tmp_path):
    cfg = _registered_config(str(tmp_path))
    
    event_1 = threading.Event()
    event_2 = threading.Event()
    
    def handler_1(payload):
        event_1.set()
        event_2.wait(timeout=2.0)
        return {"success": True}
        
    def handler_2(payload):
        return {"success": True}
        
    HANDLERS["test.slow1"] = handler_1
    HANDLERS["test.slow2"] = handler_2
    
    api_client = MagicMock()
    # Queue both commands
    api_client.get_pending_commands.return_value = [
        {"id": 201, "command_type": "test.slow1", "payload": {}},
        {"id": 202, "command_type": "test.slow2", "payload": {}}
    ]
    
    # Limit concurrency to 1 worker
    os.environ["WOLFPANEL_TEST_ASYNC"] = "1"
    os.environ["WOLFPANEL_MAX_CONCURRENT_COMMANDS"] = "1"
    
    with patch("command_runner.load_agent_token", return_value="wp_agent_dev_token"):
        process_commands(cfg, api_client)
        
        # Wait until slow1 starts
        event_1.wait(timeout=1.0)
        
        # Verify states
        states = load_command_states(cfg.var_dir)
        assert states["201"]["status"] == "running"
        assert states["202"]["status"] == "pending"  # Should be queued but not running because concurrency is 1
        
        # Unblock slow1
        event_2.set()
        time.sleep(0.5)
        
        # Now slow2 should run and complete
        states = load_command_states(cfg.var_dir)
        assert states["201"]["status"] == "success"
        assert states["202"]["status"] == "success"


def test_command_timeout(tmp_path):
    cfg = _registered_config(str(tmp_path))
    
    event = threading.Event()
    
    def slow_handler(payload):
        event.wait(timeout=5.0)
        return {"success": True}
        
    HANDLERS["test.timeout"] = slow_handler
    
    api_client = MagicMock()
    api_client.get_pending_commands.return_value = [
        # Set short timeout of 1 second in payload
        {"id": 301, "command_type": "test.timeout", "payload": {"timeout": 1}}
    ]
    
    os.environ["WOLFPANEL_TEST_ASYNC"] = "1"
    os.environ["WOLFPANEL_MAX_CONCURRENT_COMMANDS"] = "1"
    
    with patch("command_runner.load_agent_token", return_value="wp_agent_dev_token"):
        process_commands(cfg, api_client)
        
        # Wait for timeout to expire (timeout is 1s, wait 1.5s)
        time.sleep(1.5)
        
        # Verify state is failed due to timeout
        states = load_command_states(cfg.var_dir)
        assert states["301"]["status"] == "failed"
        
        # Verify result has timeout error
        api_client.post_command_result.assert_any_call(
            301,
            {"status": "failed", "result": {"error": "Command execution timed out after 1 seconds"}},
            "wp_agent_dev_token"
        )


def test_startup_recovery(tmp_path):
    cfg = _registered_config(str(tmp_path))
    
    # Manually populate command states with some running and pending commands
    update_command_state(cfg.var_dir, "401", status="running", command_type="some_action", updated_at="old_time")
    update_command_state(cfg.var_dir, "402", status="pending", command_type="some_action", updated_at="old_time")
    update_command_state(cfg.var_dir, "403", status="success", command_type="some_action", updated_at="old_time")
    
    api_client = MagicMock()
    
    with patch("command_runner.load_agent_token", return_value="wp_agent_dev_token"):
        # Run recovery
        recover_interrupted_commands(cfg, api_client)
        
        # Verify states are updated to failed
        states = load_command_states(cfg.var_dir)
        assert states["401"]["status"] == "failed"
        assert states["402"]["status"] == "failed"
        assert states["403"]["status"] == "success"  # Unchanged
        
        # Verify API was notified
        api_client.post_command_result.assert_any_call(
            401,
            {"status": "failed", "result": {"error": "Agent restarted during command execution"}},
            "wp_agent_dev_token"
        )
        api_client.post_command_result.assert_any_call(
            402,
            {"status": "failed", "result": {"error": "Agent restarted during command execution"}},
            "wp_agent_dev_token"
        )


def test_split_logging(tmp_path):
    cfg = _registered_config(str(tmp_path))
    
    # Configure logging to files inside tmp_path
    # Clean up previous handlers
    root_logger = logging.getLogger("wolfpanel")
    for handler in list(root_logger.handlers):
        root_logger.removeHandler(handler)
        
    import logger
    logger._configured = False
    setup_logging(cfg.agent_log, level=logging.DEBUG)
    
    # Get loggers
    hb_log = logging.getLogger("wolfpanel.heartbeat")
    run_log = logging.getLogger("wolfpanel.runner")
    
    hb_log.info("this is a heartbeat log message")
    run_log.info("this is a runner log message")
    
    # Flush and close logging
    logging.shutdown()
    
    agent_log_file = cfg.agent_log
    jobs_log_file = cfg.agent_log.parent / "jobs.log"
    
    assert agent_log_file.is_file()
    assert jobs_log_file.is_file()
    
    agent_content = agent_log_file.read_text(encoding="utf-8")
    jobs_content = jobs_log_file.read_text(encoding="utf-8")
    
    assert "this is a heartbeat log message" in agent_content
    assert "this is a runner log message" not in agent_content
    
    assert "this is a runner log message" in jobs_content
    assert "this is a heartbeat log message" not in jobs_content


# Cleanup at the end of tests
@pytest.fixture(autouse=True)
def clean_env():
    # Store originals
    orig_async = os.environ.get("WOLFPANEL_TEST_ASYNC")
    orig_max = os.environ.get("WOLFPANEL_MAX_CONCURRENT_COMMANDS")
    
    # Pre-test cleanup
    import command_runner
    if command_runner._worker_manager is not None:
        command_runner._worker_manager.shutdown()
        command_runner._worker_manager = None
        
    yield
    
    # Restore
    if orig_async is not None:
        os.environ["WOLFPANEL_TEST_ASYNC"] = orig_async
    else:
        os.environ.pop("WOLFPANEL_TEST_ASYNC", None)
        
    if orig_max is not None:
        os.environ["WOLFPANEL_MAX_CONCURRENT_COMMANDS"] = orig_max
    else:
        os.environ.pop("WOLFPANEL_MAX_CONCURRENT_COMMANDS", None)

    # Post-test cleanup
    if command_runner._worker_manager is not None:
        command_runner._worker_manager.shutdown()
        command_runner._worker_manager = None