"""Cloud exchange configuration file parsing."""

from __future__ import annotations

import abc
import logging
import pathlib
import sys
from logging.handlers import TimedRotatingFileHandler
from typing import Any
from typing import BinaryIO
from typing import Literal
from typing import Protocol
from typing import TypeVar

if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
    from typing import Self
else:  # pragma: <3.11 cover
    from typing_extensions import Self

from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field

from academy.exchange.cloud.backend import MailboxBackend
from academy.exchange.cloud.backend import PythonBackend
from academy.exchange.cloud.backend import RedisBackend
from academy.logging import _Formatter

if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
    import tomllib
else:  # pragma: <3.11 cover
    import tomli as tomllib  # type: ignore[import-not-found,unused-ignore]


BaseModelT = TypeVar('BaseModelT', bound=BaseModel)


class LogConfig(BaseModel):
    """Log config for exchange server.

    Attributes:
        level: Console log verbosity.
        logfile: Base file name to write to.
        logfile_level: Verbosity of logs in file.
        rotate: Weather to rotate logs or not.
        rotation_interval_days: How often in days to rotate logs.
        backup_count: How many previous log files to keep
    """

    model_config = ConfigDict(extra='forbid')
    level: int | str = logging.WARNING
    logfile: str | None = None
    logfile_level: int | str = logging.INFO
    rotate: bool = False
    rotation_interval_days: int = Field(default=1, gt=0)
    backup_count: int = Field(default=7, gt=0)

    def init_logger(self) -> logging.Logger:
        """Initialize logger from this configuration."""
        stdout_handler = logging.StreamHandler(sys.stdout)
        stdout_handler.setFormatter(_Formatter())
        stdout_handler.setLevel(self.level)
        handlers: list[logging.Handler] = [stdout_handler]

        if self.logfile is not None:
            logfile_level = (
                self.level
                if self.logfile_level is None
                else self.logfile_level
            )
            path = pathlib.Path(self.logfile)
            path.parent.mkdir(parents=True, exist_ok=True)
            file_handler: logging.Handler
            if self.rotate:
                file_handler = TimedRotatingFileHandler(
                    path,
                    when='d',
                    interval=self.rotation_interval_days,
                    backupCount=self.backup_count,
                )
            else:
                file_handler = logging.FileHandler(path)

            file_handler.setFormatter(_Formatter(color=False))
            file_handler.setLevel(logfile_level)
            handlers.append(file_handler)

        logging.basicConfig(
            datefmt='%Y-%m-%d %H:%M:%S',
            level=logging.NOTSET,
            handlers=handlers,
        )

        # This needs to be after the configuration of the root logger because
        # warnings get logged to a 'py.warnings' logger.
        logging.captureWarnings(True)

        logger = logging.getLogger()
        logger.info(
            (
                'Configured logger (stdout-level=%s, '
                'logfile=%s, logfile-level=%s)'
            ),
            logging.getLevelName(self.level)
            if isinstance(self.level, int)
            else self.level,
            self.logfile,
            logging.getLevelName(self.logfile_level)
            if isinstance(self.logfile_level, int)
            else self.logfile_level,
        )
        return logger


class ExchangeAuthConfig(BaseModel):
    """Exchange authentication configuration.

    Attributes:
        method: Authentication method.
        kwargs: Arbitrary keyword arguments to pass to the authenticator.
            The kwargs are excluded from the [`repr()`][repr] of this
            class because they often contain secrets.
    """

    model_config = ConfigDict(extra='forbid')

    method: Literal['globus'] | None = None
    kwargs: dict[str, Any] = Field(default_factory=dict, repr=False)


class BackendConfig(Protocol):
    """Config for backend of storing messages."""

    @abc.abstractmethod
    def get_backend(self) -> MailboxBackend:
        """Construct an instance of the backend from the config."""
        ...


class PythonBackendConfig(BaseModel):
    """Config for using PythonBackend."""

    model_config = ConfigDict(extra='forbid')

    message_size_limit_kb: int = 1024
    kind: Literal['python'] = Field(default='python', repr=False)

    def get_backend(self) -> MailboxBackend:
        """Construct an instance of the backend from the config."""
        return PythonBackend()


