import asyncio
import json
import os
from contextlib import suppress
from collections.abc import Mapping
from typing import Any, Optional, Type, Sequence, List

import pandas as pd
from fastapi import APIRouter, Body, HTTPException, Query, Request
from fastapi.encoders import jsonable_encoder
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, Field, create_model

from sibi_flux.dask_cluster import safe_compute
from sibi_flux.dask_cluster.client_manager import get_persistent_client
from sibi_flux.datacube._data_cube import Datacube


class BaseCubeRouter:
    """
    A reusable router class for Datacubes that provides:
    - POST /: Data retrieval with dynamic filters (optional pagination).
    - POST /stream: SSE streaming with dynamic filters.

    Attributes:
        cube_cls: The Datacube class to wrap.
        router: The APIRouter instance containing the endpoints.
        FiltersModel: Dynamic Pydantic model for validation.
        pagination: Whether to enable pagination logic on the main endpoint.
    """

    def __init__(
        self,
        cube_cls: Type[Datacube],
        prefix: str = "",
        tags: Sequence[str] | None = None,
        pagination: bool = True,
    ):
        self.cube_cls = cube_cls
        self.pagination = pagination
        self.router = APIRouter(prefix=prefix, tags=tags or [])
        self.FiltersModel = self._create_filters_model()
        self._setup_routes()

    def _generate_docstring_and_example(self, field_defs: Mapping[str, Any]) -> str:
        """
        Generates a JSON-like representation of the model for the docstring.
        """
        example_dict = {}
        for name, (type_, field) in field_defs.items():
            # Simplistic type mapping for display
            type_name = "string"
            if type_ is Optional[bool]:
                type_name = True
            elif type_ is Optional[int]:
                type_name = 0
            elif type_ is Optional[float]:
                type_name = 0.0

            example_dict[name] = type_name

        pretty_json = json.dumps(example_dict, indent=2)
        return f"Filter options:\n\n```json\n{pretty_json}\n```"

    def _create_filters_model(self) -> Type[BaseModel]:
        # Default behavior: if exclude_columns is empty/missing, exclude 'id'.
        _exclude_columns = self.cube_cls.config.get("exclude_columns", [])
        if not _exclude_columns:
            _exclude_columns = ["id"]

        try:
            # Check if introspection is allowed via env var (default: true)
            if os.environ.get("SIBI_FLUX_ROUTER_INTROSPECTION", "true").lower() == "false":
                return self._fallback_create_model(_exclude_columns)

            # Introspect dtypes by loading an empty slice/schema
            # We use n_records=1 at init to hint the backend to limit query size
            # Note: This requires the DB/Backend to be accessible at startup
            cube = self.cube_cls(n_records=1)

            # Try to load just metadata/empty df
            # If backend doesn't support limit, this might raise Error or load full data.
            # We catch exceptions to fallback.
            df = cube.load()

            if df is None:
                return self._fallback_create_model(_exclude_columns)

            _field_definitions = {}

            # Identify boolean fields from config as override/hint
            _boolean_fields = getattr(self.cube_cls, "boolean_fields", [])
            _field_map = self.cube_cls.config.get("field_map", {})
            _target_columns = set(_field_map.values()) if _field_map else None

            for col in df.columns:
                if col in _exclude_columns:
                    continue

                # If field_map is provided, use only targeted columns
                if _target_columns is not None and col not in _target_columns:
                    continue

                if col not in df.columns:
                    continue

                dtype = df.dtypes[col]
                desc = f"Filter by {col}"
                field_type: Any = Optional[str]  # Default

                if col in _boolean_fields:
                    field_type = Optional[bool]
                    desc += " (bool)"
                elif pd.api.types.is_bool_dtype(dtype):
                    field_type = Optional[bool]
                    desc += " (bool)"
                elif pd.api.types.is_integer_dtype(dtype):
                    field_type = Optional[int]
                    desc += " (int)"
                elif pd.api.types.is_float_dtype(dtype):
                    field_type = Optional[float]
                    desc += " (float)"

                _field_definitions[col] = (
                    field_type,
                    Field(default=None, description=desc),
                )

            model = create_model(
                f"{self.cube_cls.__name__}Filters",
                **_field_definitions,
                __config__=ConfigDict(extra="ignore"),
            )
            model.__doc__ = self._generate_docstring_and_example(_field_definitions)
            return model

        except Exception:
            # Fallback to config-based generation
            # print(f"DTO INTROSPECTION ERROR for {self.cube_cls.__name__}: {e}")
            return self._fallback_create_model(_exclude_columns)

    def _fallback_create_model(self, exclude_columns: Sequence[str]) -> Type[BaseModel]:
        # Retrieve fields from the Cube configuration
        _field_map = self.cube_cls.config.get("field_map", {})
        _output_columns = list(_field_map.values()) if _field_map else []
        _boolean_fields = getattr(self.cube_cls, "boolean_fields", [])

        _field_definitions = {}
        for col in _output_columns:
            if col in exclude_columns:
                continue

            if col in _boolean_fields:
                _field_definitions[col] = (
                    Optional[bool],
                    Field(default=None, description=f"Filter by {col} (bool)"),
                )
            else:
                _field_definitions[col] = (
                    Optional[str | int],
                    Field(default=None, description=f"Filter by {col}"),
                )

        FiltersModel: Any = create_model(
            f"{self.cube_cls.__name__}Filters",
            **_field_definitions,
            __config__=ConfigDict(extra="ignore"),
        )
        FiltersModel.__doc__ = self._generate_docstring_and_example(_field_definitions)
        return FiltersModel

    async def _materialize_df(self, df: Any) -> pd.DataFrame:
        if hasattr(df, "compute"):
            # Use the shared persistent client if available
            client = get_persistent_client()
            return await asyncio.to_thread(safe_compute, df, dask_client=client)
        return df

    async def _stream_data_task(
        self, cube: Datacube, filters: Mapping[str, Any]
    ) -> None:
        async with cube:
            await cube.emit("status", msg="loading")
            try:
                df = await cube.aload(**filters)
                df = await self._materialize_df(df)

                if df is not None and not df.empty:
                    records = df.to_dict(orient="records")
                    await cube.emit("data", data=jsonable_encoder(records))

                await cube.emit("status", msg="complete")

            except Exception as e:
                # Log the actual error for internal tracking (would use a real logger in production)
                print(f"DATACUBE STREAM ERROR: {str(e)}")
                # Return a sanitized error message
                await cube.emit("error", detail="An internal error occurred during data streaming.")

    def _setup_routes(self):
        # We define the functions inside closure to capture 'self' reliably
        # and use the dynamic FiltersModel for type annotation and FastAPI dependency.

        FiltersModel = self.FiltersModel

        async def _core_logic(cube, filters):
            if isinstance(filters, dict):
                filters = FiltersModel(**filters)
            filter_dict = filters.model_dump(exclude_unset=True)
            try:
                df = await cube.aload(**filter_dict)
            except Exception as e:
                # Log the actual error
                print(f"DATACUBE LOAD ERROR: {str(e)}")
                # Return sanitized error to client
                raise HTTPException(
                    status_code=500, 
                    detail="An internal error occurred while loading datacube data."
                )
            df = await self._materialize_df(df)
            return df

        if self.pagination:

            @self.router.post(
                "",
                summary=f"Get Paged {self.cube_cls.__name__} Data",
                description="Retrieve data with filtering and pagination.",
            )
            async def get_cube_data(
                filters: Any = Body(
                    default_factory=lambda: FiltersModel(),
                    description="Filters to apply.",
                    openapi_examples={},
                ),
                page: int = Query(1, ge=1, description="Page number."),
                page_size: int = Query(
                    50, ge=1, le=1000, description="Records per page."
                ),
            ) -> List[dict[str, Any]]:
                cube = self.cube_cls()
                df = await _core_logic(cube, filters)

                if df is None or df.empty:
                    return []

                start = (page - 1) * page_size
                end = start + page_size
                paged_df = df.iloc[start:end]
                return paged_df.to_dict(orient="records")

        else:

            @self.router.post(
                "",
                summary=f"Get {self.cube_cls.__name__} Data",
                description="Retrieve all data with filtering (pagination disabled).",
            )
            async def get_cube_data_all(
                filters: Any = Body(
                    default_factory=lambda: FiltersModel(),
                    description="Filters to apply.",
                    openapi_examples={},
                ),
            ) -> List[dict[str, Any]]:
                cube = self.cube_cls()
                df = await _core_logic(cube, filters)

                if df is None or df.empty:
                    return []

                return df.to_dict(orient="records")

        @self.router.post(
            "/stream",
            summary=f"Stream {self.cube_cls.__name__} Data",
            description="Stream data updates via SSE.",
        )
        async def stream_cube_data(
            request: Request,
            filters: Any = Body(
                default_factory=lambda: FiltersModel(),
                description="Filters to apply.",
                openapi_examples={},
            ),
        ) -> StreamingResponse:
            cube = self.cube_cls(auto_sse=True)
            if isinstance(filters, dict):
                filters = FiltersModel(**filters)
            filter_dict = filters.model_dump(exclude_unset=True)

            asyncio.create_task(self._stream_data_task(cube, filter_dict))

            return StreamingResponse(
                cube.get_sse().aiter_sse(), media_type="text/event-stream"
            )
