#!/usr/bin/env python3
"""
mi_ble_reader
==============

Async BLE reader for Xiaomi Mi temperature / humidity sensors.

This module keeps the original helper names (`sensor_data_query`, `test`, etc.) so
you can keep your existing code unchanged.  In addition it exposes a minimal
public API that other apps can import:

* `read_data`            - async coroutine, returns SensorReading or None
* `read_multiple`        - async coroutine, reads a list of MACs concurrently
* `read_data_sync`       - sync wrapper around `read_data` (creates its own loop)

All heavy lifting is done in the async functions; the sync wrapper is only
provided for convenience.
"""

import asyncio
import warnings, logging
warnings.filterwarnings("ignore")
logging.getLogger('apscheduler').setLevel(logging.ERROR)
from dataclasses import dataclass
from typing import List, Optional
from config import *
from bleak import BleakClient, BleakScanner

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
MI_SERVICE_UUID = "ebe0ccb0-7a0a-4b0c-8a1a-6ff2997da3a6"
TEMP_CHAR_UUID = "ebe0ccc1-7a0a-4b0c-8a1a-6ff2997da3a6"

SCAN_TIMEOUT   = 5.0   # seconds to scan for each MAC
CONNECT_TIMEOUT = 10.0
READ_TIMEOUT    = 2.0
SERVICE_DISCOVERY_TIMEOUT = 5.0 
_lock = asyncio.Lock()

# --------------------------------------------------------------------------- #
# Logging
# --------------------------------------------------------------------------- #
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# --------------------------------------------------------------------------- #
# Data model
# --------------------------------------------------------------------------- #
@dataclass
class SensorReading:
    """Single Mi sensor reading."""
    mac: str
    temperature: float
    humidity: int


# --------------------------------------------------------------------------- #
# Helpers - kept identical to your original code
# --------------------------------------------------------------------------- #
async def _discover_device(mac: str) -> Optional[str]:
    """Scan once for the given MAC and return its address if found."""
    devices = await BleakScanner.discover(timeout=SCAN_TIMEOUT)
    for d in devices:
        if d.address.lower() == mac.lower():
            log.info("Found device %s (%s)", d.name or d.address, d.address)
            return d.address
    log.warning("Device %s not found in scan", mac)
    return None


def _parse_sensor_data(raw: bytes) -> (float, int):
    """Parse the 3-byte Mi packet into temperature & humidity."""
    sign = raw[1] & (1 << 7)
    temp = ((raw[1] & 0x7F) << 8 | raw[0])
    if sign:
        temp -= 32767
    temp = temp / 100.0
    hum = raw[2]
    return temp, hum

# --------------------------------------------------------------------------- #
# Optional device discovery helper
# --------------------------------------------------------------------------- #
async def _scan_once() -> List[str]:
    """
    Scan once for nearby BLE devices and return the MAC addresses of those
    that satisfy *filter_fn*.

    Parameters
    ----------
    filter_fn : Callable[[address, name], bool] | None
        If supplied, a predicate that receives the device address and name
        and returns ``True`` when the device should be kept.
        When ``None`` all discovered addresses are returned.

    Returns
    -------
    List[str]
        A list of MAC addresses (upper-case as returned by Bleak).
    """
    log.debug("Starting one-shot BLE scan (timeout = %.1fs)", SCAN_TIMEOUT)
    devices = await BleakScanner.discover(timeout=SCAN_TIMEOUT)

    return [dev.address for dev in devices if _is_match(dev.address, dev.name)]

# this function filters MACs
def _is_match(mac: str, name: str) -> bool:
    """Return True if the MAC starts with the configured prefix and is correct model."""
    if name == "LYWSD03MMC":
        return mac.upper().startswith(MAC_PREFIX.upper())

# this helper function scans for MACs
def scan_macs():
    print("Scanning for MACs:")
    new_macs = asyncio.run(_scan_once())
    
    print(f"Found {len(new_macs)} new MACs.")
    return new_macs

