"""MCP tools for goal-state interview, retrieval, and listing.

Implements vm_goal_interview, vm_get_goal_state, and vm_list_goal_states.
These are plain async functions discovered by FileSystemProvider in server.py.
"""

from __future__ import annotations

import uuid
from dataclasses import dataclass
from typing import cast

from fastmcp.server.context import AcceptedElicitation, Context

from mcp_vm_blackbox.goal_state_store import composite_key, default_store, goal_state_key
from mcp_vm_blackbox.models import Criterion, CriterionKind, GoalState, HostProtocol, HostRosterEntry

# ---------------------------------------------------------------------------
# Module-level store singleton
# ---------------------------------------------------------------------------

_store = default_store

# ---------------------------------------------------------------------------
# State key prefixes (module-level constants)
# ---------------------------------------------------------------------------

_draft_key_prefix = "goal_draft"

# ---------------------------------------------------------------------------
# Elicitation flat dataclasses
# ---------------------------------------------------------------------------


@dataclass
class ScenarioStep:
    """Step 1: scenario name and fail_fast flag."""

    scenario_name: str
    fail_fast: bool


@dataclass
class HostEntry:
    """Step 2: a single host in the roster."""

    host_name: str
    host_ip: str
    protocol: HostProtocol
    port: int
    username: str


@dataclass
class CriterionEntry:
    """Step 3: a single success criterion for a host."""

    host_name: str
    criterion_type: CriterionKind
    param_key: str
    param_value: str


@dataclass
class TicketRefEntry:
    """Step 4: a ticket reference."""

    ticket_ref: str


@dataclass
class OverwriteConfirm:
    """Confirmation step when a committed state already exists."""

    confirm: str  # "yes" to proceed, anything else treated as no


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _has_mandatory_criteria(criteria: list[Criterion]) -> bool:
    """Return True iff at least one port_open AND one service_running criterion."""
    types = {c.criterion_type for c in criteria}
    return "port_open" in types and "service_running" in types


async def _collect_hosts(
    ctx: Context, composite_key: str, draft: GoalState
) -> tuple[list[HostRosterEntry], str | None]:
    """Elicit hosts until the user declines.

    Returns (hosts, abort_message). abort_message is non-None when the user
    declined before entering at least one host.
    """
    hosts: list[HostRosterEntry] = []
    while True:
        result = await ctx.elicit(
            f"Add a host to the roster (host {len(hosts) + 1}). Decline when done.", response_type=HostEntry
        )
        if not isinstance(result, AcceptedElicitation):
            if not hosts:
                _store.save_draft(draft)
                return ([], f"Interview aborted (no hosts entered). Draft saved for {composite_key}.")
            break
        data = cast("HostEntry", result.data)
        hosts.append(
            HostRosterEntry(
                name=data.host_name, ip=data.host_ip, protocol=data.protocol, port=data.port, username=data.username
            )
        )
    return hosts, None


async def _collect_criteria(ctx: Context, hosts: list[HostRosterEntry]) -> list[Criterion]:
    """Elicit criteria for each host until the user declines per host."""
    criteria: list[Criterion] = []
    for host in hosts:
        while True:
            result = await ctx.elicit(
                f"Add a criterion for host '{host.name}'. Decline when done with this host.",
                response_type=CriterionEntry,
            )
            if not isinstance(result, AcceptedElicitation):
                break
            data = cast("CriterionEntry", result.data)
            criteria.append(
                Criterion(
                    criterion_id=str(uuid.uuid4()),
                    host_name=data.host_name,
                    criterion_type=data.criterion_type,
                    params={data.param_key: data.param_value},
                )
            )
    return criteria


async def _collect_ticket_refs(ctx: Context) -> list[str]:
    """Elicit ticket references until the user declines (optional)."""
    ticket_refs: list[str] = []
    while True:
        result = await ctx.elicit(
            "Add a ticket reference (e.g. JIRA-1234). Decline when done.", response_type=TicketRefEntry
        )
        if not isinstance(result, AcceptedElicitation):
            break
        data = cast("TicketRefEntry", result.data)
        ticket_refs.append(data.ticket_ref)
    return ticket_refs


# ---------------------------------------------------------------------------
# MCP tools (plain async functions — FileSystemProvider registers them)
# ---------------------------------------------------------------------------


