"""Infrastructure provisioning layer — PostgreSQL (INFRA-01/02/03).

All database interactions use the ``psycopg`` native Python SDK.
No subprocess or shell-out is permitted; fail loud if no SDK path exists.

State is persisted to ``~/.mcp-vm-blackbox/infra/`` and never committed.
"""

from __future__ import annotations

import importlib
import json
import logging
import os
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from fastmcp.exceptions import ToolError
from fastmcp.tools import tool

from mcp_vm_blackbox.goal_state_store import composite_key
from mcp_vm_blackbox.models import InfraDestroyResult, InfraProvisionResult, InfraState, InfraStatusResult

__all__ = ["infra_destroy_postgres", "infra_provision_postgres", "infra_status"]

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Storage root — same parent as goal states and post-mortems
# ---------------------------------------------------------------------------

_INFRA_DIR = Path(os.environ.get("MCP_VM_BLACKBOX_INFRA_DIR", str(Path.home() / ".mcp-vm-blackbox" / "infra")))

_DEFAULT_PG_HOST = os.environ.get("MCP_VM_BLACKBOX_PG_HOST", "localhost")
_DEFAULT_PG_PORT = int(os.environ.get("MCP_VM_BLACKBOX_PG_PORT", "5432"))
_DEFAULT_PG_SUPERUSER = os.environ.get("MCP_VM_BLACKBOX_PG_SUPERUSER", "postgres")
_DEFAULT_PG_SUPERPASSWORD = os.environ.get("MCP_VM_BLACKBOX_PG_SUPERPASSWORD", "postgres")


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _state_path(scenario_key: str) -> Path:
    """Return the JSON file path for the given scenario's infra state."""
    safe_key = scenario_key.replace("/", "__")
    return _INFRA_DIR / f"{safe_key}.json"


def _load_state(scenario_key: str) -> InfraState | None:
    """Load persisted InfraState or return None if not found."""
    path = _state_path(scenario_key)
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text())
        return InfraState.model_validate(data)
    except Exception as exc:  # noqa: BLE001
        logger.warning("Failed to parse infra state at %s: %s", path, exc)
        return None


def _save_state(state: InfraState) -> Path:
    """Persist InfraState to disk and return its path."""
    path = _state_path(state.scenario_key)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(state.model_dump_json(indent=2))
    return path


def _remove_state(scenario_key: str) -> Path | None:
    """Remove the persisted state file if it exists; return its path or None."""
    path = _state_path(scenario_key)
    if path.exists():
        path.unlink()
        return path
    return None


def _load_psycopg() -> Any:
    """Import and return the psycopg module.

    Raises:
        ToolError: When psycopg is not installed.
    """
    try:
        return importlib.import_module("psycopg")
    except ImportError as exc:
        raise ToolError(
            "psycopg is not available. Install it with: pip install psycopg[binary]\n"
            "All database interactions require the native psycopg SDK — "
            "no subprocess fallback is permitted (INFRA-02)."
        ) from exc


def _build_connstr(host: str, port: int, user: str, password: str, dbname: str = "postgres") -> str:
    """Build a psycopg connection string."""
    return f"host={host} port={port} user={user} password={password} dbname={dbname} connect_timeout=5"


async def _verify_connectivity(psycopg: Any, connstr: str) -> bool:
    """Attempt a single connection to verify the database is reachable.

    Returns:
        True on success, False on failure.
    """
    try:
        async with await psycopg.AsyncConnection.connect(connstr, autocommit=True) as conn:
            await conn.execute("SELECT 1")
    except Exception as exc:  # noqa: BLE001
        logger.warning("PostgreSQL connectivity check failed: %s", exc)
        return False
    else:
        return True


# ---------------------------------------------------------------------------
# Public MCP tools
# ---------------------------------------------------------------------------


