# Copyright (C) 2026 Hector van der Aa <hector@h3cx.dev>
# Copyright (C) 2026 Association Exergie <association.exergie@gmail.com>
# SPDX-License-Identifier: GPL-3.0-or-later
import asyncio
from functools import partial
import uuid


from dynalab.state import AppState
from dynalab_protocol.models import (
    ConnectorHello,
    DynaLabHello,
    HandshakeAccepted,
    HandshakeRejected,
    Heartbeat,
    VersionDescriptor,
    machine_timestamp_ms,
)
from dynalab_protocol.wire import write_message, read_message


async def _heartbeat_loop(writer: asyncio.StreamWriter, period_ms: int = 1000) -> None:
    sequence: int = 0

    while True:
        heartbeat = Heartbeat(sequence=sequence, send_timestamp=machine_timestamp_ms())
        await write_message(writer, heartbeat)
        print(f"Sent heartbeat {sequence}")

        sequence += 1
        await asyncio.sleep(float(period_ms) / 1000)


async def _timeout_checker(
    state: AppState,
    connector_uuid: uuid.UUID,
    max_timeout: int,
    timeout_event: asyncio.Event,
) -> None:
    while True:
        last_heartbeat = state.last_heartbeat.get(uuid)
        if last_heartbeat is not None:
            if machine_timestamp_ms() > last_heartbeat + max_timeout:
                timeout_event.set()
        await asyncio.sleep(0.5)


async def _handle_connector(
    state: AppState, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
) -> None:
    peer = writer.get_extra_info("peername")
    print(f"Connector connected {peer}")

    try:
        dynalab_hello = DynaLabHello(
            instance_id=uuid.uuid4(),
            version=VersionDescriptor(type="alpha", major=0, minor=0, patch=1),
        )
        await write_message(writer, dynalab_hello)

        message = await asyncio.wait_for(read_message(reader), timeout=5.0)

        if not isinstance(message, ConnectorHello):
            raise ValueError(f"Expected connector_hello, received {message.type}")

        connector_hello = message

        print(f"Connector connected: {connector_hello}")

        if len(connector_hello.signals) == 0:
            handshake_rejected = HandshakeRejected(
                reason="Signals list is empty, cannot connecto to empty connector"
            )
            await write_message(writer, handshake_rejected)
            print("Rejecting handshake due to empty list")
            return
        handshake_accepted = HandshakeAccepted(
            accepted_signals=[signal.id for signal in connector_hello.signals]
        )
        await write_message(writer, handshake_accepted)

        heartbeat_task = asyncio.create_task(_heartbeat_loop(writer))
        timeout_event = asyncio.Event()
        timeout_task = asyncio.create_task(
            _timeout_checker(
                state,
                connector_hello.connector_uuid,
                dynalab_hello.heartbeat_timeout_ms,
                timeout_event,
            )
        )

        while not timeout_event.is_set():
            try:
                message = await asyncio.wait_for(read_message(reader), timeout=0.1)
                if isinstance(message, Heartbeat):
                    state.last_heartbeat[connector_hello.connector_uuid] = (
                        message.return_timestamp
                    )
            except TimeoutError:
                continue
            except ConnectionError:
                return

    finally:
        heartbeat_task.cancel()
        timeout_task.cancel()
        print(f"Connector disconnected {peer}")
        writer.close()
        await writer.wait_closed()


async def _json_main(state: AppState) -> None:
    server = await asyncio.start_server(
        partial(_handle_connector, state), host="127.0.0.1", port=8765
    )

    print("Json connector server listening on 127.0.0.1:8765")

    async with server:
        while not state.stop_event.is_set():
            await asyncio.sleep(0.1)


def json_thread_main(state: AppState) -> None:
    asyncio.run(_json_main(state))