"""Module providing the iRail class for interacting with the iRail API."""

import asyncio
from asyncio import Lock
from datetime import datetime
import logging
import time
from types import TracebackType
from typing import Any, Type

from aiohttp import ClientError, ClientResponse, ClientSession

from pyrail.models import (
    CompositionApiResponse,
    ConnectionsApiResponse,
    DisturbancesApiResponse,
    LiveboardApiResponse,
    StationsApiResponse,
    VehicleApiResponse,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger: logging.Logger = logging.getLogger(__name__)


class iRail:
    """A Python wrapper for the iRail API, handling rate limiting and endpoint requests.

    Attributes:
        lang (str): The language for API responses ('nl', 'fr', 'en', 'de').
        session (ClientSession, optional): The HTTP session used for API requests.
            If not provided, a new session will be created and managed internally.
            If provided, the session lifecycle must be managed externally.

    Endpoints:
        stations: Retrieve all stations.
        liveboard: Retrieve a liveboard for a station or ID.
        connections: Retrieve connections between two stations.
        vehicle: Retrieve information about a specific train.
        composition: Retrieve the composition of a train.
        disturbances: Retrieve information about current disturbances on the rail network.

    """

    # Available iRail API endpoints and their parameter requirements.
    # Each endpoint is configured with required parameters, optional parameters, and XOR
    # parameter groups (where exactly one parameter from the group must be provided).
    endpoints: dict[str, dict[str, Any]] = {
        "stations": {},
        "liveboard": {"xor": ["station", "id"], "optional": ["date", "time", "arrdep", "alerts"]},
        "connections": {
            "required": ["from", "to"],
            "optional": ["date", "time", "timesel", "typeOfTransport"],
        },
        "vehicle": {"required": ["id"], "optional": ["date", "alerts"]},
        "composition": {"required": ["id"], "optional": ["data"]},
        "disturbances": {"optional": ["lineBreakCharacter"]},
    }

    def __init__(self, lang: str = "en", session: ClientSession | None = None) -> None:
        """Initialize the iRail API client.

        Args:
            lang (str): The language for API responses. Default is 'en'.
            session (ClientSession, optional): An existing aiohttp session. Defaults to None.

        """
        self.lang: str = lang
        self.tokens: int = 3
        self.burst_tokens: int = 5
        self.last_request_time: float = time.time()
        self.lock: Lock = Lock()
        self.session: ClientSession | None = session
        self._owns_session = session is None  # Track ownership
        self.etag_cache: dict[str, str] = {}
        logger.info("iRail instance created")

    async def __aenter__(self) -> "iRail":
        """Initialize the aiohttp client session when entering the async context."""
        if self.session and not self.session.closed:
            logger.debug("Using externally provided session")
        elif not self.session:
            logger.debug("Creating new internal aiohttp session")
            self.session = ClientSession()
        return self

    async def __aexit__(
        self, exc_type: Type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
    ) -> None:
        """Close the aiohttp client session when exiting the async context."""
        if self.session:
            if self.session.closed:
                logger.debug("Session is already closed, skipping closure")
            elif not self._owns_session:
                logger.debug("Session is externally provided; not closing it")
            else:
                logger.debug("Closing aiohttp session")
                try:
                    await self.session.close()
                except Exception as e:
                    logger.error("Error while closing aiohttp session: %s", e)

    @property
    def lang(self) -> str:
        """Get the language setting for API requests.

        Returns:
            str: The current language setting ('nl', 'fr', 'en', 'de').

        """
        return self.__lang

    @lang.setter
    def lang(self, value: str) -> None:
        """Set the language for API requests.

        Args:
            value (str): The language to use. Must be one of: 'nl', 'fr', 'en', 'de'.
                        If an invalid value is provided, defaults to 'en'.

        """
        if value in ["nl", "fr", "en", "de"]:
            self.__lang = value
        else:
            self.__lang = "en"

    def clear_etag_cache(self) -> None:
        """Clear the ETag cache."""
        self.etag_cache.clear()
        logger.info("ETag cache cleared")

    def _refill_tokens(self) -> None:
        """Refill tokens for rate limiting using a token bucket algorithm.

        - Standard tokens: Refill rate of 3 tokens/second, max 3 tokens.
        - Burst tokens: Refilled only when standard tokens are full, max 5 tokens.
        """
        logger.debug("Refilling tokens")
        current_time: float = time.time()
        elapsed: float = current_time - self.last_request_time
        self.last_request_time = current_time

        # Refill tokens, 3 tokens per second, cap tokens at 3
        self.tokens = min(3, self.tokens + int(elapsed * 3))

        # Refill burst tokens only if the request rate is below 3 requests per second
        if self.tokens == 3:
            self.burst_tokens = min(5, self.burst_tokens + int(elapsed * 3))

    async def _handle_rate_limit(self) -> None:
        """Handle rate limiting using a token bucket algorithm.

        - Standard tokens: 3 requests/second.
        - Burst tokens: Additional 5 requests/second for spikes.
        - Waits and refills tokens if both are exhausted.
        """
        logger.debug("Handling rate limit")
        self._refill_tokens()
        if self.tokens < 1:
            if self.burst_tokens >= 1:
                self.burst_tokens -= 1
            else:
                logger.warning("Rate limiting active, waiting for tokens")
                wait_time = max(0, 1 - (time.time() - self.last_request_time))
                await asyncio.sleep(wait_time)
                self._refill_tokens()
                self.tokens -= 1
        else:
            self.tokens -= 1

    def _add_etag_header(self, method: str) -> dict[str, str]:
        """Add ETag header for the given method if a cached ETag is available.

        Args:
            method (str): The API endpoint for which the header is being generated.

        Returns:
            dict[str, str]: A dictionary containing HTTP headers, including the ETag header
            if a cached value exists.

        """
        headers: dict[str, str] = {"User-Agent": "pyRail (https://github.com/tjorim/pyrail; tielemans.jorim@gmail.com)"}
        if method in self.etag_cache:
            logger.debug("Adding If-None-Match header with value: %s", self.etag_cache[method])
            headers["If-None-Match"] = self.etag_cache[method]
        return headers

    def _validate_date(self, date: str | None) -> bool:
        """Validate the date format (DDMMYY).

        Args:
            date (str, optional): The date string to validate. Expected format is DDMMYY,
            e.g., '150923' for September 15, 2023.

        Returns:
            bool: True if the date is valid or None is provided, False otherwise.

        """
        if not date:
            return True
        try:
            datetime.strptime(date, "%d%m%y")
            return True
        except ValueError:
            logger.error("Invalid date format. Expected DDMMYY (e.g., 150923 for September 15, 2023), got: %s", date)
            return False

    def _validate_time(self, time: str | None) -> bool:
        """Validate the time format (HHMM).

        Args:
            time (str, optional): The time string to validate. Expected format is HHMM,
            e.g., '1430' for 2:30 PM.

        Returns:
            bool: True if the time is valid or None is provided, False otherwise.

        """
        if not time:
            return True
        try:
            datetime.strptime(time, "%H%M")
            return True
        except ValueError:
            logger.error("Invalid time format. Expected HHMM (e.g., 1430 for 2:30 PM), got: %s", time)
            return False

    def _validate_params(self, method: str, params: dict[str, Any] | None = None) -> bool:
        """Validate parameters for a specific iRail API endpoint based on predefined requirements.

        Args:
            method (str): The API endpoint method to validate parameters for.
            params (dict[str, Any], optional): Dictionary of parameters to validate. Defaults to None.

        Returns:
            bool: True if parameters are valid, False otherwise.

        Raises:
            No explicit exceptions raised. Uses logging to report validation errors.

        Validates the following conditions:
            1. Checks if the method is a known endpoint
            2. Ensures all required parameters are present
            3. Enforces XOR (exclusive OR) parameter logic where applicable
            4. Prevents unexpected parameters from being included

        Example:
            # Validates parameters for a 'get_liveboard' method
            is_valid = self.validate_params('get_liveboard', {'station': 'Brussels', 'date': '2025-01-15'})

        """
        if method not in self.endpoints:
            logger.error("Unknown API endpoint: %s", method)
            return False

        endpoint = self.endpoints[method]
        required = endpoint.get("required", [])
        xor = endpoint.get("xor", [])
        optional = endpoint.get("optional", [])

        params = params or {}

        # Validate date and time formats if present
        if "date" in params and not self._validate_date(params["date"]):
            return False
        if "time" in params and not self._validate_time(params["time"]):
            return False

        # Ensure all required parameters are present
        for param in required:
            if param not in params or params[param] is None:
                logger.error("Missing required parameter: %s for endpoint: %s", param, method)
                return False

        # Check XOR logic (only one of XOR parameters can be set)
        if xor:
            xor_values = [params.get(param) is not None for param in xor]
            if sum(xor_values) != 1:
                logger.error("Exactly one of the XOR parameters %s must be provided for endpoint: %s", xor, method)
                return False

        # Ensure no unexpected parameters are included
        all_params = required + xor + optional
        for param in params.keys():
            if param not in all_params:
                logger.error("Unexpected parameter: %s for endpoint: %s", param, method)
                return False

        return True

    async def _handle_success_response(self, response: ClientResponse, method: str) -> dict[str, Any] | None:
        """Handle a successful API response."""
        if "Etag" in response.headers:
            self.etag_cache[method] = response.headers["Etag"]
        try:
            json_data: dict[str, Any] | None = await response.json()
            if not json_data:
                logger.warning("Empty response received")
            return json_data
        except ValueError:
            logger.error("Failed to parse JSON response")
            return None

    async def _handle_response(
        self, response: ClientResponse, method: str, args: dict[str, Any] | None = None
    ) -> dict[str, Any] | None:
        """Handle the API response based on status code."""
        if response.status == 429:
            retry_after: int = int(response.headers.get("Retry-After", 1))
            logger.warning("Rate limited, retrying after %d seconds", retry_after)
            await asyncio.sleep(retry_after)
            return await self._do_request(method, args)
        elif response.status == 400:
            logger.error("Bad request: %s", await response.text())
            return None
        elif response.status == 404:
            logger.error("Endpoint %s not found, response: %s", method, await response.text())
            return None
        elif response.status == 200:
            return await self._handle_success_response(response, method)
        elif response.status == 304:
            logger.info("Data not modified, using cached data for method %s", method)
            return None
        else:
            logger.error("Request failed with status code: %s, response: %s", response.status, await response.text())
            return None

    async def _do_request(self, method: str, args: dict[str, Any] | None = None) -> dict[str, Any] | None:
        """Send an asynchronous request to the specified iRail API endpoint.

        This method handles API requests with rate limiting, parameter validation,
        and ETag-based caching. It supports conditional requests and manages
        various HTTP response scenarios.

        Args:
            method (str): The iRail API endpoint method to request.
            args (dict, optional): Additional query parameters for the request.
                Defaults to None.

        Returns:
            dict or None: Parsed JSON response from the API if successful,
                None if the request fails or encounters an error.

        Raises:
            No explicit exceptions, but logs errors for various failure scenarios.

        Notes:
            - Implements rate limiting using a token-based mechanism
            - Supports ETag-based conditional requests
            - Handles rate limit (429) responses with automatic retry
            - Logs detailed information about request processing

        """
        logger.info("Starting request to endpoint: %s", method)
        if self.session is None:
            logger.error("Session not initialized. Use 'async with' context manager to initialize the client.")
            return None
        if not self._validate_params(method, args or {}):
            logger.error("Validation failed for method: %s with args: %s", method, args)
            return None
        async with self.lock:
            await self._handle_rate_limit()

        # Construct the request URL and parameters
        url: str = "https://api.irail.be/{}/".format(method)
        params = {"format": "json", "lang": self.lang}
        if args:
            params.update(args)

        request_headers: dict[str, str] = self._add_etag_header(method)

        try:
            async with self.session.get(url, params=params, headers=request_headers) as response:
                return await self._handle_response(response, method, args)
        except ClientError as e:
            logger.error("Request failed due to an exception: %s", e)
            return None

    async def get_stations(self) -> StationsApiResponse | None:
        """Retrieve a list of all train stations from the iRail API.

        This method fetches the complete list of available train stations without any additional filtering parameters.

        Returns:
            StationsApiResponse or None: A StationsApiResponse object containing a list of all train stations and their details, or None if the request fails.

        Example:
            async with iRail() as client:
                stations = await client.get_stations()
                if stations:
                    print(f"Total stations: {len(stations)}")

        """
        stations_response_dict = await self._do_request("stations")
        if stations_response_dict is None:
            return None
        return StationsApiResponse.from_dict(stations_response_dict)

    async def get_liveboard(
        self,
        station: str | None = None,
        id: str | None = None,
        date: str | None = None,
        time: str | None = None,
        arrdep: str = "departure",
        alerts: bool = False,
    ) -> LiveboardApiResponse | None:
        """Retrieve a liveboard for a specific train station.

        Asynchronously fetches live departure or arrival information for a given station.
        Provide exactly one of 'station' or 'id'.

        Args:
            station (str): Name of the train station.
            id (str): Unique identifier of the train station.
            date (str, optional): Date for the liveboard in format 'DDMMYY'. Defaults to current date.
            time (str, optional): Time for the liveboard in format 'HHMM'. Defaults to current time.
            arrdep (str, optional): Type of board to retrieve. Either 'departure' (default) or 'arrival'.
            alerts (bool, optional): Whether to include service alerts. Defaults to False.

        Returns:
            LiveboardApiResponse or None: A LiveboardApiResponse object containing liveboard information, or None if the request fails.

        Example:
            async with iRail() as client:
                liveboard = await client.get_liveboard(station='Brussels-South')
                if liveboard:
                    print(f"Liveboard for Brussels-South: {liveboard}")

        """
        extra_params: dict[str, Any] = {
            "station": station,
            "id": id,
            "date": date,
            "time": time,
            "arrdep": arrdep,
            "alerts": "true" if alerts else "false",
        }
        liveboard_response_dict = await self._do_request(
            "liveboard", {k: v for k, v in extra_params.items() if v is not None}
        )
        if liveboard_response_dict is None:
            return None
        return LiveboardApiResponse.from_dict(liveboard_response_dict)

    async def get_connections(
        self,
        from_station: str,
        to_station: str,
        date: str | None = None,
        time: str | None = None,
        timesel: str = "departure",
        type_of_transport: str = "automatic",
    ) -> ConnectionsApiResponse | None:
        """Retrieve train connections between two stations using the iRail API.

        Args:
            from_station (str): Name or ID of the departure station
            to_station (str): Name or ID of the arrival station
            date (str, optional): Date of travel in format 'DDMMYY' (default: current date)
            time (str, optional): Time of travel in format 'HH:MM' (default: current time)
            timesel (str, optional): Time selection type, either 'departure' or 'arrival' (default: 'departure')
            type_of_transport (str, optional): Type of transport, options include 'automatic', 'trains', 'nointernationaltrains' or 'all' (default: 'automatic')

        Returns:
            ConnectionsApiResponse or None: A ConnectionsApiResponse object containing connection details, or None if the request fails or no connections are found.

        Example:
            async with iRail() as client:
                connections = await client.get_connections("Antwerpen-Centraal", "Brussel-Centraal")
                if connections:
                    print(f"Connections from Antwerpen-Centraal to Brussel-Centraal: {connections}")

        """
        extra_params: dict[str, Any] = {
            "from": from_station,
            "to": to_station,
            "date": date,
            "time": time,
            "timesel": timesel,
            "typeOfTransport": type_of_transport,
        }
        connections_response_dict = await self._do_request(
            "connections", {k: v for k, v in extra_params.items() if v is not None}
        )
        if connections_response_dict is None:
            return None
        return ConnectionsApiResponse.from_dict(connections_response_dict)

    async def get_vehicle(self, id: str, date: str | None = None, alerts: bool = False) -> VehicleApiResponse | None:
        """Retrieve detailed information about a specific train vehicle.

        Args:
            id (str): Unique identifier of the train vehicle to retrieve information for.
            date (str, optional): Specific date for which vehicle information is requested. Defaults to None (current date).
            alerts (bool, optional): Flag to include service alerts for the vehicle. Defaults to False.

        Returns:
            VehicleApiResponse or None: A VehicleApiResponse object containing vehicle details, or None if the request fails.

        Example:
            async with iRail() as client:
                vehicle_info = await client.get_vehicle("BE.NMBS.IC1832")

        """
        extra_params: dict[str, Any] = {"id": id, "date": date, "alerts": "true" if alerts else "false"}
        vehicle_response_dict = await self._do_request(
            "vehicle", {k: v for k, v in extra_params.items() if v is not None}
        )
        if vehicle_response_dict is None:
            return None
        return VehicleApiResponse.from_dict(vehicle_response_dict)

    async def get_composition(self, id: str, data: str | None = None) -> CompositionApiResponse | None:
        """Retrieve the composition details of a specific train.

        Args:
            id (str): The unique identifier of the train for which composition details are requested.
            data (str, optional): Additional data parameter to get all raw unfiltered data as iRail fetches it from the NMBS (set to 'all'). Defaults to '' (filtered data).

        Returns:
            CompositionApiResponse or None: A CompositionApiResponse object containing train composition details, or None if the request fails.

        Example:
            async with iRail() as client:
                composition = await client.get_composition('S51507')

        """
        extra_params: dict[str, str | None] = {"id": id, "data": data}
        composition_response_dict = await self._do_request(
            "composition", {k: v for k, v in extra_params.items() if v is not None}
        )
        if composition_response_dict is None:
            return None
        return CompositionApiResponse.from_dict(composition_response_dict)

    async def get_disturbances(self, line_break_character: str | None = None) -> DisturbancesApiResponse | None:
        """Retrieve information about current disturbances on the rail network.

        Args:
            line_break_character (str, optional): A custom character to use for line breaks in the disturbance description. Defaults to ''.

        Returns:
            DisturbancesApiResponse or None: A DisturbancesApiResponse object containing disturbance information, or None if the request fails or no disturbances are found.

        Example:
            async with iRail() as client:
                disturbances = await client.get_disturbances()
                if disturbances:
                    print(f"Current disturbances: {disturbances}")

        """
        extra_params = {"lineBreakCharacter": line_break_character}
        disturbances_response_dict = await self._do_request(
            "disturbances", {k: v for k, v in extra_params.items() if v is not None}
        )
        if disturbances_response_dict is None:
            return None
        return DisturbancesApiResponse.from_dict(disturbances_response_dict)