@tool
async def infra_provision_postgres(
    vm_name: str,
    scenario_name: str,
    db_name: str = "scenario_db",
    db_user: str = "scenario_user",
    db_password: str = "scenario_pass",  # noqa: S107
    pg_host: str = _DEFAULT_PG_HOST,
    pg_port: int = _DEFAULT_PG_PORT,
) -> InfraProvisionResult:
    """Provision a PostgreSQL database for a scenario (INFRA-01/02/03).

    Uses the ``psycopg`` native Python SDK to create the database and user.
    No subprocess or shell-out is permitted.  Verifies connectivity before
    returning.  Persists state to ``~/.mcp-vm-blackbox/infra/``.

    Args:
        vm_name: Target VM name — used to scope the infra state.
        scenario_name: Scenario identifier — used to scope the infra state.
        db_name: PostgreSQL database name to create.
        db_user: PostgreSQL user to create.
        db_password: Password for the new user.
        pg_host: PostgreSQL server hostname.
        pg_port: PostgreSQL server port.

    Returns:
        InfraProvisionResult with endpoint, connectivity, and state_path.

    Raises:
        ToolError: When psycopg is unavailable or the superuser connection fails.
    """
    psycopg = _load_psycopg()
    ck = composite_key(vm_name, scenario_name)
    endpoint = f"{pg_host}:{pg_port}"

    # Connect as superuser to create DB and role.
    superconn_str = _build_connstr(pg_host, pg_port, _DEFAULT_PG_SUPERUSER, _DEFAULT_PG_SUPERPASSWORD)
    try:
        async with await psycopg.AsyncConnection.connect(superconn_str, autocommit=True) as conn:
            # Create user (role) if not exists.
            await conn.execute(
                f"DO $$ BEGIN "
                f"IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '{db_user}') THEN "
                f"CREATE ROLE {db_user} LOGIN PASSWORD '{db_password}'; "
                f"END IF; END $$"
            )
            # Create database if not exists.
            exists = await (await conn.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,))).fetchone()
            if not exists:
                await conn.execute(f"CREATE DATABASE {db_name} OWNER {db_user}")
    except Exception as exc:
        raise ToolError(
            f"PostgreSQL provisioning failed for {ck}: {exc}\n"
            "Ensure the PostgreSQL server is running and superuser credentials are correct."
        ) from exc

    # Verify connectivity with the new user.
    user_connstr = _build_connstr(pg_host, pg_port, db_user, db_password, db_name)
    connected = await _verify_connectivity(psycopg, user_connstr)

    # Persist state.
    state = InfraState(
        scenario_key=ck,
        db_endpoint=endpoint,
        db_name=db_name,
        db_user=db_user,
        db_password=db_password,
        network_config={"pg_host": pg_host, "pg_port": str(pg_port)},
        provisioned_at=datetime.now(UTC).isoformat(),
    )
    state_path = _save_state(state)

    return InfraProvisionResult(
        scenario_key=ck,
        db_endpoint=endpoint,
        db_name=db_name,
        db_user=db_user,
        connected=connected,
        state_path=str(state_path),
    )


@tool
async def infra_destroy_postgres(
    vm_name: str, scenario_name: str, pg_host: str = _DEFAULT_PG_HOST, pg_port: int = _DEFAULT_PG_PORT
) -> InfraDestroyResult:
    """Tear down a PostgreSQL database provisioned for a scenario (INFRA-01/02/03).

    Drops the database and user created by ``infra_provision_postgres``.
    Removes the persisted state file.  Uses psycopg native SDK — no subprocess.

    Args:
        vm_name: Target VM name.
        scenario_name: Scenario identifier.
        pg_host: PostgreSQL server hostname.
        pg_port: PostgreSQL server port.

    Returns:
        InfraDestroyResult with destroyed status and state_path.

    Raises:
        ToolError: When psycopg is unavailable or no infra state exists.
    """
    ck = composite_key(vm_name, scenario_name)
    state = _load_state(ck)

    if state is None:
        raise ToolError(f"No provisioned infra state found for {ck}.\nRun infra_provision_postgres first.")

    psycopg = _load_psycopg()
    db_name = state.db_name
    db_user = state.db_user

    superconn_str = _build_connstr(pg_host, pg_port, _DEFAULT_PG_SUPERUSER, _DEFAULT_PG_SUPERPASSWORD)
    try:
        async with await psycopg.AsyncConnection.connect(superconn_str, autocommit=True) as conn:
            # Terminate existing connections to the DB before dropping.
            await conn.execute(
                "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()",
                (db_name,),
            )
            await conn.execute(f"DROP DATABASE IF EXISTS {db_name}")
            await conn.execute(f"DROP ROLE IF EXISTS {db_user}")
    except Exception as exc:
        raise ToolError(f"PostgreSQL teardown failed for {ck}: {exc}") from exc

    state_path = _remove_state(ck)

    return InfraDestroyResult(scenario_key=ck, destroyed=True, state_path=str(state_path) if state_path else None)


@tool
async def infra_status(vm_name: str, scenario_name: str) -> InfraStatusResult:
    """Return current infrastructure state for a scenario (INFRA-03).

    Reads persisted state from ``~/.mcp-vm-blackbox/infra/`` and optionally
    verifies live connectivity to the database using psycopg.

    Args:
        vm_name: Target VM name.
        scenario_name: Scenario identifier.

    Returns:
        InfraStatusResult with provisioned status, connectivity, and endpoint.
    """
    ck = composite_key(vm_name, scenario_name)
    state = _load_state(ck)

    if state is None:
        return InfraStatusResult(scenario_key=ck, provisioned=False)

    # Attempt connectivity check.
    connected: bool | None = None
    try:
        psycopg = _load_psycopg()
        user_connstr = _build_connstr(
            state.network_config.get("pg_host", "localhost"),
            int(state.network_config.get("pg_port", "5432")),
            state.db_user,
            state.db_password,
            state.db_name,
        )
        connected = await _verify_connectivity(psycopg, user_connstr)
    except ToolError:
        # psycopg not available — report provisioned but skip connectivity check
        connected = None

    path = _state_path(ck)
    return InfraStatusResult(
        scenario_key=ck, provisioned=True, connected=connected, db_endpoint=state.db_endpoint, state_path=str(path)
    )
