import inspect
from abc import ABC, abstractmethod
from typing import Any, Callable, Type, Union

from fastapi import Depends, FastAPI, Request
from pydantic import BaseModel, create_model


class SelectEngine(ABC):

    @abstractmethod
    def select(self, query: str, *params: Any) -> list[Any]:
        pass

    @abstractmethod
    def select_count(self, query: str, *params: Any) -> int:
        pass


class QueryContext(BaseModel):
    """
    Context object passed to middlewares, containing all query-related data
    that can be modified by middleware.
    """

    select_query: str
    select_params: list[Any] | None = None
    count_query: str | None = None
    count_params: list[Any] | None = None
    filters: BaseModel
    limit: int
    offset: int
    dependency_results: dict[str, Any]


# Middleware type that receives QueryContext and next middleware callable
# and returns the ResponseModel (data_model)
Middleware = Callable[[QueryContext, Callable[[QueryContext], Any]], Any]


def serve_select(
    app: FastAPI,
    query_engine: SelectEngine,
    select_query: str,
    select_param: Callable[[BaseModel, int, int], list[Any]] | None = None,
    count_query: str | None = None,
    count_param: Callable[[BaseModel], list[Any]] | None = None,
    path: str = "/get",
    filter_model: Type[BaseModel] | None = None,
    data_model: Type[BaseModel] | None = None,
    dependencies: list[Any] | None = None,
    middlewares: list[Middleware] | None = None,
    summary: str | None = None,
    description: str | None = None
):
    """
    Adding fastapi route to app with proper annotation:
    - taking filter_model as parameter
    - return pydantic model with two property:
        - data (data_model)
        - count (int)
    - supports FastAPI dependencies
    - supports middleware system
    """
    if filter_model is None:
        filter_model = create_model("FilterModel")
    wrapped_dependencies = _create_wrapped_fastapi_dependencies(dependencies)
    SelectResponseModel = _create_select_response_model(data_model)

    # Create the route with dependencies
    @app.get(
        path,
        response_model=SelectResponseModel,
        dependencies=wrapped_dependencies,
        summary=summary,
        description=description
    )
    async def get(
        request: Request,
        filters: filter_model = Depends(),  # type: ignore
        limit: int = 100,
        offset: int = 0,
    ):
        """
        This route is automatically generated by dbanu
        """
        # Extract dependency results from request state
        dependency_results = {}
        if request and hasattr(request.state, "dependency_results"):
            dependency_results = request.state.dependency_results
        # Build initial select parameters
        select_args = (
            select_param(filters, limit, offset)
            if select_param is not None
            else [limit, offset]
        )
        # Build initial count parameters
        count_args = count_param(filters) if count_param is not None else []
        # Create initial QueryContext
        initial_context = QueryContext(
            select_query=select_query,
            select_params=select_args,
            count_query=count_query,
            count_params=count_args,
            filters=filters,
            limit=limit,
            offset=offset,
            dependency_results=dependency_results,
        )
        query_processor = _create_query_processor(query_engine, SelectResponseModel)
        handler = _create_middleware_chain(middlewares, query_processor)
        return await handler(initial_context)


class SelectSource(BaseModel):
    model_config = {"arbitrary_types_allowed": True}
    query_engine: SelectEngine
    select_query: str
    select_param: Callable[[BaseModel, int, int], list[Any]] | None = None
    count_query: str
    count_param: Callable[[BaseModel], list[Any]] | None = None
    middlewares: list[Middleware] | None = None


def serve_union(
    app: FastAPI,
    sources: dict[str, SelectSource],
    path: str = "/get",
    filter_model: Type[BaseModel] | None = None,
    data_model: Type[BaseModel] | None = None,
    dependencies: list[Any] | None = None,
    middlewares: list[Middleware] | None = None,
    summary: str | None = None,
    description: str | None = None
):

    if filter_model is None:
        filter_model = create_model("FilterModel")
    wrapped_dependencies = _create_wrapped_fastapi_dependencies(dependencies)
    SelectResponseModel = _create_select_response_model(data_model)

    # Create the route with dependencies
    @app.get(
        path,
        response_model=SelectResponseModel,
        dependencies=wrapped_dependencies,
        summary=summary,
        description=description
    )
    async def get(
        request: Request,
        filters: filter_model = Depends(),  # type: ignore
        limit: int = 100,
        offset: int = 0,
        priority: str | None = None
    ):
        """
        This route is automatically generated by dbanu
        """
        # Extract dependency results from request state
        dependency_results = {}
        if request and hasattr(request.state, "dependency_results"):
            dependency_results = request.state.dependency_results
        
        priority_list = list(sources.keys()) if priority is None else priority.split(",")
        
        # Step 1: Get total count from each source
        source_counts = {}
        total_count = 0
        for source_name in priority_list:
            source = sources[source_name]
            count_context = QueryContext(
                select_query="",
                select_params=[],
                count_query=source.count_query,
                count_params=source.count_param,
                filters=filters,
                limit=0,
                offset=0,
                dependency_results=dependency_results,
            )
            count_processor = _create_count_processor(source.query_engine)
            handler = _create_middleware_chain(source.middlewares, count_processor)
            source_count = await handler(count_context)
            source_counts[source_name] = source_count
            total_count += source_count
        
        # Step 2: Calculate which records to fetch from each source
        final_data = []
        remaining_limit = limit
        current_offset = offset
        for source_name in priority_list:
            source = sources[source_name]
            source_count = source_counts[source_name]
            # Skip this source if offset is beyond its records
            if current_offset >= source_count:
                current_offset -= source_count
                continue
            # Calculate how many records to fetch from this source
            source_limit = min(remaining_limit, source_count - current_offset)
            if source_limit > 0:
                # Build select parameters for this specific source
                select_args = (
                    source.select_param(filters, source_limit, current_offset)
                    if source.select_param is not None
                    else [source_limit, current_offset]
                )
                # Create QueryContext for this source
                select_context = QueryContext(
                    select_query=source.select_query,
                    select_params=select_args,
                    count_query="",
                    count_params=[],
                    filters=filters,
                    limit=source_limit,
                    offset=current_offset,
                    dependency_results=dependency_results,
                )
                select_processor = _create_select_processor(source.query_engine)
                handler = _create_middleware_chain(
                    _get_combined_middlewares(middlewares, source.middlewares),
                    select_processor
                )
                source_data = await handler(select_context)
                # Add the data to final result
                final_data.extend(source_data)
                remaining_limit -= len(source_data)
                # Reset offset for next source
                current_offset = 0
            # Stop if we've reached the limit
            if remaining_limit <= 0:
                break
        return SelectResponseModel(data=final_data, count=total_count)


