from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from importlib.metadata import version
from time import monotonic
from typing import Literal, TypedDict

from bleak import BleakClient, BleakScanner
from bleak.exc import BleakDeviceNotFoundError, BleakError

from k1link.artifacts import utc_now_iso
from k1link.device_plugins.xgrids_k1.ble.scanner import discovered_device
from k1link.device_plugins.xgrids_k1.ble.wifi_provisioning import (
    AP_FALLBACK_IPV4,
    SERVICE_UUID,
    STATUS_CHARACTERISTIC_UUID,
    WRITE_CHARACTERISTIC_UUID,
    ResolvedWriteMode,
    StatusObservation,
    WifiStatus,
    WriteMode,
    parse_wifi_status,
)

PROFILE_ID = "xgrids-k1-fw3-quick-connect-ap-v1"
FRAME_LENGTH = 100
COMMAND_OFFSET = 99
ENABLE_AP_COMMAND = 1

ApActivationOutcome = Literal[
    "already_active",
    "ap_ready_observed",
    "status_changed",
    "ble_disconnected_after_write",
    "no_status_change_before_timeout",
]


class ApActivationResult(TypedDict):
    schema_version: int
    profile_id: str
    started_at_utc: str
    completed_at_utc: str
    adapter: str
    bleak_version: str
    device_macos_uuid: str
    device_name: str
    service_uuid: str
    write_characteristic_uuid: str
    status_characteristic_uuid: str
    operation: str
    write_performed: bool
    write_mode: ResolvedWriteMode | None
    write_without_response_advertised: bool
    max_write_without_response_size: int
    frame_length: int
    baseline_status: WifiStatus
    observations: list[StatusObservation]
    ready_observed: bool
    outcome: ApActivationOutcome


def build_ap_activation_frame() -> bytearray:
    """Build LixelGO's fixed 100-byte Quick Connect AP-enable frame."""

    frame = bytearray(FRAME_LENGTH)
    frame[COMMAND_OFFSET] = ENABLE_AP_COMMAND
    return frame


def is_ap_ready_status(status: WifiStatus) -> bool:
    # The reviewed LixelGO build maps byte 51 of the 7f02 response to its
    # Wi-Fi-AP-ready flag. WIFI_AP plus the fallback address describes the
    # selected control mode, but can remain stale after the beacon disappears.
    return (
        status["mode"] == "WIFI_AP"
        and status["ipv4"] == AP_FALLBACK_IPV4
        and status["reserved"] not in (None, 0)
    )


def _outcome(
    baseline: WifiStatus,
    observations: list[StatusObservation],
    disconnected: bool,
) -> ApActivationOutcome:
    if observations:
        final = observations[-1]["status"]
        if is_ap_ready_status(final):
            return "ap_ready_observed"
        if final != baseline:
            return "status_changed"
    if disconnected:
        return "ble_disconnected_after_write"
    return "no_status_change_before_timeout"


