from datetime import date
from typing import Optional
from typing import TypedDict
from uuid import UUID

from marshmallow import Schema, fields

from charts.connectors import snowflake
from charts.features import ff_chart_sql_filter

from charts.constants.chart_columns import CHART_COLUMN_NAMES
from charts.logic import sql_loader
from charts.utils.dataloader import dataloader, DataLoaderResponse, MarshmallowRedisDataLoader
from charts.utils.db.mapper import map_db_result_with_column_names


class ChartGroupsChartIdsSchema(Schema):
    slug = fields.String(required=True)
    chart_ids = fields.List(fields.UUID(required=True), required=True)
    default_chart_id = fields.UUID(required=True)


class ChartGroupsChartIdsResponse(TypedDict):
    slug: str
    chart_ids: list[UUID]
    default_chart_id: str


def _escape(slug: str) -> str:
    return slug.replace("'", "\\'")


@dataloader
def _get_chart_groups_chart_ids(chart_slugs: list[str], *, controlled_charts_ffs: list[str]) -> list[DataLoaderResponse[ChartGroupsChartIdsResponse]]:
    sql = sql_loader.load_query('chart_group_chart_ids_dataloaded')

    controlled_charts_filter = ff_chart_sql_filter(controlled_charts_ffs)

    sql = sql.replace('{ chart_filter }', controlled_charts_filter)

    params: snowflake.SnowflakeParameterType = {
        'chart_groups': [[f"({i}, '{_escape(slug)}')"] for i, slug in enumerate(chart_slugs)]
    }

    columns = (
        'slug',
        'chart_ids',
        'default_chart_id',
    )

    raw_db_result = snowflake.fetchall_nocache(sql, params)
    results = map_db_result_with_column_names(raw_db_result, columns)

    return [
        {
            'data': ChartGroupsChartIdsResponse({
                'slug': result['slug'],
                'chart_ids': result['chart_ids'].split(','),
                'default_chart_id': result['default_chart_id'],
            }) if result['slug'] else None
        }
        for result
        in results
    ]


def _key_serializer(chart_slug: str, *, controlled_charts_ffs: list[str]) -> str:
    return f'chart_group:{chart_slug}:chart_ids:controlled_chart_ffs:{sorted(controlled_charts_ffs)}'


chart_groups_chart_ids_dataloader = MarshmallowRedisDataLoader(
    dataloader=_get_chart_groups_chart_ids,
    expiry=7 * 24 * 60 * 60,
    key_serializer=_key_serializer,
    schema=ChartGroupsChartIdsSchema,
)