"""VM recording tools: start, stop, status, and frame extraction for always-on screen capture."""

from __future__ import annotations

import logging
import re
import shutil
from datetime import UTC, datetime
from pathlib import Path
from typing import Annotated

import anyio
import av
from fastmcp.tools import tool
from pydantic import Field

from mcp_vm_blackbox._helpers import _resolve, _ssh_run, default_blackbox_data_root, project_root
from mcp_vm_blackbox.models import (
    FrameExtractionResult,
    RecordingStartResult,
    RecordingStatusResult,
    RecordingStopResult,
    SummaryVideoResult,
)
from mcp_vm_blackbox.proof_demo import build_proof_demo
from mcp_vm_blackbox.vm_registry import default_registry

_log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Scratch directory for recordings and frames
# ---------------------------------------------------------------------------

_bb = default_blackbox_data_root()
_RECORDINGS_DIR = _bb / "scratch" / "recordings"
_FRAMES_DIR = _bb / "scratch" / "frames"


def _safe_vm_name(vm_name: str) -> str:
    """Sanitise a VM name for use in file paths.

    Replaces non-alphanumeric characters (except hyphens) with underscores.

    Args:
        vm_name: Raw VirtualBox VM name.

    Returns:
        Filesystem-safe version of the name.
    """
    return re.sub(r"[^a-zA-Z0-9\-]", "_", vm_name)


def _recording_path(vm_name: str) -> Path:
    """Generate a timestamped recording path for a VM.

    Args:
        vm_name: VirtualBox VM name.

    Returns:
        Path like ``scratch/recordings/<safe_name>-<YYYYMMDD-HHMMSS>-screen0.webm``.
    """
    ts = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
    safe = _safe_vm_name(vm_name)
    return _RECORDINGS_DIR / f"{safe}-{ts}-screen0.webm"


def _vboxmanage_recording_cmd(vm_name: str, *args: str) -> list[str]:
    """Build a VBoxManage controlvm recording command.

    Args:
        vm_name: VirtualBox VM name.
        *args: Additional arguments after ``controlvm <vm> recording``.

    Returns:
        Full command as a list of strings.
    """
    return ["VBoxManage", "controlvm", vm_name, "recording", *args]


# ---------------------------------------------------------------------------
# Private helpers — shared by MCP tools and vm_lifecycle hooks
# ---------------------------------------------------------------------------


async def _start_recording(
    vm_name: str, target: str = "local", fps: int = 5, video_rate_kbps: int = 256
) -> RecordingStartResult:
    """Start VBoxManage recording on a running VM.

    Issues the sequence of ``controlvm recording`` commands to configure and
    start screen capture. Updates the VM registry with recording state.

    Args:
        vm_name: VirtualBox VM name.
        target: Named target or SSH string.
        fps: Frames per second for recording.
        video_rate_kbps: Video bitrate in kbps.

    Returns:
        RecordingStartResult with status information.
    """
    cfg = _resolve(target)
    rec_path = _recording_path(vm_name)
    rec_path.parent.mkdir(parents=True, exist_ok=True)
    abs_rec_path = str(rec_path.resolve())

    # Issue VBoxManage controlvm recording commands in sequence
    commands: list[tuple[str, list[str]]] = [
        ("filename", _vboxmanage_recording_cmd(vm_name, "filename", abs_rec_path)),
        ("videofps", _vboxmanage_recording_cmd(vm_name, "videofps", str(fps))),
        ("videorate", _vboxmanage_recording_cmd(vm_name, "videorate", str(video_rate_kbps))),
        ("on", _vboxmanage_recording_cmd(vm_name, "on")),
        ("start", _vboxmanage_recording_cmd(vm_name, "start")),
    ]

    for step_name, cmd in commands:
        result = await _ssh_run(cfg, cmd)
        if result.returncode != 0:
            error_msg = f"Recording {step_name} failed: {result.stderr.strip() or result.stdout}"
            _log.error("vm_recording_start: %s", error_msg)
            return RecordingStartResult(status="error", vm_name=vm_name, error=error_msg)

    now = datetime.now(tz=UTC).isoformat()
    await default_registry.update(vm_name, recording=True, recording_path=abs_rec_path, recording_started_at=now)

    return RecordingStartResult(
        status="ok",
        vm_name=vm_name,
        recording_path=str(rec_path),
        started_at=now,
        fps=fps,
        video_rate_kbps=video_rate_kbps,
    )