@asynccontextmanager
async def device_ap_activation_session(
    device_macos_uuid: str,
    timeout_seconds: float = 15.0,
    poll_interval_seconds: float = 0.5,
    write_mode: WriteMode = "auto",
) -> AsyncIterator[ApActivationResult]:
    """Keep BLE connected around one reviewed Quick Connect AP-enable write.

    The payload is the exact fixed frame used by the reviewed LixelGO build.
    ``auto`` follows the live GATT properties because the owner-controlled K1
    advertises a write with response even though the Android client requests a
    write without response. The caller receives the result while the same BLE
    session is still alive, matching LixelGO's AP-ready -> native Wi-Fi handoff.
    No transport or command retry is attempted.
    """

    if timeout_seconds <= 0:
        raise ValueError("timeout_seconds must be positive")
    if poll_interval_seconds <= 0:
        raise ValueError("poll_interval_seconds must be positive")
    if write_mode not in ("auto", "with_response", "without_response"):
        raise ValueError(f"Unsupported write mode: {write_mode}")

    frame = build_ap_activation_frame()
    started_at = utc_now_iso()
    observations: list[StatusObservation] = []
    disconnected = False

    try:
        async with asyncio.timeout(timeout_seconds + 25.0):
            device = discovered_device(device_macos_uuid)
            if device is None:
                device = await BleakScanner.find_device_by_address(
                    device_macos_uuid,
                    timeout=min(20.0, timeout_seconds),
                )
            if device is None:
                raise BleakDeviceNotFoundError(
                    device_macos_uuid,
                    "Device was not rediscovered; keep the K1 powered and nearby.",
                )

        async with BleakClient(device, timeout=timeout_seconds, pair=False) as client:
            async with asyncio.timeout(timeout_seconds + 10.0):
                device_name = client.name
                service = client.services.get_service(SERVICE_UUID)
                write_characteristic = client.services.get_characteristic(
                    WRITE_CHARACTERISTIC_UUID
                )
                status_characteristic = client.services.get_characteristic(
                    STATUS_CHARACTERISTIC_UUID
                )
                if service is None:
                    raise ValueError(f"Reviewed K1 service not found: {SERVICE_UUID}")
                if write_characteristic is None:
                    raise ValueError(
                        "Reviewed K1 AP-control characteristic not found: "
                        f"{WRITE_CHARACTERISTIC_UUID}"
                    )
                if status_characteristic is None:
                    raise ValueError(
                        f"Reviewed K1 status characteristic not found: {STATUS_CHARACTERISTIC_UUID}"
                    )
                if write_characteristic.service_uuid != service.uuid:
                    raise ValueError(
                        "K1 AP-control characteristic is attached to an unexpected service"
                    )
                if status_characteristic.service_uuid != service.uuid:
                    raise ValueError(
                        "K1 status characteristic is attached to an unexpected service"
                    )
                if "read" not in status_characteristic.properties:
                    raise ValueError("Reviewed K1 status characteristic is not readable")

                properties = set(write_characteristic.properties)
                max_without_response = write_characteristic.max_write_without_response_size
                baseline = parse_wifi_status(
                    bytes(await client.read_gatt_char(status_characteristic))
                )
                # WIFI_AP is a control-mode status, not proof that the radio is
                # still beaconing. A physical run found the exact SSID shortly
                # after AP-enable, then found no beacon while 7f02 continued to
                # report WIFI_AP. LixelGO emits the reviewed enable frame for
                # each explicit Quick Connect action, so Mission Core does the
                # same once per operator action instead of short-circuiting on
                # a stale-ready status. There is still no automatic retry.

                resolved_write_mode: ResolvedWriteMode
                if write_mode == "auto":
                    if "write-without-response" in properties:
                        resolved_write_mode = "without_response"
                    elif "write" in properties:
                        resolved_write_mode = "with_response"
                    else:
                        raise ValueError("Reviewed K1 characteristic is not writable")
                elif write_mode == "with_response":
                    if "write" not in properties:
                        raise ValueError(
                            "Reviewed K1 characteristic does not advertise writes with response"
                        )
                    resolved_write_mode = "with_response"
                else:
                    if len(frame) > max_without_response:
                        raise ValueError(
                            "AP activation frame exceeds the negotiated "
                            "write-without-response size"
                        )
                    resolved_write_mode = "without_response"

                await client.write_gatt_char(
                    write_characteristic,
                    frame,
                    response=resolved_write_mode == "with_response",
                )
                write_completed = monotonic()
                deadline = write_completed + timeout_seconds

                while monotonic() < deadline:
                    try:
                        status = parse_wifi_status(
                            bytes(await client.read_gatt_char(status_characteristic))
                        )
                    except BleakError:
                        if not client.is_connected:
                            disconnected = True
                            break
                        raise
                    observation: StatusObservation = {
                        "observed_at_utc": utc_now_iso(),
                        "seconds_after_write": round(monotonic() - write_completed, 3),
                        "status": status,
                    }
                    if not observations or status != observations[-1]["status"]:
                        observations.append(observation)
                    if is_ap_ready_status(status):
                        break
                    await asyncio.sleep(poll_interval_seconds)

                result: ApActivationResult = {
                    "schema_version": 1,
                    "profile_id": PROFILE_ID,
                    "started_at_utc": started_at,
                    "completed_at_utc": utc_now_iso(),
                    "adapter": "CoreBluetooth",
                    "bleak_version": version("bleak"),
                    "device_macos_uuid": device_macos_uuid,
                    "device_name": device_name,
                    "service_uuid": service.uuid,
                    "write_characteristic_uuid": write_characteristic.uuid,
                    "status_characteristic_uuid": status_characteristic.uuid,
                    "operation": "single_reviewed_quick_connect_ap_activation",
                    "write_performed": True,
                    "write_mode": resolved_write_mode,
                    "write_without_response_advertised": (
                        "write-without-response" in properties
                    ),
                    "max_write_without_response_size": max_without_response,
                    "frame_length": len(frame),
                    "baseline_status": baseline,
                    "observations": observations,
                    "ready_observed": bool(
                        observations and is_ap_ready_status(observations[-1]["status"])
                    ),
                    "outcome": _outcome(baseline, observations, disconnected),
                }
            # Keep the same CoreBluetooth session alive while the caller waits
            # for and performs the host-side CoreWLAN association. LixelGO does
            # not tear down this BLE manager between its AP-ready callback and
            # native Wi-Fi connect call.
            yield result
    finally:
        frame[:] = b"\x00" * len(frame)


async def activate_device_ap_once(
    device_macos_uuid: str,
    timeout_seconds: float = 15.0,
    poll_interval_seconds: float = 0.5,
    write_mode: WriteMode = "auto",
) -> ApActivationResult:
    """Run one AP activation and release BLE immediately after its result.

    Host association flows must use :func:`device_ap_activation_session` so
    the reviewed LixelGO BLE-to-Wi-Fi handoff remains one connected session.
    """

    async with device_ap_activation_session(
        device_macos_uuid,
        timeout_seconds=timeout_seconds,
        poll_interval_seconds=poll_interval_seconds,
        write_mode=write_mode,
    ) as result:
        return result
