from __future__ import annotations

import json
import os
import queue
import threading
import time
import base64
import shutil
import uuid
from pathlib import Path

import webview
from serial.tools import list_ports

from app import Axes, MOTOR_FUNCTIONS, MavlinkWorker, motor_number_for_function


class ControlApi:
    """Small, local-only bridge between the CSS UI and the MAVLink worker."""

    def __init__(self):
        self.events: queue.Queue = queue.Queue()
        self.worker = MavlinkWorker(self.events)
        self.lock = threading.Lock()
        self.state = {
            "connected": False,
            "link": "DISCONNECTED",
            "vehicle": "No heartbeat",
            "armed": False,
            "status": "Ready. Select a USB port to begin.",
            "battery": None,
            "log_path": None,
            "servos": {str(i): {"function": None, "direction": None} for i in range(1, 17)},
        }
        self.servo_functions: dict[int, int] = {}
        self.motor_directions: dict[int, int] = {}
        self.battery_path = Path(__file__).resolve().parent / "battery_settings.json"
        self.model_path = Path(__file__).resolve().parent / "model_settings.json"
        self.model_cache = Path(__file__).resolve().parent / "runtime" / "vehicle_model.stl"
        self.model_cache_name = Path(__file__).resolve().parent / "runtime" / "vehicle_model_name.txt"
        self.model_uploads: dict[str, tuple[object, Path, str, int, int]] = {}
        self.model_lock = threading.Lock()
        # Leading underscore keeps the native WebView object out of js_api
        # serialization. Only the small public methods below are exposed.
        self._window = None
        self.maximized = False
        self.battery_config = {
            "cells": 4, "capacity_mah": 5500, "c_rating": 35,
            "full_v_per_cell": 4.2, "empty_v_per_cell": 3.3,
            "warning_percent": 20,
        }
        try:
            self.battery_config.update(json.loads(self.battery_path.read_text(encoding="utf-8")))
        except (OSError, ValueError, TypeError):
            pass

    def initial_state(self):
        return {
            "ports": [p.device for p in list_ports.comports()],
            "battery_config": self.battery_config,
            "state": self._snapshot(),
        }

    def refresh_ports(self):
        return [p.device for p in list_ports.comports()]

    def minimize_window(self):
        if self._window:
            self._window.minimize()
        return True

    def toggle_maximize_window(self):
        if not self._window:
            return False
        if self.maximized:
            self._window.restore()
        else:
            self._window.maximize()
        self.maximized = not self.maximized
        return self.maximized

    def close_window(self):
        if self._window:
            self._window.destroy()
        return True

    def begin_model_save(self, name, size):
        safe_name = Path(str(name)).name
        if not safe_name.lower().endswith(".stl"):
            raise ValueError("Choose an STL mesh")
        size = int(size)
        if not 0 < size <= 30 * 1024 * 1024:
            raise ValueError("STL must be between 1 byte and 30 MB")
        self.model_cache.parent.mkdir(exist_ok=True)
        token = uuid.uuid4().hex
        temp = self.model_cache.parent / f".{token}.upload"
        with self.model_lock:
            self.model_uploads[token] = (temp.open("wb"), temp, safe_name, size, 0)
        return token

    def append_model_chunk(self, token, data):
        try:
            chunk = base64.b64decode(data, validate=True)
        except (ValueError, TypeError) as exc:
            raise ValueError("Invalid STL data chunk") from exc
        with self.model_lock:
            upload = self.model_uploads.get(str(token))
            if not upload:
                raise ValueError("Model upload expired")
            handle, temp, name, expected, received = upload
            if received + len(chunk) > expected:
                raise ValueError("Model upload is larger than expected")
            handle.write(chunk)
            self.model_uploads[str(token)] = (handle, temp, name, expected, received + len(chunk))
        return True

    def finish_model_save(self, token):
        with self.model_lock:
            upload = self.model_uploads.pop(str(token), None)
            if not upload:
                raise ValueError("Model upload expired")
            handle, temp, name, expected, received = upload
            handle.close()
            if received != expected:
                temp.unlink(missing_ok=True)
                raise ValueError("Model upload was incomplete")
            os.replace(temp, self.model_cache)
            self.model_cache_name.write_text(name, encoding="utf-8")
        return True

    def get_saved_model_info(self):
        self._migrate_legacy_model()
        try:
            if self.model_cache.is_file() and self.model_cache.stat().st_size <= 30 * 1024 * 1024:
                name = self.model_cache_name.read_text(encoding="utf-8").strip() or "vehicle_model.stl"
                return {"name": name, "size": self.model_cache.stat().st_size}
        except (OSError, ValueError):
            pass
        return None

    def read_saved_model_chunk(self, offset, size=196608):
        offset, size = int(offset), min(int(size), 196608)
        if offset < 0 or size < 1 or not self.model_cache.is_file():
            raise ValueError("Saved model is unavailable")
        with self.model_cache.open("rb") as stream:
            stream.seek(offset)
            content = stream.read(size)
        return {"offset": offset, "data": base64.b64encode(content).decode("ascii"), "eof": offset + len(content) >= self.model_cache.stat().st_size}

    def _migrate_legacy_model(self):
        if self.model_cache.is_file():
            return
        try:
            path = Path(json.loads(self.model_path.read_text(encoding="utf-8"))["path"])
            if path.is_file() and path.suffix.lower() == ".stl" and path.stat().st_size <= 30 * 1024 * 1024:
                self.model_cache.parent.mkdir(exist_ok=True)
                shutil.copyfile(path, self.model_cache)
                self.model_cache_name.write_text(path.name, encoding="utf-8")
        except (OSError, ValueError, TypeError, KeyError):
            pass

    def connect(self, port, baud=57600, rate=20):
        if not port:
            raise ValueError("Select a USB serial port")
        self.worker.rate_hz = int(rate)
        self.worker.connect(str(port), int(baud))
        with self.lock:
            self.state.update(
                connected=True, link="CONNECTING",
                status=f"Opening {port} at {baud} baud…",
            )
        return True

    def disconnect(self):
        self.worker.disconnect()
        return True

    def set_axes(self, forward, turn, vertical):
        self.worker.set_axes(Axes(
            max(-1.0, min(1.0, float(forward))),
            0.0,
            max(-1.0, min(1.0, float(vertical))),
            max(-1.0, min(1.0, float(turn))),
        ))
        return True

    def set_gain(self, value):
        value = max(5, min(50, int(value)))
        self.worker.gain = value / 100.0
        return value

    def command(self, name):
        if name == "manual":
            self.worker.log("OPERATOR", "clicked SET MANUAL")
            self.worker.set_manual_mode()
        elif name == "arm":
            self.worker.log("OPERATOR", "clicked ARM")
            self.worker.command_arm(True)
        elif name == "disarm":
            self.worker.log("OPERATOR", "clicked DISARM")
            self.worker.command_arm(False)
        else:
            raise ValueError("Unknown command")
        return True

    def set_servo_function(self, servo, motor):
        servo, motor = int(servo), int(motor)
        if not 1 <= servo <= 16 or not 0 <= motor <= 12:
            raise ValueError("Invalid servo or motor")
        function = 0 if motor == 0 else (32 + motor if motor <= 8 else 73 + motor)
        self.worker.set_parameter(f"SERVO{servo}_FUNCTION", function)
        return True

    def reverse_motor(self, servo):
        servo = int(servo)
        motor = motor_number_for_function(self.servo_functions.get(servo))
        if motor is None:
            raise ValueError(f"SERVO{servo} is not assigned to a motor")
        current = self.motor_directions.get(motor)
        if current is None:
            raise ValueError(f"MOT_{motor}_DIRECTION has not been read yet")
        self.worker.set_parameter(f"MOT_{motor}_DIRECTION", -1 if current == 1 else 1)
        return True

    def save_battery(self, config):
        clean = {
            "cells": int(config["cells"]),
            "capacity_mah": int(config["capacity_mah"]),
            "c_rating": float(config["c_rating"]),
            "full_v_per_cell": float(config["full_v_per_cell"]),
            "empty_v_per_cell": float(config["empty_v_per_cell"]),
            "warning_percent": int(config["warning_percent"]),
        }
        if clean["cells"] < 1 or clean["capacity_mah"] <= 0 or clean["c_rating"] <= 0:
            raise ValueError("Battery cells, capacity and C rating must be positive")
        if not clean["full_v_per_cell"] > clean["empty_v_per_cell"] > 0:
            raise ValueError("Full voltage per cell must be greater than empty voltage")
        if not 0 <= clean["warning_percent"] <= 100:
            raise ValueError("Warning percentage must be between 0 and 100")
        self.battery_config = clean
        self.battery_path.write_text(json.dumps(clean, indent=2) + "\n", encoding="utf-8")
        return clean

    def open_logs(self):
        path = Path(__file__).resolve().parent / "logs"
        path.mkdir(exist_ok=True)
        os.startfile(path)
        return True

    def poll(self):
        while True:
            try:
                event = self.events.get_nowait()
            except queue.Empty:
                break
            self._consume(event)
        if self.worker.target_system and time.monotonic() - self.worker.last_heartbeat > 5:
            with self.lock:
                self.state["link"] = (
                    "DATA OK / HEARTBEAT DELAYED"
                    if time.monotonic() - self.worker.last_rx < 2 else "ALL LINK DATA LOST"
                )
        return self._snapshot()

    def _snapshot(self):
        with self.lock:
            return json.loads(json.dumps(self.state))

    def _consume(self, event):
        kind = event[0]
        with self.lock:
            if kind == "heartbeat":
                _, system, autopilot, mode, armed = event
                self.state.update(
                    connected=True, link="LINK OK", vehicle=f"{mode} • {'ARMED' if armed else 'DISARMED'}",
                    armed=bool(armed),
                )
            elif kind in ("status", "statustext"):
                self.state["status"] = event[1]
            elif kind == "ack":
                self.state["status"] = f"Command {event[1]} ACK: {event[2]}"
            elif kind == "battery":
                voltage = float(event[1])
                cells = self.battery_config["cells"]
                per_cell = voltage / cells
                low, high = self.battery_config["empty_v_per_cell"], self.battery_config["full_v_per_cell"]
                percent = max(0, min(100, (per_cell - low) * 100 / (high - low)))
                self.state["battery"] = {
                    "voltage": voltage, "per_cell": per_cell, "percent": percent,
                    "cells": cells, "low": percent <= self.battery_config["warning_percent"],
                }
            elif kind == "servo_param":
                _, param_id, value = event
                servo = int(param_id[5:param_id.index("_")])
                self.servo_functions[servo] = int(value)
                self._update_servo(servo)
            elif kind == "motor_direction":
                _, motor, value = event
                self.motor_directions[int(motor)] = int(value)
                for servo in range(1, 17):
                    self._update_servo(servo)
            elif kind == "param_saved":
                self.state["status"] = f"Confirmed: {event[1]}={event[2]}"
            elif kind == "log_path":
                self.state["log_path"] = event[1]
            elif kind == "error":
                self.state["status"] = event[1]
            elif kind == "disconnected":
                self.state.update(connected=False, link="DISCONNECTED", vehicle="No heartbeat",
                                  armed=False, battery=None, status="Disconnected")

    def _update_servo(self, servo):
        function = self.servo_functions.get(servo)
        motor = motor_number_for_function(function)
        self.state["servos"][str(servo)] = {
            "function": function,
            "label": MOTOR_FUNCTIONS.get(function, "Unassigned") if function is not None else "Waiting…",
            "motor": motor,
            "direction": self.motor_directions.get(motor) if motor else None,
        }

    def close(self):
        self.worker.disconnect()


def main():
    api = ControlApi()
    ui_path = Path(__file__).resolve().parent / "web" / "index.html"
    window = webview.create_window(
        "SUB Control", url=str(ui_path), js_api=api,
        width=1120, height=760, min_size=(900, 650),
        background_color="#111719", frameless=True, easy_drag=False,
        shadow=True,
    )
    api._window = window
    window.events.maximized += lambda: setattr(api, "maximized", True)
    window.events.restored += lambda: setattr(api, "maximized", False)
    try:
        webview.start(debug=False)
    finally:
        api.close()


if __name__ == "__main__":
    main()