# pumpkin-api-py: Pumpkin plugin API for Python
# Copyright (C) 2024 alex
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import traceback
from typing import Any, Callable, Dict, Optional, Type, TypedDict

import wit_world
import wit_world.exports
from componentize_py_types import Err
from wit_world.exports import metadata
from wit_world.imports import command, context, event, scheduler, server, text, world


class AiGoal(TypedDict):
    can_start: Callable[[server.Server, world.Entity], bool]
    should_continue: Callable[[server.Server, world.Entity], bool]
    start: Callable[[server.Server, world.Entity], None]
    tick: Callable[[server.Server, world.Entity], None]
    stop: Callable[[server.Server, world.Entity], None]


_PLUGIN_INSTANCE: Optional["Plugin"] = None
_EVENT_HANDLERS: Dict[int, Callable[[server.Server, Any], Any]] = {}
_COMMAND_HANDLERS: Dict[
    int, Callable[[command.CommandSender, server.Server, command.ConsumedArgs], int]
] = {}
_TASK_HANDLERS: Dict[int, Callable[[server.Server], None]] = {}
_AI_GOALS: Dict[int, AiGoal] = {}
_GENERATORS: Dict[int, Callable[[world.GenerationPhase, world.ChunkBuffer], None]] = {}

_NEXT_HANDLER_ID: int = 0


def _exception_message(exc: Exception) -> str:
    if hasattr(exc, "value"):
        value = getattr(exc, "value")
        if value is not None:
            return str(value)
    message = str(exc)
    if message:
        return message
    return "".join(traceback.format_exception_only(type(exc), exc)).strip()


def _format_exception(exc: Exception) -> str:
    details = "".join(
        traceback.format_exception(type(exc), exc, exc.__traceback__)
    ).strip()
    if details:
        return details
    return _exception_message(exc)


def _get_next_handler_id() -> int:
    global _NEXT_HANDLER_ID
    handler_id = _NEXT_HANDLER_ID
    _NEXT_HANDLER_ID += 1
    return handler_id


class Plugin:
    def __init__(self):
        self._pending_events = []

    def metadata(self) -> metadata.PluginMetadata:
        raise NotImplementedError

    def on_load(self, ctx: context.Context) -> None:
        for event_type, handler, priority, blocking in self._pending_events:
            self.register_event(ctx, event_type, handler, priority, blocking)

    def on_unload(self, ctx: context.Context) -> None:
        pass

    def register_event(
        self,
        ctx: context.Context,
        event_type: event.EventType,
        handler: Callable[[server.Server, Any], Any],
        priority=event.EventPriority.NORMAL,
        blocking=True,
    ):
        handler_id = _get_next_handler_id()
        _EVENT_HANDLERS[handler_id] = handler
        ctx.register_event(handler_id, event_type, priority, blocking)

    def register_handler(
        self,
        handler: Callable[
            [command.CommandSender, server.Server, command.ConsumedArgs], int
        ],
    ) -> int:
        handler_id = _get_next_handler_id()
        _COMMAND_HANDLERS[handler_id] = handler
        return handler_id

    def register_command(
        self,
        ctx: context.Context,
        cmd: command.Command,
        handler: Callable[
            [command.CommandSender, server.Server, command.ConsumedArgs], int
        ],
        permission: str = "",
        extra_nodes=None,
    ):
        handler_id = self.register_handler(handler)
        cmd.execute_with_handler_id(handler_id)
        if extra_nodes:
            for node in extra_nodes:
                node.execute_with_handler_id(handler_id)
        ctx.register_command(cmd, permission)

    def register_ai_goal(self, goal: AiGoal) -> int:
        goal_id = _get_next_handler_id()
        _AI_GOALS[goal_id] = goal
        return goal_id

    def register_chunk_generator(
        self, generator: Callable[[world.GenerationPhase, world.ChunkBuffer], None]
    ) -> int:
        generator_id = _get_next_handler_id()
        _GENERATORS[generator_id] = generator
        return generator_id

    def schedule_delayed_task(
        self, delay_ticks: int, handler: Callable[[server.Server], None]
    ) -> int:
        handler_id = _get_next_handler_id()
        _TASK_HANDLERS[handler_id] = handler
        return scheduler.schedule_delayed_task(handler_id, delay_ticks)

    def schedule_repeating_task(
        self,
        delay_ticks: int,
        period_ticks: int,
        handler: Callable[[server.Server], None],
    ) -> int:
        handler_id = _get_next_handler_id()
        _TASK_HANDLERS[handler_id] = handler
        return scheduler.schedule_repeating_task(handler_id, delay_ticks, period_ticks)

    def event_handler(
        self,
        event_type: event.EventType,
        priority=event.EventPriority.NORMAL,
        blocking=True,
    ):
        def decorator(func):
            self._pending_events.append((event_type, func, priority, blocking))
            return func

        return decorator

    def handle_ipc_message(self, sender: str, message: bytes) -> bytes:
        raise NotImplementedError


