from datetime import datetime, timezone
from typing import Sequence

import anyio
import pytest
from loguru import logger

from exo.master.main import Master
from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.backends import Backend
from exo.shared.types.commands import (
    CommandId,
    ForwarderCommand,
    ForwarderDownloadCommand,
    PlaceInstance,
    TextGeneration,
)
from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
from exo.shared.types.events import (
    Event,
    GlobalForwarderEvent,
    IndexedEvent,
    InstanceCreated,
    LocalForwarderEvent,
    NodeGatheredInfo,
    TaskCreated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import (
    MemoryUsage,
)
from exo.shared.types.tasks import TaskStatus
from exo.shared.types.tasks import TextGeneration as TextGenerationTask
from exo.shared.types.text_generation import (
    InputMessage,
    InputMessageContent,
    TextGenerationTaskParams,
)
from exo.shared.types.worker.instances import (
    InstanceMeta,
    MlxRingInstance,
    ShardAssignments,
)
from exo.shared.types.worker.runners import ShardWithId
from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
from exo.utils.channels import channel
from exo.utils.info_gatherer.info_gatherer import NodeBackends


class MockAggregator:
    def dump(self) -> dict[str, str]:
        return {}


class MockStorage:
    async def get(self, _: str) -> None:
        return None

    async def put(self, _1: str, _2: str) -> None:
        return None

    async def delete(self, _: str) -> None:
        return None

    async def dump(self, _: str) -> dict[str, str]:
        return {}


@pytest.mark.asyncio
async def test_master():
    node_id = NodeId("yoooo")
    session_id = SessionId(master_node_id=node_id, election_clock=0)

    ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
    command_sender, co_receiver = channel[ForwarderCommand]()
    local_event_sender, le_receiver = channel[LocalForwarderEvent]()
    fcds, _fcdr = channel[ForwarderDownloadCommand]()
    ev_send, ev_recv = channel[Event]()

    async def mock_event_router():
        idx = 0
        sid = SystemId()
        with ev_recv as master_events:
            async for event in master_events:
                await local_event_sender.send(
                    LocalForwarderEvent(
                        origin=sid,
                        origin_idx=idx,
                        session=session_id,
                        event=event,
                    )
                )
                idx += 1

    all_events: list[IndexedEvent] = []

    def _get_events() -> Sequence[IndexedEvent]:
        orig_events = global_event_receiver.collect()
        for e in orig_events:
            all_events.append(
                IndexedEvent(
                    event=e.event,
                    idx=len(all_events),  # origin=e.origin,
                )
            )
        return all_events

    master = Master(
        node_id,
        session_id,
        event_sender=ev_send,
        global_event_sender=ge_sender,
        local_event_receiver=le_receiver,
        command_receiver=co_receiver,
        download_command_sender=fcds,
        aggregator=MockAggregator(),  # pyright: ignore[reportArgumentType]
        storage=MockStorage(),  # pyright: ignore[reportArgumentType]
    )
    logger.info("run the master")
    async with anyio.create_task_group() as tg:
        tg.start_soon(master.run)
        tg.start_soon(mock_event_router)

        # inject a NodeGatheredInfo event
        logger.info("inject a NodeGatheredInfo event")
        await local_event_sender.send(
            LocalForwarderEvent(
                origin_idx=0,
                origin=SystemId("Worker"),
                session=session_id,
                event=(
                    NodeGatheredInfo(
                        when=str(datetime.now(tz=timezone.utc)),
                        node_id=node_id,
                        info=MemoryUsage(
                            ram_total=Memory.from_bytes(678948 * 1024),
                            ram_available=Memory.from_bytes(678948 * 1024),
                            swap_total=Memory.from_bytes(0),
                            swap_available=Memory.from_bytes(0),
                        ),
                    )
                ),
            )
        )
        await local_event_sender.send(
            LocalForwarderEvent(
                origin_idx=1,
                origin=SystemId("Worker"),
                session=session_id,
                event=(
                    NodeGatheredInfo(
                        when=str(datetime.now(tz=timezone.utc)),
                        node_id=node_id,
                        info=NodeBackends(backends=[Backend.MlxMetal]),
                    )
                ),
            )
        )

        # wait for initial topology event
        logger.info("wait for initial topology event")
        while len(list(master.state.topology.list_nodes())) == 0:
            await anyio.sleep(0.001)
        while len(master.state.node_memory) == 0:
            await anyio.sleep(0.001)
        while len(master.state.node_backends) == 0:
            await anyio.sleep(0.001)

        logger.info("inject a CreateInstance Command")
        await command_sender.send(
            ForwarderCommand(
                origin=SystemId("API"),
                command=(
                    PlaceInstance(
                        command_id=CommandId(),
                        model_card=ModelCard(
                            model_id=ModelId("llama-3.2-1b"),
                            n_layers=16,
                            storage_size=Memory.from_bytes(678948),
                            hidden_size=7168,
                            supports_tensor=True,
                            tasks=[ModelTask.TextGeneration],
                            backends=[Backend.MlxMetal],
                        ),
                        sharding=Sharding.Pipeline,
                        instance_meta=InstanceMeta.MlxRing,
                        min_nodes=1,
                    )
                ),
            )
        )
        logger.info("wait for an instance")
        while len(master.state.instances.keys()) == 0:
            await anyio.sleep(0.001)
        logger.info("inject a TextGeneration Command")
        await command_sender.send(
            ForwarderCommand(
                origin=SystemId("API"),
                command=(
                    TextGeneration(
                        command_id=CommandId(),
                        task_params=TextGenerationTaskParams(
                            model=ModelId("llama-3.2-1b"),
                            input=[
                                InputMessage(
                                    role="user",
                                    content=InputMessageContent("Hello, how are you?"),
                                )
                            ],
                        ),
                    )
                ),
            )
        )
        while len(_get_events()) < 4:
            await anyio.sleep(0.01)

        events = _get_events()
        assert len(events) == 4
        assert events[0].idx == 0
        assert events[1].idx == 1
        assert events[2].idx == 2
        assert events[3].idx == 3
        assert isinstance(events[0].event, NodeGatheredInfo)
        assert isinstance(events[1].event, NodeGatheredInfo)
        assert isinstance(events[2].event, InstanceCreated)
        created_instance = events[2].event.instance
        assert isinstance(created_instance, MlxRingInstance)
        runner_id = created_instance.shard_assignments.shards[0].runner_id
        # Validate the shard assignments
        expected_shard_assignments = ShardAssignments(
            model_id=ModelId("llama-3.2-1b"),
            shards=[
                ShardWithId(
                    node_id,
                    runner_id,
                    PipelineShardMetadata(
                        start_layer=0,
                        end_layer=16,
                        n_layers=16,
                        model_card=ModelCard(
                            model_id=ModelId("llama-3.2-1b"),
                            n_layers=16,
                            storage_size=Memory.from_bytes(678948),
                            hidden_size=7168,
                            supports_tensor=True,
                            tasks=[ModelTask.TextGeneration],
                            backends=[Backend.MlxMetal],
                        ),
                        device_rank=0,
                        world_size=1,
                    ),
                )
            ],
            primary_output_node=0,
        )
        assert created_instance.shard_assignments == expected_shard_assignments
        # For single-node, hosts_by_node should have one entry with self-binding
        assert len(created_instance.hosts_by_node) == 1
        assert node_id in created_instance.hosts_by_node
        assert len(created_instance.hosts_by_node[node_id]) == 1
        assert created_instance.hosts_by_node[node_id][0].ip == "0.0.0.0"
        assert created_instance.ephemeral_port > 0
        assert isinstance(events[3].event, TaskCreated)
        assert events[3].event.task.task_status == TaskStatus.Pending
        assert isinstance(events[3].event.task, TextGenerationTask)
        assert events[3].event.task.task_params == TextGenerationTaskParams(
            model=ModelId("llama-3.2-1b"),
            input=[
                InputMessage(
                    role="user", content=InputMessageContent("Hello, how are you?")
                )
            ],
        )

        ev_send.close()
        await master.shutdown()
