"""
serve_union implementation with priority-based pagination
"""

import inspect
from typing import Any, Callable, Type

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

from dbanu.api.dependencies import create_wrapped_fastapi_dependencies
from dbanu.core.engine import QueryContext, SelectEngine
from dbanu.core.middleware import create_middleware_chain, get_combined_middlewares
from dbanu.core.response import create_select_response_model
from dbanu.utils.pagination import calculate_union_pagination


class SelectSource(BaseModel):
    """Configuration for a single data source in union queries"""

    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[Any] | 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[Any] | None = None,
    summary: str | None = None,
    description: str | None = None,
):
    """
    Create a union endpoint that combines results from multiple sources
    with priority-based pagination
    """
    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 [
                source.strip() for source in priority.split(",") if source.strip() != ""
            ]
        )

        # 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
        fetch_plan = calculate_union_pagination(
            source_counts, priority_list, limit, offset
        )

        # Step 3: Fetch data according to the plan
        final_data = []
        for source_name, (source_limit, source_offset) in fetch_plan.items():
            source = sources[source_name]
            # Build select parameters for this specific source
            select_args = (
                source.select_param(filters, source_limit, source_offset)
                if source.select_param is not None
                else [source_limit, source_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=source_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)

        return SelectResponseModel(data=final_data, count=total_count)


def _create_select_processor(query_engine: SelectEngine):
    """Create a select processor for the middleware chain"""

    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):
    """Create a count processor for the middleware chain"""

    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