def register_plugin(plugin_class: Type[Plugin]):
    global _PLUGIN_INSTANCE
    _PLUGIN_INSTANCE = plugin_class()


class WitWorldImpl(wit_world.WitWorld):
    def init_plugin(self) -> None:
        pass

    def on_load(self, context: context.Context) -> None:
        if _PLUGIN_INSTANCE:
            try:
                # In WIT, on-load returns result<_, string>
                # componentize-py handles this by expecting None or raising Err
                _PLUGIN_INSTANCE.on_load(context)
            except Err as e:
                raise Err(_exception_message(e))
            except Exception as e:
                raise Err(_format_exception(e))

    def on_unload(self, context: context.Context) -> None:
        if _PLUGIN_INSTANCE:
            try:
                _PLUGIN_INSTANCE.on_unload(context)
            except Err as e:
                raise Err(_exception_message(e))
            except Exception as e:
                raise Err(_format_exception(e))

    def handle_event(
        self, event_id: int, server: server.Server, event: event.Event
    ) -> event.Event:
        if event_id in _EVENT_HANDLERS:
            # Pass the inner event data to the handler
            result = _EVENT_HANDLERS[event_id](server, event.value)
            # If the handler returns something, assume it's the updated event data
            if result is not None:
                event.value = result
        return event

    def handle_command(
        self,
        command_id: int,
        sender: command.CommandSender,
        server: server.Server,
        args: command.ConsumedArgs,
    ) -> int:
        if command_id in _COMMAND_HANDLERS:
            try:
                return _COMMAND_HANDLERS[command_id](sender, server, args)
            except Exception as e:
                if hasattr(e, "value") and isinstance(
                    e.value,
                    (
                        command.CommandError_InvalidConsumption,
                        command.CommandError_InvalidRequirement,
                        command.CommandError_PermissionDenied,
                        command.CommandError_CommandFailed,
                    ),
                ):
                    raise e
                raise Err(
                    command.CommandError_CommandFailed(
                        text.TextComponent.text(_format_exception(e))
                    )
                )

        raise Err(
            command.CommandError_CommandFailed(
                text.TextComponent.text(
                    f"no handler registered for command id {command_id}"
                )
            )
        )

    def handle_task(self, handler_id: int, server: server.Server) -> None:
        if handler_id in _TASK_HANDLERS:
            _TASK_HANDLERS[handler_id](server)

    def handle_ipc_message(self, sender: str, message: bytes) -> bytes:
        if _PLUGIN_INSTANCE:
            try:
                return _PLUGIN_INSTANCE.handle_ipc_message(sender, message)
            except Err as e:
                raise Err(_exception_message(e))
            except Exception as e:
                raise Err(_format_exception(e))

        raise Err("Internal error: No plugin instance available?")

    def handle_ai_goal_can_start(
        self, goal_id: int, server: server.Server, entity: world.Entity
    ) -> bool:
        if goal_id in _AI_GOALS:
            return _AI_GOALS[goal_id]["can_start"](server, entity)

        raise Exception(f"No AI goal registered for ID {goal_id}")

    def handle_ai_goal_should_continue(
        self, goal_id: int, server: server.Server, entity: world.Entity
    ) -> bool:
        if goal_id in _AI_GOALS:
            return _AI_GOALS[goal_id]["should_continue"](server, entity)

        raise Exception(f"No AI goal registered for ID {goal_id}")

    def handle_ai_goal_start(
        self, goal_id: int, server: server.Server, entity: world.Entity
    ) -> None:
        if goal_id in _AI_GOALS:
            _AI_GOALS[goal_id]["start"](server, entity)
            return
        raise Exception(f"No AI goal registered for ID {goal_id}")

    def handle_ai_goal_tick(
        self, goal_id: int, server: server.Server, entity: world.Entity
    ) -> None:
        if goal_id in _AI_GOALS:
            _AI_GOALS[goal_id]["tick"](server, entity)
            return
        raise Exception(f"No AI goal registered for ID {goal_id}")

    def handle_ai_goal_stop(
        self, goal_id: int, server: server.Server, entity: world.Entity
    ) -> None:
        if goal_id in _AI_GOALS:
            _AI_GOALS[goal_id]["stop"](server, entity)
            return
        raise Exception(f"No AI goal registered for ID {goal_id}")

    def handle_generate_phase(
        self, generator_id: int, phase: world.GenerationPhase, chunk: world.ChunkBuffer
    ) -> None:
        if generator_id in _AI_GOALS:
            _GENERATORS[generator_id](phase, chunk)
            return
        raise Exception(f"No generator registered for ID {generator_id}")


class MetadataImpl:
    def get_metadata(self) -> metadata.PluginMetadata:
        if _PLUGIN_INSTANCE:
            return _PLUGIN_INSTANCE.metadata()
        return metadata.PluginMetadata("unknown", "0.0.0", [], "No metadata", [], [])