"""Streaming generator for Prism.

Generates the event bus, SSE endpoints, and WebSocket handlers
for real-time streaming of model CRUD events.
"""

from __future__ import annotations

from pathlib import Path

from prisme.generators.base import GeneratedFile, GeneratorBase, create_init_file
from prisme.spec.stack import FileStrategy
from prisme.utils.template_engine import TemplateRenderer


class StreamingGenerator(GeneratorBase):
    """Generator for streaming infrastructure (event bus, SSE, WebSocket)."""

    REQUIRED_TEMPLATES = [
        "backend/streaming/event_bus.py.jinja2",
    ]

    def __init__(self, *args, **kwargs) -> None:  # type: ignore[no-untyped-def]
        super().__init__(*args, **kwargs)
        backend_base = Path(self.generator_config.backend_output)
        package_name = self.get_package_name()
        package_base = backend_base / package_name
        self.streaming_path = package_base / self.generator_config.streaming_path
        self.generated_path = package_base / self.generator_config.streaming_generated_path
        self.renderer = TemplateRenderer()
        self.renderer.validate_templates_exist(self.REQUIRED_TEMPLATES)

    def generate_files(self) -> list[GeneratedFile]:
        """Generate all streaming infrastructure files."""
        files: list[GeneratedFile] = []

        # Collect all custom events across models
        custom_events = self._collect_custom_events()

        # Always generate the event bus
        files.append(self._generate_event_bus(custom_events))

        # Get enabled models for per-model endpoints
        model_names = [m.name for m in self.spec.models if m.expose]

        # Generate SSE handler and routes if SSE is enabled
        sse_enabled = (
            self.exposure_config
            and getattr(self.exposure_config, "sse", None)
            and self.exposure_config.sse.enabled
        )
        if sse_enabled:
            sse_path = self.exposure_config.sse.path
            keepalive = 30
            files.append(self._generate_sse_handler(keepalive))
            files.append(self._generate_sse_routes(model_names, sse_path))

        # Generate WebSocket handler and routes if WebSocket is enabled
        ws_enabled = (
            self.exposure_config
            and getattr(self.exposure_config, "websocket", None)
            and self.exposure_config.websocket.enabled
        )
        if ws_enabled:
            ws_path = self.exposure_config.websocket.path
            heartbeat = 30
            files.append(self._generate_ws_handler(heartbeat))
            files.append(self._generate_ws_routes(model_names, ws_path))

        # Generate __init__ files
        files.append(self._generate_generated_init(bool(sse_enabled), bool(ws_enabled)))
        files.append(self._generate_streaming_init(bool(sse_enabled), bool(ws_enabled)))

        return files

    def _collect_custom_events(self) -> list[str]:
        """Collect all custom events defined across models."""
        custom = set()
        # Get project-level streaming config
        if (
            self.exposure_config
            and getattr(self.exposure_config, "streaming", None)
            and self.exposure_config.streaming.enabled
        ):
            # No custom_events on StreamingBusConfig - those are per-model
            pass

        # Collect per-model custom events
        for model in self.spec.models:
            model_custom = model.get_delivery_override("streaming_custom_events")
            if model_custom:
                custom.update(model_custom)

        return sorted(custom)

    def _generate_event_bus(self, custom_events: list[str]) -> GeneratedFile:
        """Generate the central event bus module."""
        content = self.renderer.render_file(
            "backend/streaming/event_bus.py.jinja2",
            context={
                "custom_events": custom_events,
            },
        )
        return GeneratedFile(
            path=self.generated_path / "event_bus.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="Streaming event bus",
        )

    def _generate_sse_handler(self, keepalive_interval: int) -> GeneratedFile:
        """Generate the SSE handler module."""
        content = self.renderer.render_file(
            "backend/streaming/sse_handler.py.jinja2",
            context={
                "keepalive_interval": keepalive_interval,
            },
        )
        return GeneratedFile(
            path=self.streaming_path / "sse_handler.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="SSE event handler",
        )

    def _generate_sse_routes(self, model_names: list[str], sse_path: str) -> GeneratedFile:
        """Generate the SSE routes module."""
        content = self.renderer.render_file(
            "backend/streaming/sse_routes.py.jinja2",
            context={
                "model_names": model_names,
                "sse_path": sse_path,
            },
        )
        return GeneratedFile(
            path=self.streaming_path / "sse_routes.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="SSE streaming routes",
        )

    def _generate_ws_handler(self, heartbeat_interval: int) -> GeneratedFile:
        """Generate the WebSocket handler module."""
        content = self.renderer.render_file(
            "backend/streaming/ws_handler.py.jinja2",
            context={
                "heartbeat_interval": heartbeat_interval,
            },
        )
        return GeneratedFile(
            path=self.streaming_path / "ws_handler.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="WebSocket event handler",
        )

    def _generate_ws_routes(self, model_names: list[str], ws_path: str) -> GeneratedFile:
        """Generate the WebSocket routes module."""
        content = self.renderer.render_file(
            "backend/streaming/ws_routes.py.jinja2",
            context={
                "model_names": model_names,
                "ws_path": ws_path,
            },
        )
        return GeneratedFile(
            path=self.streaming_path / "ws_routes.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="WebSocket streaming routes",
        )

    def _generate_generated_init(self, sse_enabled: bool, ws_enabled: bool) -> GeneratedFile:
        """Generate __init__.py for _generated folder."""
        imports = ["from .event_bus import EventBus, EventType, StreamEvent, event_bus"]
        exports = ["EventBus", "EventType", "StreamEvent", "event_bus"]
        return create_init_file(
            self.generated_path,
            imports,
            exports,
            "Generated streaming infrastructure.",
        )

    def _generate_streaming_init(self, sse_enabled: bool, ws_enabled: bool) -> GeneratedFile:
        """Generate __init__.py for streaming folder."""
        imports = [
            "from ._generated.event_bus import EventBus, EventType, StreamEvent, event_bus",
        ]
        exports = ["EventBus", "EventType", "StreamEvent", "event_bus"]

        if sse_enabled:
            imports.append("from .sse_handler import create_sse_response")
            imports.append("from .sse_routes import router as sse_router")
            exports.extend(["create_sse_response", "sse_router"])

        if ws_enabled:
            imports.append("from .ws_handler import handle_websocket")
            imports.append("from .ws_routes import router as ws_router")
            exports.extend(["handle_websocket", "ws_router"])

        return create_init_file(
            self.streaming_path,
            imports,
            exports,
            "Streaming infrastructure for real-time updates.",
        )


__all__ = ["StreamingGenerator"]
