"""EventBusMixin — event bus, subscriptions, campaigns, and scheduled tasks."""
from __future__ import annotations

import importlib
import inspect
import json
import logging
from datetime import UTC, datetime
from typing import Any

from logos_adk.crew_checkpoint_store import JSONObject, ensure_json_object
from logos_adk.errors import RoutingError
from logos_adk.event_bus import EventBus, EventRecord, InMemoryEventBus
from logos_adk.runnable import Runnable
from logos_adk.subscription_registry import subscription_handler_registry
from logos_adk.team_config import EventBusConfig
from logos_adk.types import Response, Usage

from logos_adk.team._types import AgentLike, _TeamEventSubscription, CampaignState

_SCHEDULED_TASK_TOPIC_PREFIX = "__scheduled_task__:"


class EventBusMixin:
    """Event bus management, subscriptions, campaigns, and scheduled tasks."""

    # -- public API ----------------------------------------------------------

    def subscribe(
        self,
        agent_name: str,
        topic: str,
        *,
        handler: Any | None = None,
        publish_to: str | None = None,
        name: str | None = None,
    ) -> Any:
        """Subscribe an agent or custom handler to a topic on the shared event bus."""
        self._validate_topic_access(agent_name, topic, action="subscribe")
        if publish_to is not None:
            self._validate_topic_access(agent_name, publish_to, action="publish")
        target_agent = self.get_agent(agent_name)
        if target_agent is None:
            raise ValueError(f"Agent '{agent_name}' is not registered in team '{self.name}'")
        target = target_agent if handler is None else self._resolve_subscription_handler(target_agent, handler)
        self._ensure_event_bus()
        self._topic_subscriptions.append(
            _TeamEventSubscription(
                agent_name=agent_name,
                topic=topic.strip(),
                handler_name=name or getattr(target, "name", None) or getattr(target, "__name__", None) or agent_name,
                target=target,
                publish_to=publish_to.strip() if isinstance(publish_to, str) else publish_to,
            )
        )
        return self

    def publish_event(
        self,
        from_agent: str,
        topic: str,
        payload: dict[str, Any],
        *,
        campaign_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        dedupe_key: str | None = None,
        available_in_seconds: float = 0.0,
    ) -> EventRecord:
        """Publish an event into the team's shared event bus."""
        if from_agent not in self._agents:
            raise RoutingError(f"Agent '{from_agent}' is not registered in team '{self.name}'")
        topic_name = topic.strip()
        self._validate_topic_access(from_agent, topic_name, action="publish")
        event_bus = self._ensure_event_bus()
        return event_bus.publish(
            topic_name,
            ensure_json_object(payload, path=f"team_event[{topic_name}]"),
            campaign_id=campaign_id,
            correlation_id=correlation_id,
            run_id=run_id,
            publisher=from_agent,
            dedupe_key=dedupe_key,
            available_in_seconds=available_in_seconds,
        )

    def schedule_task(
        self,
        from_agent: str,
        task: str,
        *,
        to_agent: str | None = None,
        delay_seconds: float | None = None,
        due_at: datetime | None = None,
        session_id: str | None = None,
        campaign_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        dedupe_key: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> EventRecord:
        """Schedule a future agent task without manually wiring event topics."""
        from logos_adk.topology import validate_send

        target_agent = (to_agent or from_agent).strip()
        if target_agent not in self._agents:
            raise ValueError(f"Agent '{target_agent}' is not registered in team '{self.name}'")
        if from_agent != target_agent:
            validate_send(from_agent, target_agent, self._topology)
        payload = ensure_json_object(
            {
                "task": task,
                "target_agent": target_agent,
                "source_agent": from_agent,
                "session_id": session_id,
                "metadata": metadata or {},
            },
            path="team.schedule_task.payload",
        )
        return self.schedule_event(
            from_agent,
            self._scheduled_task_topic(target_agent),
            payload,
            delay_seconds=delay_seconds,
            due_at=due_at,
            campaign_id=campaign_id,
            correlation_id=correlation_id or session_id,
            run_id=run_id,
            dedupe_key=dedupe_key,
        )

    def schedule_event(
        self,
        from_agent: str,
        topic: str,
        payload: dict[str, Any],
        *,
        delay_seconds: float | None = None,
        due_at: datetime | None = None,
        campaign_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        dedupe_key: str | None = None,
    ) -> EventRecord:
        """Schedule a future event on the team's shared event bus."""
        if delay_seconds is not None and due_at is not None:
            raise ValueError("Pass either delay_seconds or due_at, not both")
        scheduled_delay = self._scheduled_delay_seconds(delay_seconds=delay_seconds, due_at=due_at)
        return self.publish_event(
            from_agent,
            topic,
            payload,
            campaign_id=campaign_id,
            correlation_id=correlation_id,
            run_id=run_id,
            dedupe_key=dedupe_key,
            available_in_seconds=scheduled_delay,
        )

    def event_bus(self) -> EventBus:
        """Return the team's event bus, creating an in-memory one if needed."""
        return self._ensure_event_bus()

    def list_campaign_events(self, campaign_id: str) -> list[EventRecord]:
        """List all persisted events associated with a campaign."""
        if not isinstance(campaign_id, str) or not campaign_id.strip():
            raise ValueError("campaign_id must be a non-empty string")
        return self._ensure_event_bus().list_events(campaign_id=campaign_id.strip())

    def list_campaign_states(
        self,
        *,
        topic: str | None = None,
        status: str | None = None,
    ) -> list[CampaignState]:
        """List computed state for every persisted campaign known to the event bus."""
        records = self._ensure_event_bus().list_events(topic=topic, status=status)
        campaigns: dict[str, list[EventRecord]] = {}
        for record in records:
            if not isinstance(record.campaign_id, str) or not record.campaign_id.strip():
                continue
            campaigns.setdefault(record.campaign_id.strip(), []).append(record)
        states = [
            self._campaign_state_from_records(campaign_id, campaign_records)
            for campaign_id, campaign_records in campaigns.items()
        ]
        return sorted(
            states,
            key=lambda item: (
                item.last_updated_at or "",
                item.campaign_id,
            ),
            reverse=True,
        )

    def campaign_state(self, campaign_id: str) -> CampaignState:
        """Compute campaign state from persisted event bus records."""
        if not isinstance(campaign_id, str) or not campaign_id.strip():
            raise ValueError("campaign_id must be a non-empty string")
        normalized_campaign_id = campaign_id.strip()
        return self._campaign_state_from_records(
            normalized_campaign_id,
            self.list_campaign_events(normalized_campaign_id),
        )

    def _campaign_state_from_records(
        self,
        campaign_id: str,
        records: list[EventRecord],
    ) -> CampaignState:
        """Aggregate campaign counters from an already-loaded event list."""
        ready_events = 0
        errored_events = 0
        next_due_at: str | None = None
        last_updated_at: str | None = None
        now = datetime.now(UTC)
        for record in records:
            if record.last_error:
                errored_events += 1
            if record.status == "queued":
                available_at = datetime.fromisoformat(record.available_at)
                if available_at <= now:
                    ready_events += 1
                if next_due_at is None or record.available_at < next_due_at:
                    next_due_at = record.available_at
            if last_updated_at is None or record.updated_at > last_updated_at:
                last_updated_at = record.updated_at
        return CampaignState(
            campaign_id=campaign_id,
            total_events=len(records),
            queued_events=sum(1 for record in records if record.status == "queued"),
            ready_events=ready_events,
            claimed_events=sum(1 for record in records if record.status == "claimed"),
            completed_events=sum(1 for record in records if record.status == "completed"),
            errored_events=errored_events,
            next_due_at=next_due_at,
            last_updated_at=last_updated_at,
        )

    # -- event round dispatch ------------------------------------------------

    async def _run_event_round(self) -> tuple[Response | None, int]:
        if self._event_bus is None or not self._topic_subscriptions:
            return None, 0
        processed = 0
        last_response: Response | None = None
        for subscription in self._topic_subscriptions:
            event = self._event_bus.claim(
                self._event_consumer_name(subscription),
                [subscription.topic],
            )
            if event is None:
                continue
            processed += 1
            response = await self._dispatch_event_subscription(subscription, event)
            if response is not None:
                last_response = response
        return last_response, processed

    async def _dispatch_event_subscription(
        self,
        subscription: _TeamEventSubscription,
        event: EventRecord,
    ) -> Response | None:
        consumer = self._event_consumer_name(subscription)
        try:
            response, response_payload = await self._invoke_event_target(subscription.target, event)
            if subscription.publish_to is not None:
                self._validate_topic_access(subscription.agent_name, subscription.publish_to, action="publish")
                self._ensure_event_bus().publish(
                    subscription.publish_to,
                    response_payload,
                    campaign_id=event.campaign_id,
                    correlation_id=event.correlation_id,
                    run_id=event.run_id,
                    publisher=subscription.agent_name,
                )
            self._ensure_event_bus().ack(event.id, consumer)
            return response
        except Exception as exc:
            self._ensure_event_bus().requeue(
                event.id,
                consumer,
                error=f"{type(exc).__name__}: {exc}",
            )
            raise

    async def _invoke_event_target(
        self,
        target: Any,
        event: EventRecord,
    ) -> tuple[Response | None, JSONObject]:
        if isinstance(target, Runnable):
            prompt = self._event_prompt_from_payload(event.payload)
            response = await target.run(
                prompt,
                session_id=event.correlation_id,
                correlation_id=event.correlation_id,
            )
            return response, self._response_to_event_payload(response)
        result = target(event)
        if inspect.isawaitable(result):
            result = await result
        if isinstance(result, Response):
            return result, self._response_to_event_payload(result)
        payload = self._coerce_event_payload(result)
        response = Response(
            content=payload.get("content", json.dumps(payload, ensure_ascii=False)),
            tool_calls=[],
            usage=Usage(0, 0),
            raw={"event_payload": dict(payload)},
        )
        return response, payload

    # -- internal helpers ----------------------------------------------------

    def _ensure_event_bus(self) -> EventBus:
        if self._event_bus is None:
            self._event_bus = InMemoryEventBus()
        return self._event_bus

    @staticmethod
    def _build_event_bus_from_config(config: EventBusConfig | None) -> EventBus | None:
        if config is None:
            return None
        if config.backend == "memory":
            return InMemoryEventBus()
        if config.backend == "sqlite":
            from logos_adk.event_bus import SQLiteEventBus

            return SQLiteEventBus(config.path or ".logos_adk/event_bus.db")
        if config.backend == "postgresql":
            from logos_adk.event_bus import PostgreSQLEventBus

            return PostgreSQLEventBus(config.dsn)
        raise ValueError(f"Unsupported event bus backend '{config.backend}'")

    def _resolve_topic_policy(
        self,
        topic_policy: dict[str, list[str]] | None,
    ) -> dict[str, set[str]]:
        if topic_policy is None:
            return {name: {"*"} for name in self._agents}
        policy: dict[str, set[str]] = {}
        known_agents = set(self._agents)
        for agent_name, topics in topic_policy.items():
            if agent_name not in known_agents:
                raise ValueError(
                    f"Agent '{agent_name}' in topic_policy is not in team '{self.name}': {sorted(known_agents)}"
                )
            policy[agent_name] = {topic.strip() for topic in topics if isinstance(topic, str) and topic.strip()}
        for agent_name in self._agents:
            policy.setdefault(agent_name, set())
        return policy

    def _validate_topic_access(self, agent_name: str, topic: str, *, action: str) -> None:
        if agent_name not in self._agents:
            raise RoutingError(f"Agent '{agent_name}' is not registered in team '{self.name}'")
        if self._is_internal_topic(topic):
            return
        allowed_topics = self._topic_policy.get(agent_name, set())
        if "*" in allowed_topics or topic in allowed_topics:
            return
        raise RoutingError(
            f"Agent '{agent_name}' cannot {action} topic '{topic}'. "
            f"Allowed topics: {sorted(allowed_topics)}"
        )

    def _resolve_subscription_handler(self, agent: AgentLike, handler: Any) -> Any:
        if isinstance(handler, str):
            if not hasattr(agent, handler):
                registered = subscription_handler_registry.get(handler)
                if registered is not None:
                    return registered
                imported = self._load_subscription_handler_from_path(handler)
                if imported is not None:
                    return imported
                raise ValueError(
                    f"Agent '{agent.name}' has no handler '{handler}', and it is not registered in subscription_handler_registry"
                )
            return getattr(agent, handler)
        return handler

    def _load_subscription_handler_from_path(self, path: str) -> Any | None:
        if ":" in path:
            module_name, attr_name = path.split(":", 1)
        elif "." in path:
            module_name, attr_name = path.rsplit(".", 1)
        else:
            return None
        if not module_name or not attr_name:
            return None
        try:
            module = importlib.import_module(module_name)
        except ImportError:
            return None
        except Exception:
            logging.getLogger(__name__).debug(
                "Failed to load subscription handler from '%s'",
                path,
                exc_info=True,
            )
            return None
        candidate = getattr(module, attr_name, None)
        return candidate if callable(candidate) else None

    def _register_agent_topic_subscriptions(self, agent: AgentLike) -> None:
        pending = list(getattr(agent, "_pending_topic_subscriptions", []))
        if not pending:
            return
        for item in pending:
            self.subscribe(
                agent.name,
                item["topic"],
                handler=item.get("handler"),
                publish_to=item.get("publish_to"),
                name=item.get("name"),
            )
        agent._pending_topic_subscriptions.clear()

    def _register_builtin_scheduled_task_subscriptions(self) -> None:
        for agent_name in self._agents:
            topic = self._scheduled_task_topic(agent_name)
            if any(
                item.agent_name == agent_name and item.topic == topic
                for item in self._topic_subscriptions
            ):
                continue
            self._topic_subscriptions.append(
                _TeamEventSubscription(
                    agent_name=agent_name,
                    topic=topic,
                    handler_name="scheduled_task",
                    target=self._make_scheduled_task_handler(agent_name),
                )
            )

    def _scheduled_task_topic(self, agent_name: str) -> str:
        return f"{_SCHEDULED_TASK_TOPIC_PREFIX}{agent_name}"

    def _is_internal_topic(self, topic: str) -> bool:
        return topic.startswith(_SCHEDULED_TASK_TOPIC_PREFIX)

    def _make_scheduled_task_handler(self, agent_name: str) -> Any:
        async def _run_scheduled_task(event: EventRecord) -> Response:
            payload = event.payload
            task = payload.get("task")
            if not isinstance(task, str) or not task.strip():
                raise ValueError("scheduled task payload requires non-empty 'task'")
            target_agent_name = payload.get("target_agent")
            if target_agent_name != agent_name:
                raise ValueError(
                    f"Scheduled task for '{target_agent_name}' claimed by '{agent_name}'"
                )
            session_id = payload.get("session_id")
            if session_id is not None and not isinstance(session_id, str):
                session_id = None
            target_agent = self._agents[agent_name]
            response = await target_agent.run(
                task,
                session_id=session_id or event.correlation_id,
                correlation_id=event.correlation_id,
            )
            raw = dict(response.raw or {})
            raw["scheduled_task"] = {
                "source_agent": payload.get("source_agent"),
                "target_agent": agent_name,
                "campaign_id": event.campaign_id,
                "event_id": event.id,
            }
            return Response(
                content=response.content,
                content_parts=response.content_parts,
                tool_calls=response.tool_calls,
                usage=response.usage,
                raw=raw,
                reasoning=response.reasoning,
                status=response.status,
            )

        return _run_scheduled_task

    def _event_consumer_name(self, subscription: _TeamEventSubscription) -> str:
        return f"{self.name}:{subscription.agent_name}:{subscription.handler_name}"

    def _scheduled_delay_seconds(
        self,
        *,
        delay_seconds: float | None,
        due_at: datetime | None,
    ) -> float:
        if due_at is not None:
            if due_at.tzinfo is None:
                due_at = due_at.replace(tzinfo=UTC)
            return max(0.0, (due_at.astimezone(UTC) - datetime.now(UTC)).total_seconds())
        return max(0.0, float(delay_seconds or 0.0))

    def _event_prompt_from_payload(self, payload: JSONObject) -> str:
        prompt = payload.get("prompt")
        if isinstance(prompt, str) and prompt:
            return prompt
        return json.dumps(payload, ensure_ascii=False)

    def _response_to_event_payload(self, response: Response) -> JSONObject:
        payload: dict[str, Any] = {"content": response.content}
        if response.raw:
            payload["raw"] = dict(response.raw)
        return ensure_json_object(payload, path="team_event.response")

    def _coerce_event_payload(self, value: Any) -> JSONObject:
        if value is None:
            return {"status": "ok"}
        if isinstance(value, dict):
            return ensure_json_object(value, path="team_event.payload")
        return {"output": str(value)}
