"""Goal state verification engine for vm-blackbox.

Provides vm_verify_goal_state — a phase-gated MCP tool that checks all criteria
in a GoalState against real hosts and returns a structured VerificationReport.

Gate behaviour (MCP-03): the tool raises ToolError unless run_status for the
composite key vm_name/scenario_name is "completed" in session state. This ensures
the tool is functionally invisible until a scenario run has finished.

Per-criterion dispatch (GOAL-06):
  WinRM hosts → WinRMClient.run_ps_raw()      (PowerShell scripts embedded below)
  SSH hosts   → _ssh_exec() for shell cmds,
                asyncio TCP probe for port_open
  local hosts → direct Python check (port_open via asyncio TCP probe only)
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING

from fastmcp.exceptions import ToolError

from mcp_vm_blackbox.goal_state_store import composite_key, default_store, goal_state_key, run_status_key
from mcp_vm_blackbox.models import CriterionResult, GoalState, HostRosterEntry, VerificationReport
from mcp_vm_blackbox.ssh_exec import _ssh_exec
from mcp_vm_blackbox.winrm_client import WinRMClient

if TYPE_CHECKING:
    from fastmcp.server.context import Context

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# State key templates
# ---------------------------------------------------------------------------

_store = default_store

# ---------------------------------------------------------------------------
# PowerShell script templates (embedded — no external files)
# ---------------------------------------------------------------------------

PORT_OPEN_PS = """\
$t = New-Object System.Net.Sockets.TcpClient
try {{
    $t.Connect("{host}", {port})
    if ($t.Connected) {{ Write-Output "OPEN" }} else {{ Write-Output "CLOSED" }}
}} catch {{
    Write-Output "CLOSED"
}} finally {{
    $t.Close()
}}
"""

SERVICE_PS = (
    '$s = Get-Service -Name "{service}" -ErrorAction SilentlyContinue; if ($s) {{ $s.Status }} else {{ "NOT_FOUND" }}'
)

PROCESS_PS = (
    '$p = Get-Process -Name "{process}" -ErrorAction SilentlyContinue; if ($p) {{ "RUNNING" }} else {{ "NOT_FOUND" }}'
)

DB_CONNECT_PS = """\
$conn = New-Object System.Data.SqlClient.SqlConnection("{conn_string}")
try {{ $conn.Open(); Write-Output "CONNECTED" }}
catch {{ Write-Output "FAILED: $($_.Exception.Message)" }}
finally {{ $conn.Close() }}
"""

# ---------------------------------------------------------------------------
# TCP port probe helper (used for SSH and local protocol port_open checks)
# ---------------------------------------------------------------------------


async def _tcp_port_open(host: str, port: int, connect_timeout: float = 5.0) -> tuple[bool, str]:
    """TCP connect probe. Returns (is_open, error_message). Never raises."""
    try:
        _, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=connect_timeout)
    except (TimeoutError, OSError) as exc:
        return False, str(exc)
    else:
        writer.close()
        await writer.wait_closed()
        return True, ""


# ---------------------------------------------------------------------------
# Per-protocol criterion dispatch
# ---------------------------------------------------------------------------


def _winrm_check(
    criterion_id: str, criterion_type: str, params: dict[str, str], host_entry: HostRosterEntry
) -> CriterionResult:
    """Dispatch a criterion check via WinRM PowerShell."""
    client = WinRMClient(host=host_entry.ip, port=host_entry.port, username=host_entry.username)

    if criterion_type == "port_open":
        host_param = params.get("host", host_entry.ip)
        port_param = params.get("port", "80")
        script = PORT_OPEN_PS.format(host=host_param, port=port_param)
        stdout, stderr, exit_code = client.run_ps_raw(script)
        actual = stdout.strip()
        passed = actual == "OPEN"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="OPEN",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    if criterion_type == "service_running":
        service = params.get("service_name", "")
        script = SERVICE_PS.format(service=service)
        stdout, stderr, exit_code = client.run_ps_raw(script)
        actual = stdout.strip()
        passed = actual == "Running"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="Running",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    if criterion_type == "process_present":
        process = params.get("process_name", "")
        script = PROCESS_PS.format(process=process)
        stdout, stderr, exit_code = client.run_ps_raw(script)
        actual = stdout.strip()
        passed = actual == "RUNNING"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="RUNNING",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    if criterion_type == "db_connect":
        conn_string = params.get("conn_string", "")
        script = DB_CONNECT_PS.format(conn_string=conn_string)
        stdout, stderr, exit_code = client.run_ps_raw(script)
        actual = stdout.strip()
        passed = actual == "CONNECTED"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="CONNECTED",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    # Unsupported criterion type for WinRM
    return CriterionResult(
        criterion_id=criterion_id,
        criterion_type=criterion_type,
        host=host_entry.name,
        expected="supported",
        actual=f"criterion_type '{criterion_type}' not supported for WinRM protocol in Phase 2",
        passed=False,
        stdout="",
        stderr="",
        exit_code=1,
    )


async def _ssh_check(
    criterion_id: str, criterion_type: str, params: dict[str, str], host_entry: HostRosterEntry
) -> CriterionResult:
    """Dispatch a criterion check via SSH (or TCP probe for port_open)."""
    if criterion_type == "port_open":
        # TCP probe — no SSH command needed
        host_param = params.get("host", host_entry.ip)
        port_param = int(params.get("port", "80"))
        is_open, error = await _tcp_port_open(host_param, port_param)
        actual = "OPEN" if is_open else f"CLOSED: {error}"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="OPEN",
            actual=actual,
            passed=is_open,
            stdout=actual,
            stderr=error,
            exit_code=0 if is_open else 1,
        )

    if criterion_type == "service_running":
        service = params.get("service_name", "")
        command = f"systemctl is-active {service}"
        stdout, stderr, exit_code = await _ssh_exec(
            host=host_entry.ip,
            port=host_entry.port,
            username=host_entry.username,
            command=command,
            identity_file=host_entry.identity_file,
        )
        actual = stdout.strip()
        passed = actual == "active"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="active",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    if criterion_type == "process_present":
        process = params.get("process_name", "")
        command = f"pgrep -x {process}"
        stdout, stderr, exit_code = await _ssh_exec(
            host=host_entry.ip,
            port=host_entry.port,
            username=host_entry.username,
            command=command,
            identity_file=host_entry.identity_file,
        )
        passed = exit_code == 0
        actual = "RUNNING" if passed else "NOT_FOUND"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="RUNNING",
            actual=actual,
            passed=passed,
            stdout=stdout,
            stderr=stderr,
            exit_code=exit_code,
        )

    # db_connect and other types are not supported for SSH in Phase 2
    return CriterionResult(
        criterion_id=criterion_id,
        criterion_type=criterion_type,
        host=host_entry.name,
        expected="supported",
        actual=f"criterion_type '{criterion_type}' not supported for SSH protocol in Phase 2",
        passed=False,
        stdout="",
        stderr="",
        exit_code=1,
    )


async def _local_check(
    criterion_id: str, criterion_type: str, params: dict[str, str], host_entry: HostRosterEntry
) -> CriterionResult:
    """Dispatch a criterion check locally (TCP probe only for Phase 2)."""
    if criterion_type == "port_open":
        host_param = params.get("host", "127.0.0.1")
        port_param = int(params.get("port", "80"))
        is_open, error = await _tcp_port_open(host_param, port_param)
        actual = "OPEN" if is_open else f"CLOSED: {error}"
        return CriterionResult(
            criterion_id=criterion_id,
            criterion_type=criterion_type,
            host=host_entry.name,
            expected="OPEN",
            actual=actual,
            passed=is_open,
            stdout=actual,
            stderr=error,
            exit_code=0 if is_open else 1,
        )

    return CriterionResult(
        criterion_id=criterion_id,
        criterion_type=criterion_type,
        host=host_entry.name,
        expected="supported",
        actual=f"criterion_type '{criterion_type}' not supported for local protocol in Phase 2",
        passed=False,
        stdout="",
        stderr="",
        exit_code=1,
    )


# ---------------------------------------------------------------------------
# Main MCP tool function
# ---------------------------------------------------------------------------


async def vm_verify_goal_state(vm_name: str, scenario_name: str, ctx: Context) -> dict:
    """Verify all criteria in a GoalState against real hosts.

    Gate (MCP-03): raises ToolError unless run_status for vm_name/scenario_name
    is "completed" in session state. Call vm-scenario-runner skill to complete a run
    before verification is available.

    Args:
        vm_name: Target VM name (composite key part 1).
        scenario_name: Scenario name (composite key part 2).
        ctx: FastMCP context providing session state access.

    Returns:
        VerificationReport.model_dump() — JSON-serializable dict.

    Raises:
        ToolError: If the run gate is not open, or goal state cannot be found.
    """
    ck = composite_key(vm_name, scenario_name)

    # --- Phase gate ---
    run_status = await ctx.get_state(run_status_key(ck))
    if run_status != "completed":
        raise ToolError(
            f"Verification gate not open for {ck}.\n"
            "A scenario run must complete before verification tools are available.\n"
            "Use vm-scenario-runner skill to dispatch a run."
        )

    # --- Load GoalState: session state first, file fallback ---
    goal_dict = await ctx.get_state(goal_state_key(ck))
    goal: GoalState | None = None
    if goal_dict is not None:
        try:
            goal = GoalState.model_validate(goal_dict)
        except ValueError:
            logger.warning("Failed to parse goal state from session for %s; falling back to file store", ck)

    if goal is None:
        goal = _store.load(vm_name, scenario_name)

    if goal is None:
        raise ToolError(
            f"Goal state not found for {ck}.\nUse vm_set_goal_state to define the goal state before verification."
        )

    # --- Build host roster index ---
    host_index: dict[str, HostRosterEntry] = {h.name: h for h in goal.hosts}

    # --- Per-criterion dispatch ---
    results: list[CriterionResult] = []
    fail_fast_triggered = False

    for criterion in goal.criteria:
        host_entry = host_index.get(criterion.host_name)
        if host_entry is None:
            result = CriterionResult(
                criterion_id=criterion.criterion_id,
                criterion_type=criterion.criterion_type,
                host=criterion.host_name,
                expected="host in roster",
                actual=f"host '{criterion.host_name}' not found in goal state host roster",
                passed=False,
                stdout="",
                stderr="",
                exit_code=1,
            )
        elif host_entry.protocol == "WinRM":
            result = _winrm_check(
                criterion_id=criterion.criterion_id,
                criterion_type=criterion.criterion_type,
                params=criterion.params,
                host_entry=host_entry,
            )
        elif host_entry.protocol == "SSH":
            result = await _ssh_check(
                criterion_id=criterion.criterion_id,
                criterion_type=criterion.criterion_type,
                params=criterion.params,
                host_entry=host_entry,
            )
        else:  # local
            result = await _local_check(
                criterion_id=criterion.criterion_id,
                criterion_type=criterion.criterion_type,
                params=criterion.params,
                host_entry=host_entry,
            )

        results.append(result)

        if goal.fail_fast and not result.passed:
            fail_fast_triggered = True
            break

    overall_passed = all(r.passed for r in results) and not fail_fast_triggered

    report = VerificationReport(
        vm_name=vm_name,
        scenario_name=scenario_name,
        passed=overall_passed,
        results=results,
        fail_fast_triggered=fail_fast_triggered,
    )
    return report.model_dump()
