from typing import Any, Dict, Optional, Type, Generic, TypeVar, Callable, Awaitable
from pspf.connectors.base import StreamingBackend
from pspf.schema import BaseEvent, SchemaRegistry
from pspf.processor import BatchProcessor
from pspf.settings import settings
from pspf.telemetry import TelemetryManager
from pydantic import BaseModel
from pspf.utils.logging import get_logger

logger = get_logger("Stream")

T = TypeVar("T", bound=BaseModel)

class Stream(Generic[T]):
    """
    High-level Facade for Stream Processing using Composition.

    Manages the lifecycle of a stream processor, including connection handling,
    schema validation, and the processing loop.

    Attributes:
        backend (StreamingBackend): The storage backend for stream operations.
        schema (Optional[Type[T]]): Pydantic model for data validation.
        processor (BatchProcessor): Internal processor for batch handling.
        telemetry (TelemetryManager): Observability manager.
    """
    def __init__(self, 
                 backend: StreamingBackend,
                 schema: Optional[Type[T]] = None):
        """
        Initialize the Stream facade.

        Args:
            backend (ValkeyStreamBackend): Configured backend instance.
            schema (Optional[Type[T]]): Pydantic model class for validation. 
                                        If None, uses dynamic SchemaRegistry.
        """
        self.backend = backend
        self.schema = schema
        self.processor = BatchProcessor(backend)
        self.telemetry = TelemetryManager()
        
    async def health(self) -> Dict[str, str]:
        """
        Perform a health check on the stream components.

        Pings the backend to ensure connectivity.

        Returns:
            Dict[str, str]: A dictionary containing status ('ok' or 'error') 
                            and details.
        """
        try:
            # Assuming backend has a ping method (we will add it)
            await self.backend.ping()
            # Could also check if group exists
            return {"status": "ok", "backend": "connected"}
        except Exception as e:
            logger.error(f"Health check failed: {e}")
            return {"status": "error", "reason": str(e)}

    async def __aenter__(self) -> "Stream[T]":
        """
        Async context manager entry.
        
        Connects to the backend and ensures the consumer group exists.
        """
        # We delegate connection management to the backend's connector 
        # if the user hasn't already connected it.
        # But logically, the backend should be ready. 
        # We will ensure the group exists here.
        await self.backend.connect() # Idempotent-ish check inside
        await self.backend.ensure_group_exists()
        return self

    async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None:
        """
        Async context manager exit.
        
        Trigger shutdown of the processor and closes the backend connection.
        """
        if self.processor:
            await self.processor.shutdown()
        # We close the connection here as we "took ownership" in aenter
        await self.backend.close()

    async def emit(self, event: BaseModel) -> str:
        """
        Produce an event to the stream.

        Serializes the Pydantic model and injects tracing context before 
        sending to the backend.

        Args:
            event (BaseModel): The event object to publish.

        Returns:
            str: The message ID generated by the backend (e.g., "1678888888000-0").
        """
        data = event.model_dump(mode='json')
        # Ensure event_type is present for deserialization
        if "event_type" not in data:
            data["event_type"] = event.__class__.__name__

        # Inject Trace Context
        # We add a hidden field to carry the trace context
        self.telemetry.inject_context(data)
            
        msg_id = await self.backend.add_event(data)
        return msg_id

    async def run(self, handler: Callable[[T], Awaitable[None]], batch_size: int = 10) -> None:
        """
        Start the infinite processing loop.

        Continuously fetches batches of messages, validates them against the schema,
        and invokes the user-provided handler.

        Args:
            handler (Callable[[T], Awaitable[None]]): Async function to process each event.
            batch_size (int): Number of messages to fetch in each batch. Default is 10.

        Raises:
             Exception: Any unhandled exception from the processor or validation logic.
        """
        # Register the user's handler with the processor
        # We wrap it to handle deserialization
        
        async def typed_handler(msg_id: str, raw_data: Dict[str, Any]) -> None:
            # 1. Deserialize / Validate
            if self.schema:
                try:
                    event = self.schema.model_validate(raw_data)
                except Exception as e:
                    logger.error(f"Schema validation failed for msg {msg_id}: {e}")
                    # Re-raise so processor handles DLO logic
                    raise e
            else:
                # Dynamic validation via Registry or fallback
                # Note: raw_data might need cleanup if it has Redis artifacts? 
                # (Valkey decode_responses=True handles bytes->str)
                try:
                    # cast is needed because validate returns BaseModel but we expect T
                    from typing import cast
                    event = cast(T, SchemaRegistry.validate(raw_data))
                except Exception as e:
                    logger.error(f"Dynamic validation failed for msg {msg_id}: {e}")
                    raise e
                
            # Inject metadata
            if isinstance(event, BaseEvent):
                event.offset = msg_id
            
            # 2. Call User Logic
            # We explicitly cast event to T (or close enough)
            await handler(event) # type: ignore

        logger.info(f"Starting stream processor for {self.backend.stream_key}...")
        await self.processor.run_loop(typed_handler, batch_size=batch_size)
