import logging
import os
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, get_args, get_origin

import csp
from ccflow import BaseModel
from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
from csp.impl.types.tstype import isTsType
from pydantic import Field, model_validator

from csp_gateway.server import (
    GatewayChannels,
    GatewayModule,
)

from .adapter import FileDropAdapterConfiguration, FileDropType, filedrop_adapter_def

__all__ = (
    "ReadFileDropConfiguration",
    "ReadFileDrop",
)

log = logging.getLogger(__name__)


class ReadFileDropConfiguration(BaseModel):
    """The configuration of a filedrop adapter for a directory and filetype"""

    dir_path: Path = Field(description="Directory to monitor for new files")
    channel_name: str = Field(description="Name of the channel to send the structs to")
    filedrop_type: Union[FileDropType, str] = Field(description="The type of files to expect and accordingly read i.e. parquet, json, etc")
    extensions: List[str] = Field(default=[], description="List of extensions to decide which files to read, empty list means all extensions")
    loader: Optional[Callable[[str], List[Any]]] = Field(default=None, description="custom loader for loading files into list of struct like data")
    deserializer: Optional[Callable[[Any], Any]] = Field(
        default=None, description="deserialize each data point to data type expected by channel type"
    )
    subscribe_with_struct_id: bool = Field(
        default=False,
        description=("If False, replaces the id field on GatewayStructs from files with one autogenerated by the current Gateway."),
    )
    subscribe_with_struct_timestamp: bool = Field(
        default=False,
        description=("If False, replaces the timestamp field on the GatewayStruct with a timestamp autogenerated by the current Gateway."),
    )

    @model_validator(mode="after")
    def check_filedrop_type(self) -> "ReadFileDropConfiguration":
        if isinstance(self.filedrop_type, str):
            try:
                self.filedrop_type = FileDropType[self.filedrop_type]
            except KeyError:
                raise ValueError(f"{self.filedrop_type} is not a valid FileDropType")
        if self.filedrop_type == FileDropType.CUSTOM and self.loader is None:
            raise ValueError("loader must be set if filedrop_type is CUSTOM")
        return self


T = TypeVar("T")
K = TypeVar("K")


class ReadFileDrop(GatewayModule):
    """The module to read data from files dropped in specific directories and send them as structs to specific channels"""

    configs: List[ReadFileDropConfiguration] = Field(description="List of configs for the directories to monitor")

    @csp.node
    def handle_list_basket(self, data: csp.ts[List[T]], list_size: int) -> csp.OutputBasket(List[csp.ts[T]], shape="list_size"):
        if csp.ticked(data):
            return data

    @csp.node
    def handle_dict_basket(self, data: csp.ts[Dict[K, T]], dict_keys: List[K]) -> csp.OutputBasket(Dict[K, csp.ts[T]], shape="dict_keys"):
        if csp.ticked(data):
            return data

    def connect(self, channels: GatewayChannels):
        channel_data = {}
        channel_basket_types = {}
        for config in self.configs:
            # ensure that directories exist
            os.makedirs(config.dir_path, exist_ok=True)
            channel_data[config.channel_name] = []
        for config in self.configs:
            context = {
                "force_new_id": not config.subscribe_with_struct_id,
                "force_new_timestamp": not config.subscribe_with_struct_timestamp,
            }
            adapter_config = FileDropAdapterConfiguration(
                dir_path=str(config.dir_path),
                filedrop_type=config.filedrop_type,
                extensions=config.extensions,
                type_adapter_args=context,
                loader=config.loader,
                deserializer=config.deserializer,
            )
            channel_type = channels.get_outer_type(config.channel_name)
            if isTsType(channel_type):
                non_ts_type = channel_type.typ
                channel_basket_types[config.channel_name] = ""
            else:
                normalized_type = ContainerTypeNormalizer.normalize_type(channel_type)
                if get_origin(normalized_type) is list:
                    inner_type = get_args(normalized_type)[0]
                    if not isTsType(inner_type):
                        raise ValueError(f"Channel type for {config.channel_name} should be of the form List[TsType], got: {channel_type}")
                    non_ts_type = List[inner_type.typ]
                    channel_basket_types[config.channel_name] = "list"
                elif get_origin(normalized_type) is dict:
                    key_type, inner_type = get_args(normalized_type)
                    if not isTsType(inner_type):
                        raise ValueError(f"Channel type for {config.channel_name} should be of the form Dict[KeyType, TsType], got: {channel_type}")
                    non_ts_type = Dict[key_type, inner_type.typ]
                    channel_basket_types[config.channel_name] = "dict"
                else:
                    raise Exception(f"Channel type cannot be handled: {channel_type}")
            data = filedrop_adapter_def(config=adapter_config, ts_typ=non_ts_type)
            channel_data[config.channel_name].append(data)
        for channel_name, data_list in channel_data.items():
            data = csp.flatten(data_list)
            if channel_basket_types[channel_name] == "list":
                channels.set_channel(channel_name, self.handle_list_basket(data, list_size=channels.dynamic_keys()[channel_name]))
            elif channel_basket_types[channel_name] == "dict":
                channels.set_channel(channel_name, self.handle_dict_basket(data, dict_keys=channels.dynamic_keys()[channel_name]))
            else:
                channels.set_channel(channel_name, data)