async def _stop_recording(vm_name: str, target: str = "local") -> RecordingStopResult:
    """Stop VBoxManage recording on a running VM.

    Issues ``controlvm recording stop`` then ``recording off``.
    Updates the VM registry to clear recording state.

    Args:
        vm_name: VirtualBox VM name.
        target: Named target or SSH string.

    Returns:
        RecordingStopResult with status information.
    """
    cfg = _resolve(target)

    # Get the current recording path from the registry before stopping
    entry = await default_registry.get(vm_name)
    rec_path = entry.recording_path if entry else None

    commands: list[tuple[str, list[str]]] = [
        ("stop", _vboxmanage_recording_cmd(vm_name, "stop")),
        ("off", _vboxmanage_recording_cmd(vm_name, "off")),
    ]

    for step_name, cmd in commands:
        result = await _ssh_run(cfg, cmd)
        if result.returncode != 0:
            error_msg = f"Recording {step_name} failed: {result.stderr.strip() or result.stdout}"
            _log.error("vm_recording_stop: %s", error_msg)
            return RecordingStopResult(status="error", vm_name=vm_name, recording_path=rec_path, error=error_msg)

    # Calculate duration if we have a start time
    duration_s: float | None = None
    if entry and entry.recording_started_at:
        try:
            started = datetime.fromisoformat(entry.recording_started_at)
            duration_s = (datetime.now(tz=UTC) - started).total_seconds()
        except ValueError:
            pass

    # Check file size on disk
    file_size: int | None = None
    if rec_path:
        ap = anyio.Path(rec_path)
        if await ap.exists():
            stat = await ap.stat()
            file_size = stat.st_size

    # Update registry: recording off, preserve path
    await default_registry.update(vm_name, recording=False, recording_started_at=None)

    return RecordingStopResult(
        status="ok", vm_name=vm_name, recording_path=rec_path, duration_s=duration_s, file_size_bytes=file_size
    )


# ---------------------------------------------------------------------------
# MCP Tools
# ---------------------------------------------------------------------------