async def vm_goal_interview(vm_name: str, ctx: Context) -> str:
    """Run a linear 4-step elicitation interview to capture goal state for a VM scenario.

    Steps:
      1. ScenarioStep  — scenario name + fail_fast
      2. HostEntry loop — collect hosts until decline
      3. CriterionEntry loop per host — collect criteria until decline
      4. TicketRefEntry loop — collect ticket refs until decline (optional)

    On completion: validates mandatory criteria, promotes draft to committed,
    sets session state, deletes draft key.

    On abort at any step: saves draft, returns abort message.
    """
    # Step 1: Scenario
    step1_result = await ctx.elicit("Enter scenario name and fail_fast setting.", response_type=ScenarioStep)
    if not isinstance(step1_result, AcceptedElicitation):
        draft = GoalState(vm_name=vm_name, scenario_name="__aborted__", status="draft")
        _store.save_draft(draft)
        return f"Interview aborted. Draft saved for {vm_name}/__aborted__."

    scenario_name = cast("ScenarioStep", step1_result.data).scenario_name
    fail_fast = cast("ScenarioStep", step1_result.data).fail_fast
    ck = composite_key(vm_name, scenario_name)

    goal_key = goal_state_key(ck)
    draft_key = f"{_draft_key_prefix}:{ck}"

    # Check for existing committed state — confirm overwrite
    if _store.load(vm_name, scenario_name) is not None:
        overwrite_result = await ctx.elicit(
            f"Goal state for {ck} already committed. Overwrite? (decline to cancel)", response_type=OverwriteConfirm
        )
        if not isinstance(overwrite_result, AcceptedElicitation):
            return f"Overwrite cancelled. Existing goal state for {ck} preserved."
        ow_data = cast("OverwriteConfirm", overwrite_result.data)
        if ow_data.confirm.strip().lower() != "yes":
            return f"Overwrite declined. Existing goal state for {ck} preserved."

    # Check for draft to resume
    if _store.load_draft(vm_name, scenario_name) is not None:
        await ctx.elicit(f"Resuming draft for {ck}... (accept to continue)", response_type=ScenarioStep)

    # Build working draft
    draft = GoalState(vm_name=vm_name, scenario_name=scenario_name, status="draft", fail_fast=fail_fast)
    _store.save_draft(draft)

    # Step 2: Host roster
    hosts, abort_msg = await _collect_hosts(ctx, ck, draft)
    if abort_msg is not None:
        return abort_msg

    draft.hosts = hosts
    _store.save_draft(draft)

    # Step 3: Criteria
    criteria = await _collect_criteria(ctx, hosts)
    draft.criteria = criteria
    _store.save_draft(draft)

    # Step 4: Ticket refs
    ticket_refs = await _collect_ticket_refs(ctx)
    draft.ticket_refs = ticket_refs

    # Mandatory criteria check
    if not _has_mandatory_criteria(criteria):
        _store.save_draft(draft)
        return (
            f"Interview incomplete. Goal state for {composite_key} requires at least "
            "one port check AND one service check required. Draft saved."
        )

    # Promote to committed
    _store.save_draft(draft)
    await ctx.set_state(goal_key, _store.promote_draft(vm_name, scenario_name).model_dump())
    await ctx.delete_state(draft_key)

    return (
        f"Goal state committed for {composite_key}. "
        f"{len(hosts)} host(s), {len(criteria)} criterion/criteria, "
        f"{len(ticket_refs)} ticket ref(s)."
    )


async def vm_get_goal_state(vm_name: str, scenario_name: str, ctx: Context) -> str:
    """Retrieve the committed goal state for a VM/scenario combination.

    Checks session state first (ctx.get_state), then falls back to file store.
    Does not re-interview.
    """
    ck = composite_key(vm_name, scenario_name)
    goal_key = goal_state_key(ck)

    session_data = await ctx.get_state(goal_key)
    if session_data is not None:
        goal = GoalState.model_validate(session_data)
        return (
            f"Goal state for {ck} (from session):\n"
            f"  fail_fast={goal.fail_fast}\n"
            f"  hosts={[h.name for h in goal.hosts]}\n"
            f"  criteria={len(goal.criteria)}\n"
            f"  tickets={goal.ticket_refs}"
        )

    goal = _store.load(vm_name, scenario_name)
    if goal is None:
        return f"No goal state found for {ck}."

    return (
        f"Goal state for {ck} (from file):\n"
        f"  fail_fast={goal.fail_fast}\n"
        f"  hosts={[h.name for h in goal.hosts]}\n"
        f"  criteria={len(goal.criteria)}\n"
        f"  tickets={goal.ticket_refs}"
    )


def vm_list_goal_states() -> str:
    """List all committed goal states from the file store, plus any drafts.

    Returns a formatted string with status for each state found.
    """
    committed = _store.load_all()
    lines: list[str] = []

    if committed:
        lines.append("Committed goal states:")
        lines.extend(
            f"  - {composite_key(g.vm_name, g.scenario_name)} [{len(g.hosts)} hosts, {len(g.criteria)} criteria]"
            for g in committed
        )
    else:
        lines.append("No committed goal states found.")

    draft_pairs = _store.list_drafts()
    if draft_pairs:
        lines.append("Draft goal states:")
        lines.extend(f"  - {vm}/{scenario} [draft]" for vm, scenario in draft_pairs)

    return "\n".join(lines)