# --------------------------------------------------------------------------- #
# Core async read functions (unchanged from the refactored version)
# --------------------------------------------------------------------------- #
async def read_data(mac: str) -> Optional[SensorReading]:
    address = await _discover_device(mac)
    if not address:
        return None

    try:
        async with BleakClient(address, timeout=CONNECT_TIMEOUT) as client:
            # <---  NEW BLOCK  ----------------------------------------------
            try:
                await asyncio.wait_for(client.get_services(),
                                      timeout=SERVICE_DISCOVERY_TIMEOUT)
            except asyncio.TimeoutError:
                log.warning("Service discovery timed out for %s", mac)
                return None
            
            except Exception as e:
                print(f"Discovery failed - {e}")
            # ---------------------------------------------------------------

            # If we get here, services are available.
            if MI_SERVICE_UUID.lower() not in (
                s.uuid.lower() for s in client.services
            ):
                log.error("Mi service not found on %s", mac)
                return None

            raw = await asyncio.wait_for(
                client.read_gatt_char(TEMP_CHAR_UUID),
                timeout=READ_TIMEOUT,
            )
            temp, hum = _parse_sensor_data(raw)
            log.info("Read from %s: %.1f°C, %d%%", mac, temp, hum)
            return SensorReading(mac=mac, temperature=temp, humidity=hum)

    except asyncio.TimeoutError:
        # This is the *connect* timeout (BleakClient)
        log.error("Connection to %s timed out after %.1fs", mac, CONNECT_TIMEOUT)
    except Exception as exc:
        log.exception("Error reading from %s: %s", mac, exc)
    return None


async def read_multiple(mac_list: List[str]) -> List[SensorReading]:
    """
    Read all MACs concurrently.  Skips any that fail.
    """
    tasks = [asyncio.create_task(read_data(mac)) for mac in mac_list]
    results = await asyncio.gather(*tasks, return_exceptions=False)
    readings = [r for r in results if r is not None]
    failed = [mac for mac, r in zip(mac_list, results) if r is None]
    if failed:
        log.warning("Failed to read: %s", ", ".join(failed))
    return readings


# --------------------------------------------------------------------------- #
# Sync wrapper (kept for legacy code that cannot run an event loop)
# --------------------------------------------------------------------------- #
def read_data_sync(mac: str) -> Optional[SensorReading]:
    """
    Convenience wrapper that runs ``read_data`` in a fresh event-loop.
    Use this from non-async code.
    """
    return asyncio.run(read_data(mac))


# --------------------------------------------------------------------------- #
# Original helper functions - unchanged
# --------------------------------------------------------------------------- #
def sensor_data_query(mac_list: List[str]) -> List[dict]:
    """Legacy function; returns a list of dicts like the original code."""
    result = []
    for mac in mac_list:
        try:
            this_result = asyncio.run(read_data(mac))
        except Exception as e:
            print(f"Error reading BLE data - {e}")
            this_result = None
        if this_result:
            result.append(
                {"name": mac, "temp": this_result.temperature, "humidity": this_result.humidity}
            )
    print(f"Data query complete - returning result")
    return result


def test(test_count = 0) -> None:
    """Legacy test harness - identical to your original script."""
    # first test scanner
    test_loop_count = 0
    found_macs = scan_macs()
    print("Found these MACs on a scan:")
    print(found_macs)
    test_result = False
    sensors = []
    mac_list = [
        "A4:C1:38:96:53:30",
        "A4:C1:38:11:F4:2A",
        "A4:C1:38:24:15:B0",
        "A4:C1:38:D7:8D:26",
    ]

    print("Starting BLE test...")
    print("  Scanning MACs:")
    for mac in mac_list:
        print(f"    {mac}")
    test_list = mac_list
    while not test_result:
        print(f"  Test loop {test_loop_count} start. Initializing run...")
        test_run_results = sensor_data_query(test_list)
        test_loop_count = test_loop_count + 1
        for test_run_result in test_run_results:
            sensors.append(test_run_result)
        passed_list = [device["name"] for device in sensors]
        print(f"    Passed list - {passed_list}")
        if len(mac_list) == len(passed_list):
            print(f"  Test run {test_count} completed, yee-haw")
            for sensor in sensors:
                print(f"    {sensor['name']} - Temp: {round(sensor['temp'],1)} | Humidity: {sensor['humidity']}")
            test_result = True
        else:
            print(f"  Test run incomplete/failed, re-running, loop {test_loop_count}. Updating test list")
            for sensor in sensors:
                print(f"    {sensor['name']} - Temp: {round(sensor['temp'],1)} | Humidity: {sensor['humidity']}")
            mac_list_new = []
            for mac in mac_list:
                if mac not in passed_list:
                    mac_list_new.append(mac)
            test_list = mac_list_new
            print(f"  Re-do list - {test_list}")


# --------------------------------------------------------------------------- #
# If run as a script, execute the legacy test
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    test_count = 0
    print("Welp, this will keep going forever. Have fun :)")
    while True:
        print(f"Starting test run {test_count}")
        test_count = test_count + 1
        try:
            test(test_count)
        except Exception as e:
            print(f"Test failed brah - {e}")