"""
S2 — Bluetooth module (connected device).

Spec table (640×480): hero @ y60 (icon 56, divider @148), action rows
52px @ 162/222/282/342. Focus default on first row. A enters, B → Home.

On device this screen is backed by a live bluetoothctl status call. Design
preview keeps a mock device only when AAHEAD_LIVE_BACKEND is disabled.
"""

from __future__ import annotations

import pygame

from ..design import tokens as t
from ..shared import widgets as w
from ..shared import fonts
from ..shared import backend


ROW_X, ROW_W, ROW_H = 16, 608, 52
ROW_Y = (154, 210, 266, 322, 378)


class BluetoothScreen:
    name = "bluetooth"

    def __init__(self, app, device=None):
        self.app = app
        self.paired_devices = []
        self.adapter_present = True
        self.audio_system_installed = True
        self.audio_setup_needed = False
        self.audio_setup_label = "Set up Bluetooth audio"
        if device:
            self.device = device
        elif backend.live_enabled():
            live, paired = self._live_state()
            self.paired_devices = paired
            self.device = live or {
                "name": "Bluetooth",
                "icon": "ph-bluetooth",
                "connected": False,
                "battery": "",
            }
        else:
            self.device = {
                "name": "AirPods Pro",
                "icon": "ph-bluetooth",
                "connected": True,
                "battery": "50%",
            }
        # rows: (label, icon, action, danger)
        self.rows = self._rows_for_device()
        self.focus = 0
        self.job_id = None
        self.job_action = ""
        self.poll_ms = 0.0

    def on_resume(self):
        live = self._live_device()
        if live:
            self.device = live
        self.rows = self._rows_for_device()
        self._clamp_focus()
        backend.gui_log("bluetooth.resumed", paired=len(self.paired_devices), device=self.device.get("name"))

    def _clamp_focus(self):
        visible_count = min(len(self.rows), len(ROW_Y))
        self.focus = max(0, min(self.focus, max(0, visible_count - 1)))

    def _rows_for_device(self):
        if self.device.get("no_adapter"):
            return [
                ("Install BT + Wi-Fi drivers", "ph-download-simple", "drivers", False),
                ("How to add Bluetooth", "ph-bluetooth", "no_adapter", False),
            ]
        needs_setup = self.adapter_present and self.audio_setup_needed
        setup_row = [(self.audio_setup_label, "ph-wrench", "setup", False)] if needs_setup else []
        if self.paired_devices and not self.device.get("connected"):
            rows = setup_row + [
                ("Connect audio", "ph-bluetooth", "connect", False),
                ("Paired devices", "ph-headphones", "paired", False),
                ("Pair new device", "ph-plus-circle", "pair", False),
                ("Audio settings", "ph-sliders-horizontal", "audio", False),
            ]
            return rows[:len(ROW_Y)]
        if not self.device.get("mac") and not self.device.get("connected"):
            return (setup_row + [
                ("Pair new device", "ph-plus-circle", "pair", False),
                ("Audio settings", "ph-sliders-horizontal", "audio", False),
                ("Install BT + Wi-Fi drivers", "ph-download-simple", "drivers", False),
            ])[:len(ROW_Y)]
        if self.device.get("connected") and not self.device.get("audio_ready", self.device.get("connected")):
            rows = setup_row + [
                ("Connect audio", "ph-bluetooth", "connect", False),
                ("Audio settings", "ph-sliders-horizontal", "audio", False),
                ("Disconnect", "ph-bluetooth-x", "disconnect", True),
                ("Forget device", "ph-trash", "forget", False),
                ("Pair another device", "ph-plus-circle", "pair", False),
            ]
            return rows[:len(ROW_Y)]
        first = ("Disconnect", "ph-bluetooth-x", "disconnect", True) if self.device.get("connected") else ("Connect audio", "ph-bluetooth", "connect", False)
        rows = setup_row + [
            first,
            ("Audio settings", "ph-sliders-horizontal", "audio", False),
            ("Forget device", "ph-trash", "forget", False),
            ("Pair another device", "ph-plus-circle", "pair", False),
        ]
        if self.paired_devices:
            insert_at = 2 if setup_row else 1
            rows.insert(insert_at, ("Paired devices", "ph-headphones", "paired", False))
        return rows[:len(ROW_Y)]

    def _live_state(self):
        env = backend.bluetooth_status()
        if not env:
            return None, []
        data = env.get("data") or {}
        self.adapter_present = bool(data.get("adapter_present"))
        self.audio_system_installed = bool(data.get("audio_system_installed"))
        self.audio_setup_needed = bool(
            data.get("audio_setup_needed", self.adapter_present and not self.audio_system_installed)
        )
        self.audio_setup_label = str(data.get("audio_setup_label") or "Set up Bluetooth audio")
        audio_route_ready = bool(data.get("audio_route_ready"))
        connected_macs = {str(d.get("mac") or "").upper() for d in data.get("connected") or []}
        paired = []
        for dev in data.get("paired") or []:
            mac = str(dev.get("mac") or "").upper()
            paired.append({
                "name": dev.get("name") or mac or "Bluetooth device",
                "mac": mac,
                "icon": "ph-headphones",
                "connected": mac in connected_macs,
                "audio_ready": bool(dev.get("audio_ready")),
                "battery": "",
            })
        dev = data.get("device")
        if not isinstance(dev, dict):
            if data.get("adapter_present"):
                return {
                    "name": "Bluetooth",
                    "icon": "ph-bluetooth",
                    "connected": False,
                    "battery": "",
                }, paired
            return {
                "name": "No Bluetooth adapter",
                "icon": "ph-bluetooth",
                "connected": False,
                "battery": "",
                "no_adapter": True,
            }, paired
        battery_pct = dev.get("battery_pct")
        return {
            "name": dev.get("name") or "Bluetooth device",
            "mac": dev.get("mac") or "",
            "icon": "ph-bluetooth",
            "connected": bool(dev.get("connected")),
            "audio_ready": bool(dev.get("audio_ready") or (dev.get("connected") and audio_route_ready)),
            "battery": f"{battery_pct}%" if battery_pct is not None else "",
            "transport_health": dev.get("transport_health") if isinstance(dev.get("transport_health"), dict) else {},
        }, paired

    def _live_device(self):
        dev, paired = self._live_state()
        self.paired_devices = paired
        return dev

    def on_action(self, action):
        if not self.rows:
            if action == "b":
                self.app.pop()
            return
        if action == "up":
            self.focus = max(0, self.focus - 1)
        elif action == "down":
            self.focus = min(min(len(self.rows), len(ROW_Y)) - 1, self.focus + 1)
        elif action == "b":
            self.app.pop()
        elif action == "a":
            self._activate(self.rows[self.focus][2])

    def _activate(self, key):
        if key == "audio":
            from .audio import AudioScreen
            self.app.push(AudioScreen(self.app, self.device))
        elif key == "paired":
            self.app.push(PairedDevicesScreen(self.app, self.paired_devices))
        elif key == "setup":
            from .dialogs import ConfirmDialog
            repair = self.audio_setup_label.startswith("Repair")
            self.app.push(ConfirmDialog(
                self.app, self,
                title="Repair Bluetooth audio?" if repair else "Start audio setup?",
                body=("AAhead will refresh the local audio runtime and only continue "
                      "to full setup if Bluetooth audio is still unavailable."
                      if repair else
                      "AAhead will install bundled Bluetooth audio components, "
                      "restart Bluetooth, and route system/game audio through "
                      "PulseAudio. Keep the dongle plugged in. Rollback is available."),
                confirm_label="Repair" if repair else "Start setup",
                body_max_chars=48,
                card_width=460,
                on_confirm=self._start_audio_setup))
        elif key == "no_adapter":
            from .bt_states import S2bNoAdapter
            self.app.push(S2bNoAdapter(self.app))
        elif key == "drivers":
            from .drivers import DriversScreen
            self.app.push(DriversScreen(self.app))
        elif key == "pair":
            from .pair import PairScreen
            self.app.push(PairScreen(self.app))
        elif key == "connect":
            self._start_device_job("connect")
        elif key == "disconnect":
            self._start_device_job("disconnect")
        elif key == "forget":
            from .dialogs import ConfirmDialog
            name = self.device["name"]
            self.app.push(ConfirmDialog(
                self.app, self,
                title="Forget device?",
                body=f"{name} will need to be paired again before you can use it.",
                confirm_label="Forget", danger=True,
                on_confirm=lambda: self._start_device_job("forget")))

    def _start_audio_setup(self):
        from .bt_states import S6Setup
        backend.gui_log("bluetooth.setup_confirmed")
        self.app.push(S6Setup(self.app))

    def _start_device_job(self, action):
        if self.job_id:
            self.app.toast("Bluetooth action is still running")
            return
        mac = self.device.get("mac") or ""
        if not mac:
            if action == "disconnect":
                self.device["connected"] = False
                self.app.toast(f"Disconnected {self.device['name']}")
            else:
                self.app.toast("No Bluetooth device selected")
            return
        calls = {
            "connect": backend.bluetooth_connect,
            "disconnect": backend.bluetooth_disconnect,
            "forget": backend.bluetooth_forget,
        }
        env = calls[action](mac)
        job = (((env or {}).get("data") or {}).get("job") or {})
        if not job.get("job_id"):
            self.app.toast((env or {}).get("title") or "Bluetooth action unavailable")
            return
        self.job_id = job["job_id"]
        self.job_action = action
        backend.gui_log("bluetooth.job_started", action=action, job_id=self.job_id, mac=mac)
        self.app.toast(job.get("title") or "Working…", ms=120000)

    def update(self, dt_ms):
        if not self.job_id:
            return
        self.poll_ms += dt_ms
        if self.poll_ms < 450:
            return
        self.poll_ms = 0.0
        env = backend.job_status(self.job_id)
        if not env:
            return
        job = ((env.get("data") or {}).get("job") or {})
        state = job.get("state")
        if state not in ("completed", "failed"):
            progress = job.get("progress") or {}
            label = progress.get("phase") or job.get("title") or "Working…"
            self.app.toast(label, ms=1200)
            return
        backend.gui_log("bluetooth.job_done", action=self.job_action, job_id=self.job_id, state=state, title=job.get("title"))
        self.app.toast(job.get("title") or ("Done" if state == "completed" else "Bluetooth action failed"))
        if state == "completed":
            if self.job_action == "disconnect":
                self.device["connected"] = False
            elif self.job_action == "forget":
                self.app.pop()
                self.job_id = None
                return
            else:
                live = self._live_device()
                if live:
                    self.device = live
            self.rows = self._rows_for_device()
            self._clamp_focus()
        self.job_id = None
        self.job_action = ""

    def draw(self, surf):
        w.status_bar(surf, title="Bluetooth", back=True)
        # hero
        dev = self.device
        if dev.get("connected") and dev.get("audio_ready"):
            state = ("Audio ready", t.Color.SUCCESS)
        elif dev.get("connected"):
            health = dev.get("transport_health") if isinstance(dev.get("transport_health"), dict) else {}
            health_state = health.get("state")
            if health_state == "service_discovery_stuck":
                state = ("Linked · services not resolved", t.Color.WARNING)
            elif health_state == "a2dp_transport_missing":
                state = ("Linked · A2DP transport missing", t.Color.WARNING)
            else:
                state = ("Linked · audio route missing", t.Color.WARNING)
        else:
            state = ("Disconnected", t.Color.TEXT_MUTED)
        subline = [state]
        if dev.get("battery"):
            subline.append((f"{dev['battery']} battery", t.Color.ACCENT_CYAN))
        w.hero_zone(surf, 60, dev["icon"], dev["name"], subline, height=100)

        # action rows
        visible_rows = self.rows[:len(ROW_Y)]
        if len(self.rows) > len(ROW_Y):
            backend.gui_log("bluetooth.rows_truncated", rows=len(self.rows), visible=len(visible_rows))
        for i, (label, icon, _key, danger) in enumerate(visible_rows):
            rect = pygame.Rect(ROW_X, ROW_Y[i], ROW_W, ROW_H)
            focused = (i == self.focus)
            icol = t.Color.ERROR if (danger and not focused) else None
            w.list_row(surf, rect, label, icon, focused=focused,
                       dim=(i >= 2), icon_color=icol)

        w.footer_hints(surf, [("A", "select"), ("B", "back")])


