from __future__ import annotations

import json
from datetime import datetime
from threading import Lock
from typing import Any, Callable, Protocol

from one_run.constants import JSON_SCHEMA_VERSION
from one_run.run_dir import RunPaths


class EventSink(Protocol):
    def emit(self, paths: RunPaths, payload: dict[str, Any]) -> None:
        ...


class JsonlEventSink:
    def emit(self, paths: RunPaths, payload: dict[str, Any]) -> None:
        paths.events.parent.mkdir(parents=True, exist_ok=True)
        with paths.events.open("a", encoding="utf-8") as f:
            f.write(json.dumps(payload))
            f.write("\n")


class EventManager:
    def __init__(self, sinks: list[EventSink]) -> None:
        self._sinks = sinks

    def emit(self, paths: RunPaths, payload: dict[str, Any]) -> None:
        for sink in self._sinks:
            sink.emit(paths, payload)


EventSinkFactory = Callable[[], EventSink]

_sink_factories: list[EventSinkFactory] = []
_sink_factories_lock = Lock()


def register_event_sink_factory(factory: EventSinkFactory) -> Callable[[], None]:
    with _sink_factories_lock:
        _sink_factories.append(factory)

    def _unregister() -> None:
        with _sink_factories_lock:
            try:
                _sink_factories.remove(factory)
            except ValueError:
                return

    return _unregister


def build_event_payload(
    *,
    run_id: str,
    event: str,
    state: str | None = None,
    message: str | None = None,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "schema_version": JSON_SCHEMA_VERSION,
        "timestamp": datetime.now().isoformat(),
        "run_id": run_id,
        "event": event,
    }
    if state is not None:
        payload["state"] = state
    if message is not None:
        payload["message"] = message
    if extra:
        payload.update(extra)
    return payload


def default_event_manager() -> EventManager:
    with _sink_factories_lock:
        extra = [factory() for factory in _sink_factories]
    return EventManager([JsonlEventSink(), *extra])


def emit_event(
    paths: RunPaths,
    *,
    run_id: str,
    event: str,
    state: str | None = None,
    message: str | None = None,
    extra: dict[str, Any] | None = None,
) -> None:
    payload = build_event_payload(
        run_id=run_id,
        event=event,
        state=state,
        message=message,
        extra=extra,
    )
    default_event_manager().emit(paths, payload)