def _create_select_response_model(data_model: Type[BaseModel] | None = None):
    actual_data_model = create_model("DataModel") if data_model is None else data_model

    class SelectResponseModel(BaseModel):
        data: list[actual_data_model] | list[Any]  # type: ignore
        count: int

    return SelectResponseModel


def _create_select_processor(query_engine: SelectEngine):

    async def process_select(context: QueryContext) -> list[Any]:
        select_params = context.select_params or []
        # Check if select method is a coroutine
        select_method = query_engine.select
        if inspect.iscoroutinefunction(select_method):
            return await select_method(context.select_query, *select_params)
        return select_method(context.select_query, *select_params)

    return process_select


def _create_count_processor(query_engine: SelectEngine):

    async def process_count(context: QueryContext) -> int:
        count_params = context.count_params or []
        if context.count_query is None:
            raise ValueError(f"select_count is not defined at this context: {context}")
        select_count_method = query_engine.select_count
        if inspect.iscoroutinefunction(select_count_method):
            return await select_count_method(context.count_query, *count_params)
        return select_count_method(context.count_query, *count_params)
    
    return process_count


def _create_query_processor(query_engine: SelectEngine, response_model: type[BaseModel]):

    async def process_query(context: QueryContext):
        select_params = context.select_params or []
        # Check if select method is a coroutine
        select_method = query_engine.select
        if inspect.iscoroutinefunction(select_method):
            data = await select_method(context.select_query, *select_params)
        else:
            data = select_method(context.select_query, *select_params)
        if context.count_query:
            count_params = context.count_params or []
            # Check if select_count method is a coroutine
            select_count_method = query_engine.select_count
            if inspect.iscoroutinefunction(select_count_method):
                total = await select_count_method(context.count_query, *count_params)
            else:
                total = select_count_method(context.count_query, *count_params)
            return response_model(data=data, count=total)
        return response_model(data=data, count=len(data))

    return process_query


def _get_combined_middlewares(
    global_middlewares: list[Middleware] | None,
    local_middlewares: list[Middleware] | None,
) -> list[Middleware] | None:
    normalized_global_middlewares = global_middlewares or []
    normalized_local_middlewares = local_middlewares or []
    return list(normalized_global_middlewares) + list(normalized_local_middlewares)


def _create_middleware_chain(
    middlewares: list[Middleware] | None,
    final_handler: Callable[[QueryContext], Any]
):
    handler = final_handler
    if middlewares is None:
        middlewares = []
    for middleware in reversed(middlewares):
        handler = _make_middleware_wrapper(middleware, handler)
    return handler


def _make_middleware_wrapper(
    current_middleware: Middleware,
    next_handler: Callable[[QueryContext], Any]
):
    async def wrapper(context: QueryContext):
        # Check if the middleware is async
        if inspect.iscoroutinefunction(current_middleware):
            return await current_middleware(context, lambda ctx: next_handler(ctx))
        else:
            return current_middleware(context, lambda ctx: next_handler(ctx))
    return wrapper


def _create_wrapped_fastapi_dependencies(dependencies: list[Any] | None):
    if dependencies is None:
        return []
    wrapped_dependencies = []
    for dep in dependencies:
        if hasattr(dep, "dependency"):
            # Create a closure-safe wrapper for each dependency
            def create_wrapped_dependency(original_dep):
                async def wrapped_dependency(request: Request):
                    result = await original_dep.dependency
                    # Store the result in request state
                    if not hasattr(request.state, "dependency_results"):
                        request.state.dependency_results = {}
                    # Use the dependency function name as key
                    dep_name = original_dep.dependency.__name__
                    request.state.dependency_results[dep_name] = result
                    return result
                return wrapped_dependency

            wrapped_dependencies.append(Depends(create_wrapped_dependency(dep)))
        else:
            wrapped_dependencies.append(dep)

    return wrapped_dependencies