@tool
async def vm_recording_start(
    vm_name: Annotated[str, Field(description="VirtualBox VM name.")],
    fps: Annotated[int, Field(description="Frames per second (default 5).")] = 5,
    video_rate_kbps: Annotated[int, Field(description="Video bitrate in kbps (default 256).")] = 256,
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> RecordingStartResult:
    """Start continuous screen recording on a running VM.

    Configures and starts VBoxManage ``controlvm recording`` to write
    WebM/VP8 video continuously to disk. The recording runs until
    explicitly stopped or the VM is shut down.

    Args:
        vm_name: VirtualBox VM name as shown in ``VBoxManage list vms``.
        fps: Frames per second for the recording (default 5).
        video_rate_kbps: Video bitrate in kbps (default 256).
        target: Named target or raw SSH string.

    Returns:
        RecordingStartResult with status, recording path, and timestamps.
    """
    return await _start_recording(vm_name, target=target, fps=fps, video_rate_kbps=video_rate_kbps)


@tool
async def vm_recording_stop(
    vm_name: Annotated[str, Field(description="VirtualBox VM name.")],
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> RecordingStopResult:
    """Stop continuous screen recording on a running VM.

    Stops the VBoxManage recording and finalises the WebM file.
    The recording path is preserved in the registry for later
    frame extraction.

    Args:
        vm_name: VirtualBox VM name as shown in ``VBoxManage list vms``.
        target: Named target or raw SSH string.

    Returns:
        RecordingStopResult with status, file path, duration, and size.
    """
    return await _stop_recording(vm_name, target=target)


@tool
async def vm_recording_status(
    vm_name: Annotated[str, Field(description="VirtualBox VM name.")],
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> RecordingStatusResult:
    """Check recording status for a VM.

    Reads the VM registry for persisted recording state and checks
    the recording file on disk for current size.

    Args:
        vm_name: VirtualBox VM name as shown in ``VBoxManage list vms``.
        target: Named target or raw SSH string.

    Returns:
        RecordingStatusResult with current recording state.
    """
    entry = await default_registry.get(vm_name)
    if entry is None:
        return RecordingStatusResult(vm_name=vm_name, recording=False)

    # Calculate elapsed time if recording is active
    elapsed_s: float | None = None
    if entry.recording and entry.recording_started_at:
        try:
            started = datetime.fromisoformat(entry.recording_started_at)
            elapsed_s = (datetime.now(tz=UTC) - started).total_seconds()
        except ValueError:
            pass

    # Check file size on disk
    file_size: int | None = None
    if entry.recording_path:
        ap = anyio.Path(entry.recording_path)
        if await ap.exists():
            stat = await ap.stat()
            file_size = stat.st_size

    return RecordingStatusResult(
        vm_name=vm_name,
        recording=entry.recording,
        recording_path=entry.recording_path,
        recording_started_at=entry.recording_started_at,
        elapsed_s=elapsed_s,
        file_size_bytes=file_size,
    )


def _parse_time_offset(ts: str) -> float:
    """Parse a time offset string to seconds.

    Accepts ``HH:MM:SS``, ``MM:SS``, or plain seconds.

    Args:
        ts: Time string in HH:MM:SS, MM:SS, or seconds format.

    Returns:
        Time offset in seconds as a float.

    Raises:
        ValueError: When the format is not recognised.
    """
    parts = ts.strip().split(":")
    if len(parts) == 3:  # noqa: PLR2004
        return int(parts[0]) * 3600 + int(parts[1]) * 60 + float(parts[2])
    if len(parts) == 2:  # noqa: PLR2004
        return int(parts[0]) * 60 + float(parts[1])
    return float(ts)


async def _resolve_source_path(recording_path: str, vm_name: str) -> str | None:
    """Resolve the recording source path from argument or registry.

    Args:
        recording_path: Explicit path, or empty string to look up from registry.
        vm_name: VirtualBox VM name for registry lookup.

    Returns:
        Resolved path string, or None if not found.
    """
    if recording_path:
        return recording_path
    entry = await default_registry.get(vm_name)
    if entry and entry.recording_path:
        return entry.recording_path
    return None


def _decode_frames(
    container: av.container.InputContainer, start_s: float, end_s: float, interval_s: float, output_dir: Path
) -> tuple[list[str], str | None]:
    """Decode frames from an open container within a time range.

    Args:
        container: Open PyAV input container.
        start_s: Start offset in seconds.
        end_s: End offset in seconds.
        interval_s: Minimum interval between extracted frames in seconds.
        output_dir: Directory to write PNG frames into.

    Returns:
        Tuple of (frame paths, error message or None).
    """
    frames: list[str] = []
    start_us = int(start_s * 1_000_000)
    container.seek(start_us)

    last_extracted = -interval_s  # ensure first valid frame is captured
    count = 0

    try:
        for frame in container.decode(video=0):
            if frame.time is None:
                continue
            if frame.time < start_s:
                continue
            if frame.time > end_s:
                break
            if (frame.time - last_extracted) >= interval_s:
                img = frame.to_image()
                frame_path = output_dir / f"frame_{count:04d}.png"
                img.save(str(frame_path))
                frames.append(str(frame_path))
                last_extracted = frame.time
                count += 1
    except av.error.FFmpegError as exc:
        return frames, f"Decode error (partial results): {exc}"

    return frames, None


@tool
async def vm_extract_frames(
    vm_name: Annotated[str, Field(description="VirtualBox VM name.")],
    timestamp_start: Annotated[str, Field(description="Start timestamp (HH:MM:SS offset).")],
    timestamp_end: Annotated[str, Field(description="End timestamp (HH:MM:SS offset).")],
    interval_ms: Annotated[int, Field(description="Interval between frames in milliseconds (default 1000).")] = 1000,
    recording_path: Annotated[
        str, Field(description="Path to WebM file. If empty, uses current/last recording from registry.")
    ] = "",
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string ("user@host").')
    ] = "local",
) -> FrameExtractionResult:
    """Extract frames from a WebM recording at regular intervals.

    Uses PyAV to decode the WebM/VP8 recording and save frames as PNG
    images. Frames are written to ``scratch/frames/<vm_name>/<timestamp>/``.

    Args:
        vm_name: VirtualBox VM name.
        timestamp_start: Start time as HH:MM:SS offset into the recording.
        timestamp_end: End time as HH:MM:SS offset into the recording.
        interval_ms: Milliseconds between extracted frames (default 1000).
        recording_path: Path to the WebM file. If empty, resolves from
            the VM registry's current or last recording path.
        target: Named target or raw SSH string.

    Returns:
        FrameExtractionResult with paths to extracted PNG frames.
    """
    # Resolve recording path
    source_path = await _resolve_source_path(recording_path, vm_name)
    if source_path is None:
        return FrameExtractionResult(
            vm_name=vm_name,
            error="No recording path found. Provide recording_path or ensure VM has a recording in the registry.",
        )

    # Parse time offsets
    try:
        start_s = _parse_time_offset(timestamp_start)
        end_s = _parse_time_offset(timestamp_end)
    except ValueError as exc:
        return FrameExtractionResult(
            vm_name=vm_name, source_recording=source_path, error=f"Invalid timestamp format: {exc}"
        )

    interval_s = interval_ms / 1000.0

    # Prepare output directory
    ts_label = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
    output_dir = _FRAMES_DIR / _safe_vm_name(vm_name) / ts_label
    await anyio.Path(output_dir).mkdir(parents=True, exist_ok=True)

    # Open container and extract frames
    try:
        container = av.open(source_path)
    except FileNotFoundError:
        return FrameExtractionResult(
            vm_name=vm_name, source_recording=source_path, error=f"Recording file not found: {source_path}"
        )
    except av.error.InvalidDataError:
        return FrameExtractionResult(
            vm_name=vm_name, source_recording=source_path, error=f"Invalid media container: {source_path}"
        )

    try:
        frames, error = _decode_frames(container, start_s, end_s, interval_s, output_dir)
    finally:
        container.close()

    return FrameExtractionResult(
        vm_name=vm_name,
        frames=frames,
        frame_count=len(frames),
        source_recording=source_path,
        time_range=f"{timestamp_start}--{timestamp_end}",
        error=error,
    )


@tool
async def vm_build_summary_video(
    job_id: Annotated[str, Field(description="Job identifier.")],
    vm_name: Annotated[str, Field(description="VM name associated with the job.")],
    evidence_root: Annotated[str, Field(description="Evidence root directory.")] = ".mcp-vm-blackbox/evidence",
) -> SummaryVideoResult:
    """Build a summary video from a completed job's evidence and continuous recording.

    Wraps ``build_proof_demo()`` to produce a highlight reel from the job's
    step log and attached evidence. Before calling validation, copies or
    symlinks the continuous recording into
    ``evidence/{vm_name}/{job_id}/recordings/`` so the validation gate finds it.

    On validation failure (``ValueError`` from ``build_proof_demo``), returns a
    ``SummaryVideoResult`` with ``error`` set rather than propagating the
    exception. All other exceptions propagate to the FastMCP
    ``ErrorHandlingMiddleware``.

    Args:
        job_id: Unique job identifier for the completed job.
        vm_name: VirtualBox VM name associated with the job.
        evidence_root: Path to the evidence root directory. Relative paths are
            resolved against the project root.

    Returns:
        SummaryVideoResult with video_path, duration_s, step_count, and
        evidence_count on success, or error set on validation failure.
    """
    # Resolve evidence_root to an absolute Path.
    evidence_root_path = Path(evidence_root)
    if not evidence_root_path.is_absolute():
        evidence_root_path = project_root() / evidence_root_path

    # If the VM registry has a recording_path, copy it into the recordings
    # directory so build_proof_demo validation can find it.
    entry = await default_registry.get(vm_name)
    if entry and entry.recording_path:
        recordings_dir = evidence_root_path / vm_name / job_id / "recordings"
        await anyio.Path(recordings_dir).mkdir(parents=True, exist_ok=True)
        src = Path(entry.recording_path)
        src_resolved = await anyio.Path(src).resolve()
        dest = recordings_dir / src.name
        if await anyio.Path(src).exists() and not await anyio.Path(dest).exists():
            try:
                await anyio.Path(dest).symlink_to(src_resolved)
            except OSError as exc:
                _log.warning("vm_build_summary_video: symlink failed, falling back to copy: %s", exc)
                shutil.copy2(src, dest)

    # Delegate to build_proof_demo; catch only ValueError (validation failure).
    try:
        manifest = await build_proof_demo(
            job_id=job_id, vm_name=vm_name, evidence_root=evidence_root_path, save_path=None, video_enabled=True
        )
    except ValueError as exc:
        _log.warning("vm_build_summary_video: validation failure for job %s: %s", job_id, exc)
        return SummaryVideoResult(job_id=job_id, vm_name=vm_name, error=str(exc))

    # Derive counts from the manifest evidence list.
    # source_evidence_count is pre-computed by build_manifest() and covers
    # screenshots + recordings that are actually present on disk.
    evidence_count = manifest.source_evidence_count + manifest.overlay_evidence_count
    # step_count: count only screenshot evidence that is present on disk.
    # job_record items are always present=True but do not represent captured
    # moments; including them would overcount steps.
    step_count = sum(1 for item in manifest.evidence if item.source_kind == "screenshot" and item.present)

    return SummaryVideoResult(
        job_id=job_id,
        vm_name=vm_name,
        video_path=manifest.video_output_path,
        step_count=step_count,
        evidence_count=evidence_count,
        # duration_s is always None: ReplayBundle has no duration field.
        # It will remain None until the Remotion renderer exposes elapsed time.
    )
