from typing import List, Literal, Optional, Type, Union
import pyarrow as pa
import pypika
from deltalake import DeltaTable
from aiocache import Cache, cached
from aiocache.serializers import PickleSerializer
from fastapi import (
    APIRouter,
    BackgroundTasks,
    Depends,
    Header,
    HTTPException,
    Query,
    Request,
)
from pydantic import BaseModel

from bmsdna.lakeapi.context import get_context_by_engine
from bmsdna.lakeapi.context.df_base import ResultData, get_sql
from bmsdna.lakeapi.core.config import BasicConfig, Config, Configs
from bmsdna.lakeapi.core.datasource import (
    Datasource,
    filter_df_based_on_params,
    filter_partitions_based_on_params,
)
from bmsdna.lakeapi.core.log import get_logger
from bmsdna.lakeapi.core.model import create_parameter_model, create_response_model
from bmsdna.lakeapi.core.partition_utils import should_hide_colname
from bmsdna.lakeapi.core.response import create_response
from bmsdna.lakeapi.core.types import (
    OutputFileType,
    Engines,
)
from bmsdna.lakeapi.core.env import CACHE_EXPIRATION_TIME_SECONDS

cache = cached(ttl=CACHE_EXPIRATION_TIME_SECONDS, cache=Cache.MEMORY, serializer=PickleSerializer())

logger = get_logger(__name__)


async def get_partitions(datasource: Datasource, params: BaseModel, config: Config) -> Optional[list]:
    parts = (
        await filter_partitions_based_on_params(
            DeltaTable(datasource.uri).metadata(),
            params.dict(exclude_unset=True) if params else {},
            config.params or [],
        )
        if config.datasource.file_type == "delta"
        else None
    )
    return parts


def remove_search(prm_dict: dict, config: Config):
    if not config.search or len(prm_dict) == 0:
        return prm_dict
    search_cols = [s.name.lower() for s in config.search]
    return {k: v for k, v in prm_dict.items() if k.lower() not in search_cols}


async def get_params_filter_expr(columns: List[str], config: Config, params: BaseModel) -> Optional[pypika.Criterion]:
    expr = await filter_df_based_on_params(
        remove_search(params.dict(exclude_unset=True) if params else {}, config),
        config.params if config.params else [],
        columns,
    )
    return expr


def get_response_model(config: Config, metamodel: ResultData) -> Optional[Type[BaseModel]]:
    response_model: Optional[type[BaseModel]] = None
    try:
        response_model = create_response_model(config.tag + "_" + config.name, metamodel)
    except Exception as err:
        logger.warning(f"Could not get response type for f{config.route}. Error:{err}")
        response_model = None
    return response_model


def exclude_cols(columns: List[str]) -> List[str]:
    columns = [c for c in columns if not should_hide_colname(c)]
    return columns


def is_complex_type(schema: pa.Schema, col_name: str):
    f = schema.field(col_name)
    return pa.types.is_nested(f.type)


def split_csv(csv_str: str) -> list[str]:
    import csv

    reader = csv.reader([csv_str], delimiter=",", quotechar='"')
    for i in reader:
        return i
    raise ValueError("cannot happen")