class PairedDevicesScreen:
    name = "paired-devices"

    def __init__(self, app, devices):
        self.app = app
        self.devices = list(devices or [])
        self.focus = 0
        self.job_id = None
        self.job_action = ""
        self.poll_ms = 0.0
        self._refresh_live()

    def on_resume(self):
        self._refresh_live()

    def _refresh_live(self):
        env = backend.bluetooth_status()
        if not env:
            return
        data = env.get("data") or {}
        connected_macs = {str(d.get("mac") or "").upper() for d in data.get("connected") or []}
        devices = []
        for dev in data.get("paired") or []:
            mac = str(dev.get("mac") or "").upper()
            devices.append({
                "name": dev.get("name") or mac or "Bluetooth device",
                "mac": mac,
                "icon": "ph-headphones",
                "connected": mac in connected_macs,
                "audio_ready": False,
                "battery": "",
            })
        current = data.get("device") if isinstance(data.get("device"), dict) else None
        if current:
            current_mac = str(current.get("mac") or "").upper()
            for dev in devices:
                if dev["mac"] == current_mac:
                    dev["connected"] = bool(current.get("connected"))
                    dev["audio_ready"] = bool(current.get("audio_ready"))
                    dev["battery"] = f"{current.get('battery_pct')}%" if current.get("battery_pct") is not None else ""
                    break
        self.devices = devices
        self.focus = min(self.focus, max(0, len(self.devices) - 1))

    def on_action(self, action):
        if action == "up" and self.devices:
            self.focus = max(0, self.focus - 1)
        elif action == "down" and self.devices:
            self.focus = min(len(self.devices) - 1, self.focus + 1)
        elif action == "b":
            self.app.pop()
        elif action == "a":
            if not self.devices:
                from .pair import PairScreen
                self.app.push(PairScreen(self.app))
                return
            self._start("connect")
        elif action == "x":
            self._start("forget")

    def _start(self, action):
        if self.job_id or not self.devices:
            return
        dev = self.devices[self.focus]
        mac = dev.get("mac") or ""
        if not mac:
            self.app.toast("No Bluetooth device selected")
            return
        env = backend.bluetooth_forget(mac) if action == "forget" else backend.bluetooth_connect(mac)
        job = (((env or {}).get("data") or {}).get("job") or {})
        if not job.get("job_id"):
            self.app.toast((env or {}).get("title") or "Bluetooth action unavailable")
            return
        self.job_id = job["job_id"]
        self.job_action = action
        self.app.toast("Forgetting device…" if action == "forget" else "Connecting…", ms=120000)

    def update(self, dt_ms):
        if not self.job_id:
            return
        self.poll_ms += dt_ms
        if self.poll_ms < 450:
            return
        self.poll_ms = 0.0
        env = backend.job_status(self.job_id)
        if not env:
            return
        job = ((env.get("data") or {}).get("job") or {})
        state = job.get("state")
        if state not in ("completed", "failed"):
            progress = job.get("progress") or {}
            label = progress.get("phase") or job.get("title") or ("Forgetting…" if self.job_action == "forget" else "Connecting…")
            self.app.toast(label, ms=1200)
            return
        self.app.toast(job.get("title") or ("Done" if state == "completed" else "Bluetooth action failed"))
        if state == "completed" and self.job_action == "forget" and self.devices:
            self.devices.pop(self.focus)
            self.focus = min(self.focus, max(0, len(self.devices) - 1))
        elif self.job_action == "connect":
            self._refresh_live()
        self.job_id = None
        self.job_action = ""

    def draw(self, surf):
        w.status_bar(surf, title="Paired devices", back=True)
        if not self.devices:
            title = fonts.render(t.Type.TITLE, "No paired devices", t.Color.TEXT_PRIMARY)
            surf.blit(title, (320 - title.get_width() // 2, 190))
            sub = fonts.render(t.Type.BODY_SM, "Pair headphones or speakers first.", t.Color.TEXT_SECONDARY)
            surf.blit(sub, (320 - sub.get_width() // 2, 220))
        for i, dev in enumerate(self.devices[:5]):
            rect = pygame.Rect(ROW_X, 60 + i * 72, ROW_W, 64)
            sub = "Audio ready" if dev.get("audio_ready") else "Linked" if dev.get("connected") else "Paired"
            w.device_row(surf, rect, dev.get("name") or "Bluetooth device", sub,
                         dev.get("icon") or "ph-headphones", 3, focused=(i == self.focus))
        if self.devices:
            hints = [("A", "connect"), ("X", "forget"), ("B", "back")]
        else:
            hints = [("A", "pair"), ("B", "back")]
        w.footer_hints(surf, hints)
