from __future__ import annotations

import asyncio
import ipaddress
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

PROFILE_ID = "xgrids-k1-fw3-wifi-v1"
SERVICE_UUID = "00007f00-0000-1000-8000-00805f9b34fb"
WRITE_CHARACTERISTIC_UUID = "00007f01-0000-1000-8000-00805f9b34fb"
STATUS_CHARACTERISTIC_UUID = "00007f02-0000-1000-8000-00805f9b34fb"
FRAME_LENGTH = 99
SSID_SLOT_LENGTH = 32
PASSWORD_SLOT_LENGTH = 64
AP_FALLBACK_IPV4 = "192.168.56.1"
ProvisioningOutcome = Literal[
    "lan_address_observed",
    "status_changed",
    "no_status_change_before_timeout",
    "ble_disconnected_after_write",
]
WriteMode = Literal["auto", "with_response", "without_response"]
ResolvedWriteMode = Literal["with_response", "without_response"]


class WifiStatus(TypedDict):
    value_length: int
    mode: str | None
    ipv4: str | None
    status_code: int
    reserved: int | None
    trailer_hex: str


class StatusObservation(TypedDict):
    observed_at_utc: str
    seconds_after_write: float
    status: WifiStatus


class WifiProvisioningResult(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_mode: ResolvedWriteMode
    write_without_response_advertised: bool
    max_write_without_response_size: int
    frame_length: int
    baseline_status: WifiStatus
    observations: list[StatusObservation]
    outcome: ProvisioningOutcome


class WifiStatusReadResult(TypedDict):
    schema_version: int
    profile_id: str
    observed_at_utc: str
    adapter: str
    bleak_version: str
    device_macos_uuid: str
    device_name: str
    service_uuid: str
    status_characteristic_uuid: str
    operation: Literal["single_reviewed_wifi_status_read"]
    write_performed: Literal[False]
    status: WifiStatus


def build_wifi_provisioning_frame(ssid: str, password: str) -> bytearray:
    """Build the deterministic 99-byte frame used by LixelGO for K1 Wi-Fi setup."""
    ssid_bytes = ssid.encode("utf-8")
    password_bytes = password.encode("utf-8")

    if not ssid_bytes:
        raise ValueError("SSID must not be empty")
    if not password_bytes:
        raise ValueError("Wi-Fi password must not be empty")
    if len(ssid_bytes) > SSID_SLOT_LENGTH:
        raise ValueError("SSID must be at most 32 UTF-8 bytes")
    if len(password_bytes) > PASSWORD_SLOT_LENGTH:
        raise ValueError("Wi-Fi password must be at most 64 UTF-8 bytes")

    frame = bytearray(FRAME_LENGTH)
    frame[0] = len(ssid_bytes)
    frame[1 : 1 + len(ssid_bytes)] = ssid_bytes
    frame[33] = len(password_bytes)
    frame[34 : 34 + len(password_bytes)] = password_bytes
    frame[98] = 0
    return frame


def parse_wifi_status(value: bytes) -> WifiStatus:
    """Parse the non-secret status frame returned by the K1 read characteristic."""
    if len(value) < 51:
        raise ValueError("K1 Wi-Fi status must contain at least 51 bytes")

    mode_length = value[0]
    if mode_length > SSID_SLOT_LENGTH:
        raise ValueError("K1 Wi-Fi status mode length is invalid")
    try:
        mode = value[1 : 1 + mode_length].decode("utf-8") if mode_length else None
    except UnicodeDecodeError as exc:
        raise ValueError("K1 Wi-Fi status mode is not valid UTF-8") from exc

    address_length = value[33]
    address_start = 34
    address_end = address_start + address_length
    if address_end > len(value):
        raise ValueError("K1 Wi-Fi status address length exceeds the frame")

    ipv4: str | None = None
    if address_length:
        try:
            address = ipaddress.ip_address(value[address_start:address_end])
        except ValueError:
            address = None
        if isinstance(address, ipaddress.IPv4Address):
            ipv4 = str(address)

    return {
        "value_length": len(value),
        "mode": mode,
        "ipv4": ipv4,
        "status_code": value[50],
        "reserved": value[51] if len(value) > 51 else None,
        "trailer_hex": value[52:].hex() if len(value) > 52 else "",
    }


def _outcome(
    baseline: WifiStatus,
    observations: list[StatusObservation],
    disconnected: bool,
) -> ProvisioningOutcome:
    if observations:
        final = observations[-1]["status"]
        if final["ipv4"] not in (None, AP_FALLBACK_IPV4):
            return "lan_address_observed"
        if final != baseline:
            return "status_changed"
    if disconnected:
        return "ble_disconnected_after_write"
    return "no_status_change_before_timeout"


async def read_wifi_status_once(
    device_macos_uuid: str,
    *,
    timeout_seconds: float = 20.0,
) -> WifiStatusReadResult:
    """Read the K1's current DHCP status over BLE without writing a characteristic."""

    if timeout_seconds <= 0:
        raise ValueError("timeout_seconds must be positive")
    async with asyncio.timeout(timeout_seconds + 5.0):
        device = discovered_device(device_macos_uuid)
        if device is None:
            device = await BleakScanner.find_device_by_address(
                device_macos_uuid,
                timeout=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:
            service = client.services.get_service(SERVICE_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 status_characteristic is None:
                raise ValueError(
                    f"Reviewed K1 status characteristic not found: {STATUS_CHARACTERISTIC_UUID}"
                )
            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")
            value = bytes(await client.read_gatt_char(status_characteristic))
            return {
                "schema_version": 1,
                "profile_id": PROFILE_ID,
                "observed_at_utc": utc_now_iso(),
                "adapter": "CoreBluetooth",
                "bleak_version": version("bleak"),
                "device_macos_uuid": device_macos_uuid,
                "device_name": client.name or "",
                "service_uuid": service.uuid,
                "status_characteristic_uuid": status_characteristic.uuid,
                "operation": "single_reviewed_wifi_status_read",
                "write_performed": False,
                "status": parse_wifi_status(value),
            }


async def provision_wifi_once(
    device_macos_uuid: str,
    ssid: str,
    password: str,
    timeout_seconds: float = 45.0,
    poll_interval_seconds: float = 1.0,
    write_mode: WriteMode = "auto",
) -> WifiProvisioningResult:
    """Perform one reviewed provisioning write and poll the K1 status characteristic."""
    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_wifi_provisioning_frame(ssid, password)
    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:
                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(
                        f"Reviewed K1 write characteristic not found: {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 write 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
                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(
                            "Provisioning frame exceeds the negotiated write-without-response size"
                        )
                    resolved_write_mode = "without_response"

                baseline_value = bytes(await client.read_gatt_char(status_characteristic))
                baseline = parse_wifi_status(baseline_value)

                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:
                        value = bytes(await client.read_gatt_char(status_characteristic))
                    except BleakError:
                        if not client.is_connected:
                            disconnected = True
                            break
                        raise
                    status = parse_wifi_status(value)
                    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 status["ipv4"] not in (None, AP_FALLBACK_IPV4):
                        break
                    await asyncio.sleep(poll_interval_seconds)

                return {
                    "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_wifi_provisioning_write",
                    "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,
                    "outcome": _outcome(baseline, observations, disconnected),
                }
    finally:
        frame[:] = b"\x00" * len(frame)