def create_config_endpoint(
    metamodel: Optional[ResultData],
    apimethod: Literal["get", "post"],
    config: Config,
    router: APIRouter,
    response_model: Type,
    basic_config: BasicConfig,
    configs: Configs,
):
    route = config.route

    query_model = create_parameter_model(
        metamodel, config.tag + "_" + config.name + "_" + apimethod, config.params, config.search, apimethod
    )

    api_method_mapping = {
        "post": router.post(
            route,
            tags=[config.tag],
            response_model=response_model,
            operation_id=config.tag + "_" + config.name + "_" + apimethod,
            name=config.name,
        ),
        "get": router.get(
            route,
            tags=[config.tag],
            response_model=response_model,
            operation_id=config.tag + "_" + config.name + "_" + apimethod,
            name=config.name,
        ),
    }

    api_method = api_method_mapping[apimethod]
    has_complex = True
    if metamodel is not None:
        has_complex = any((pa.types.is_struct(t) or pa.types.is_list(t) for t in metamodel.arrow_schema().types))

    @api_method
    async def data(
        request: Request,
        params: query_model = (Depends() if apimethod == "get" else None),  # type: ignore
        Accept: Union[str, None] = Header(default=None),
        limit: Optional[int] = 100,
        offset: Optional[int] = 0,
        select: Union[str, None] = Query(title="$select", alias="$select", default=None, include_in_schema=False),
        distinct: bool = Query(title="$distinct", alias="$distinct", default=False, include_in_schema=False),
        engine: Engines = Query(title="$engine", alias="$engine", default="duckdb", include_in_schema=False),
        format: Optional[OutputFileType] = "json",
        jsonify_complex: bool = Query(title="jsonify_complex", include_in_schema=has_complex, default=False),
    ):  # type: ignore
        logger.debug(f"{params.dict(exclude_unset=True) if params else None}Union[ ,  ]{request.url.path}")

        engine = engine or basic_config.default_engine

        logger.debug(f"Engine: {engine}")

        with get_context_by_engine(engine, chunk_size=config.chunk_size or basic_config.default_chunk_size) as context:
            realdataframe = Datasource(
                config.version_str, config.tag, config.name, config.datasource, context, basic_config=basic_config
            )
            parts = await get_partitions(realdataframe, params, config)
            df = realdataframe.get_df(parts or None)

            expr = await get_params_filter_expr(df.columns(), config, params)
            base_schema = df.arrow_schema()
            new_query = df.query_builder()
            new_query = new_query.where(expr) if expr is not None else new_query

            searches = {}
            if config.search is not None and params is not None:
                search_dict = {c.name.lower(): c for c in config.search}
                searches = {
                    k: (v, search_dict[k.lower()])
                    for k, v in params.dict(exclude_unset=True).items()
                    if k.lower() in search_dict and v is not None and len(v) >= basic_config.min_search_length
                }

            import pypika

            columns = exclude_cols(df.columns())
            if select:
                columns = [
                    c for c in columns if c in split_csv(select)
                ]  # split , is a bit naive, we might want to support real CSV with quotes here
            if config.datasource.exclude and len(config.datasource.exclude) > 0:
                columns = [c for c in columns if c not in config.datasource.exclude]
            if config.datasource.sortby and len(searches) == 0:
                for s in config.datasource.sortby:
                    new_query = new_query.orderby(
                        s.by,
                        order=pypika.Order.desc if s.direction and s.direction.lower() == "desc" else pypika.Order.asc,
                    )
            if has_complex and format in ["csv", "excel", "scsv", "csv4excel"]:
                jsonify_complex = True
            if jsonify_complex:
                new_query = new_query.select(
                    *[
                        pypika.Field(c)
                        if not is_complex_type(base_schema, c)
                        else context.json_function(pypika.Field(c)).as_(c)
                        for c in columns
                    ]
                )
            else:
                new_query = new_query.select(*columns)

            if distinct:
                assert len(columns) <= 3  # reduce complexity here
                new_query = new_query.distinct()

            if not (limit == -1 and config.allow_get_all_pages):
                limit = 1000 if limit == -1 else limit
                new_query = new_query.offset(offset or 0).limit(limit)

            if len(searches) > 0 and config.search is not None:
                import pypika.queries
                import pypika.terms

                source_view = realdataframe.tablename
                context.init_search(source_view, config.search)
                score_sum = None
                for search_key, (search_val, search_cfg) in searches.items():
                    score_sum = (
                        context.search_score_function(source_view, search_val, search_cfg, alias=None)
                        if score_sum is None
                        else score_sum + context.search_score_function(source_view, search_val, search_cfg, alias=None)
                    )
                assert score_sum is not None
                new_query = new_query.select(score_sum.as_("search_score"))
                new_query = new_query.where(pypika.terms.NotNullCriterion(pypika.queries.Field("search_score")))

                new_query = new_query.orderby(pypika.Field("search_score"), order=pypika.Order.desc)

            logger.debug(f"Query: {get_sql(new_query)}")

            df2 = context.execute_sql(new_query)

            try:
                return await create_response(
                    request.url,
                    format or request.headers["Accept"],
                    df2,
                    context,
                    basic_config=basic_config,
                    close_context=True,
                )
            except Exception as err:
                logger.error("Error in creating response", exc_info=err)
                raise HTTPException(status_code=500)