class RedisBackendConfig(BaseModel):
    """Config for RedisBackend.

    Attributes:
        hostname: Redis host
        port: Redis port
        kwargs: Any additional args to Redis
    """

    model_config = ConfigDict(extra='forbid')

    hostname: str = 'localhost'
    port: int = 6379
    message_size_limit_kb: int = Field(default=1024, gt=0, le=1024 * 512)
    mailbox_expiration_d: float = Field(default=7, gt=0)
    gravestone_expiration_d: float = Field(default=365, gt=0)
    kwargs: dict[str, Any] = Field(default_factory=dict, repr=False)
    kind: Literal['redis'] = Field(default='redis', repr=False)

    def get_backend(self) -> MailboxBackend:
        """Construct an instance of the backend from the config."""
        return RedisBackend(
            self.hostname,
            self.port,
            message_size_limit_kb=self.message_size_limit_kb,
            kwargs=self.kwargs,
            mailbox_expiration_s=int(self.mailbox_expiration_d * 24 * 3600),
            gravestone_expiration_s=int(
                self.gravestone_expiration_d * 24 * 3600,
            ),
        )


BackendConfigT = PythonBackendConfig | RedisBackendConfig


class ExchangeServingConfig(BaseModel):
    """Exchange serving configuration.

    Attributes:
        host: Network interface the server binds to.
        port: Network port the server binds to.
        certfile: Certificate file (PEM format) use to enable TLS.
        keyfile: Private key file. If not specified, the key will be
            taken from the certfile.
        auth: Authentication configuration.
        log_file: Location to write logs.
        log_level: Verbosity of logs.
    """

    model_config = ConfigDict(extra='forbid')

    host: str = 'localhost'
    port: int = 8700
    certfile: str | None = None
    keyfile: str | None = None
    auth: ExchangeAuthConfig = Field(default_factory=ExchangeAuthConfig)
    backend: BackendConfigT = Field(default_factory=PythonBackendConfig)
    logger: LogConfig = Field(default_factory=LogConfig)

    @classmethod
    def from_toml(cls, filepath: str | pathlib.Path) -> Self:
        """Parse an TOML config file.

        Example:
            Minimal config without SSL and without authentication.
            ```toml title="exchange.toml"
            port = 8700
            ```

            ```python
            from academy_exchange.config import ExchangeServingConfig

            config = ExchangeServingConfig.from_toml('exchange.toml')
            ```

        Example:
            Serve with SSL and Globus Auth.
            ```toml title="relay.toml"
            host = "0.0.0.0"
            port = 8700
            certfile = "/path/to/cert.pem"
            keyfile = "/path/to/privkey.pem"

            [auth]
            method = "globus"

            [auth.kwargs]
            client_id = "..."
            client_secret = "..."
            ```

        Note:
            Omitted values will be set to their defaults (if they are an
            optional value with a default).
            ```toml title="relay.toml"
            [serving]
            certfile = "/path/to/cert.pem"
            ```

            ```python
            from academy_exchange.config import ExchangeServingConfig

            config = ExchangeServingConfig.from_config('relay.toml')
            assert config.certfile == '/path/to/cert.pem'
            assert config.keyfile is None
            ```
        """
        with open(filepath, 'rb') as f:
            return load(cls, f)


def load(model: type[BaseModelT], fp: BinaryIO) -> BaseModelT:
    """Parse TOML from a binary file to a data class.

    Args:
        model: Config model type to parse TOML using.
        fp: File-like bytes stream to read in.

    Returns:
        Model initialized from TOML file.
    """
    return loads(model, fp.read().decode())


def loads(model: type[BaseModelT], data: str) -> BaseModelT:
    """Parse TOML string to data class.

    Args:
        model: Config model type to parse TOML using.
        data: TOML string to parse.

    Returns:
        Model initialized from TOML file.
    """
    data_dict = tomllib.loads(data)
    return model.model_validate(data_dict)
