import asyncio
from collections.abc import Callable, Coroutine
from functools import partial
import json
from json import JSONDecodeError
import logging
import os
import shutil
from typing import Any, NotRequired, TypedDict, cast

import aiohttp
from aiohttp import ClientError
from awesomeversion import AwesomeVersion
from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.storage import STORAGE_DIR
from homeassistant.loader import async_get_integration

from custom_components.powercalc.const import (
    API_URL,
    BUILT_IN_LIBRARY_DIR,
    DOMAIN,
    LIBRARY_DISCOVERY_LOW_PRIORITY_DOMAINS,
)
from custom_components.powercalc.helpers import async_cache, clear_async_cache
from custom_components.powercalc.power_profile.error import LibraryLoadingError, ProfileDownloadError
from custom_components.powercalc.power_profile.loader.protocol import Loader, ModelMetadata
from custom_components.powercalc.power_profile.power_profile import DeviceType, DiscoveryBy

_LOGGER = logging.getLogger(__name__)

ENDPOINT_LIBRARY = f"{API_URL}/library"
ENDPOINT_DOWNLOAD = f"{API_URL}/download"

TIMEOUT_SECONDS = 30


class LibraryModel(TypedDict):
    id: str
    name: NotRequired[str]
    aliases: NotRequired[list[str]]
    legacy_ids: NotRequired[list[str]]
    hash: str
    device_type: NotRequired[DeviceType]
    discovery_by: NotRequired[DiscoveryBy]
    min_version: NotRequired[str]


class LibraryManufacturer(TypedDict):
    name: str
    dir_name: str
    aliases: NotRequired[list[str]]
    models: list[LibraryModel]


class RemoteLoader(Loader):
    retry_timeout = 3

    def __init__(self, hass: HomeAssistant) -> None:
        self.hass = hass
        self.library_contents: dict[str, Any] = {}
        self.model_infos: dict[str, LibraryModel] = {}
        self.manufacturer_models: dict[str, list[LibraryModel]] = {}
        self.model_lookup: dict[str, dict[str, list[LibraryModel]]] = {}
        self.manufacturer_lookup: dict[str, set[str]] = {}
        self.profile_hashes: dict[str, str] = {}

    async def initialize(self, prefer_cached: bool = False) -> None:
        """Initialize the loader.

        Pass `prefer_cached` to keep the network off the critical path, using the library.json
        already in local storage when there is one. Only the very first run has to download.
        """

        integration = await async_get_integration(self.hass, DOMAIN)
        powercalc_version = AwesomeVersion(str(integration.version))

        self._clear_caches()
        self.library_contents = await self.load_library_json(prefer_cached)
        self.profile_hashes = await self.hass.async_add_executor_job(self._load_profile_hashes)

        self.model_infos.clear()
        self.model_lookup.clear()
        self.manufacturer_models.clear()
        self.manufacturer_lookup.clear()

        manufacturers: list[LibraryManufacturer] = self.library_contents.get("manufacturers", [])

        for manufacturer in manufacturers:
            self._index_manufacturer(manufacturer, powercalc_version)

    def get_discovery_low_priority_domains(self) -> set[str]:
        """Get the low priority discovery integration domains declared by library metadata."""
        return set(self.library_contents.get(LIBRARY_DISCOVERY_LOW_PRIORITY_DOMAINS, []))

    def _index_manufacturer(self, manufacturer: LibraryManufacturer, powercalc_version: AwesomeVersion) -> None:
        """Register a manufacturer, its aliases and all of its supported models in the lookup tables."""
        manufacturer_name = str(manufacturer.get("dir_name"))
        models: list[LibraryModel] = manufacturer.get("models", []) or []

        # manufacturer alias map (alias -> {canonical manufacturer_name})
        self.manufacturer_lookup.setdefault(manufacturer_name.lower(), set()).add(manufacturer_name)
        for alias in manufacturer.get("aliases", []) or []:
            self.manufacturer_lookup.setdefault(str(alias).lower(), set()).add(manufacturer_name)

        # per-manufacturer model lookup
        kept_models: list[LibraryModel] = []
        lookup: dict[str, list[LibraryModel]] = {}

        for model in models:
            model_id = str(model.get("id"))
            self.model_infos[f"{manufacturer_name}/{model_id}"] = model

            if self._is_unsupported_version(manufacturer_name, model_id, model, powercalc_version):
                continue

            kept_models.append(model)
            self._add_model_to_lookup(lookup, model, model_id.lower())

        self.manufacturer_models[manufacturer_name] = kept_models
        self.model_lookup[manufacturer_name] = lookup

    @staticmethod
    def _is_unsupported_version(
        manufacturer_name: str,
        model_id: str,
        model: LibraryModel,
        powercalc_version: AwesomeVersion,
    ) -> bool:
        """Check whether the model requires a newer powercalc version than the one installed."""
        min_version = model.get("min_version")
        if not min_version or powercalc_version >= AwesomeVersion(min_version):
            return False

        _LOGGER.debug(
            "Skipping model %s/%s as it requires powercalc version %s (current: %s)",
            manufacturer_name,
            model_id,
            min_version,
            powercalc_version,
        )
        return True

    @staticmethod
    def _add_model_to_lookup(lookup: dict[str, list[LibraryModel]], model: LibraryModel, model_id_lower: str) -> None:
        """Bucket a model by its id and aliases. Exact ids take priority over aliases."""
        # Exact id bucket first (highest priority)
        lookup.setdefault(model_id_lower, []).insert(0, model)

        # Alias buckets afterwards (lower priority)
        for alias in model.get("aliases", []) or []:
            alias_lower = str(alias).lower()
            if alias_lower == model_id_lower:
                continue
            # Append to the end to ensure aliased models are always last
            lookup.setdefault(alias_lower, []).append(model)

    def _clear_caches(self) -> None:
        """Clear cached lookups backed by mutable library state."""
        clear_async_cache(self.get_manufacturer_listing)
        clear_async_cache(self.find_manufacturers)
        clear_async_cache(self.get_model_listing)
        clear_async_cache(self.find_model)
        clear_async_cache(self.find_model_migration)
        clear_async_cache(self.load_model)

    async def load_library_json(self, prefer_cached: bool = False) -> dict[str, Any]:
        """Load library.json, from local storage or from the download API.

        With `prefer_cached` the locally stored copy wins when it exists, so the caller never
        waits on the network. The periodic library update refreshes it later.
        """
        if prefer_cached:
            cached_library = await self.hass.async_add_executor_job(self._read_local_library_json)
            if cached_library is not None:
                _LOGGER.debug("Loaded library.json from local storage")
                return cached_library
            _LOGGER.debug("No library.json in local storage yet, downloading it")

        try:
            return cast(dict[str, Any], await self.download_with_retry(self._download_remote_library_json))
        except ProfileDownloadError:
            _LOGGER.debug("Failed to download library.json, falling back to local copy")
            return await self.hass.async_add_executor_job(self._load_local_library_json)

    def _get_library_json_path(self) -> str:
        """Retrieve the local storage path for the library.json file."""
        return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, "library.json"))

    def _read_local_library_json(self) -> dict[str, Any] | None:
        """Read library.json from local storage, None when it has not been downloaded yet."""
        local_path = self._get_library_json_path()
        if not os.path.exists(local_path):
            return None
        with open(local_path) as f:
            return cast(dict[str, Any], json.load(f))

    def _load_local_library_json(self) -> dict[str, Any]:
        """Load library.json from local storage, raising when it is not there."""
        library_json = self._read_local_library_json()
        if library_json is None:
            raise ProfileDownloadError("Local library.json file not found")
        return library_json

    async def _download_remote_library_json(self) -> dict[str, Any] | None:
        """
        Download library.json from Github.
        On success, save it to local storage as a fallback for internet connection issues.
        """
        _LOGGER.debug("Loading library.json from github")

        local_path = self._get_library_json_path()
        session = async_get_clientsession(self.hass)

        try:
            async with asyncio.timeout(TIMEOUT_SECONDS), session.get(ENDPOINT_LIBRARY) as resp:
                if resp.status != 200:
                    raise ProfileDownloadError(
                        f"Failed to download library.json, unexpected status code: {resp.status}",
                    )

                data = await resp.read()

        except (TimeoutError, ClientError) as err:
            raise ProfileDownloadError(f"Failed to download library.json: {err}") from err

        def _save_to_local_storage(data: bytes) -> None:
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            with open(local_path, "wb") as f:
                f.write(data)

        await self.hass.async_add_executor_job(_save_to_local_storage, data)

        return cast(dict[str, Any], json.loads(data))

    @async_cache
    async def get_manufacturer_listing(
        self,
        device_types: set[DeviceType] | None,
        discovery_by: DiscoveryBy | None = None,
    ) -> set[tuple[str, str]]:
        """Get listing of available manufacturers."""

        return {
            (manufacturer["dir_name"], manufacturer["full_name"])
            for manufacturer in self.library_contents.get("manufacturers", [])
            if any(
                self._model_matches_filters(model, device_types, discovery_by)
                # Use the indexed models, so models requiring a newer Powercalc version are left out here as well.
                for model in self.manufacturer_models.get(str(manufacturer.get("dir_name")), [])
            )
        }

    @async_cache
    async def find_manufacturers(self, search: str) -> set[str]:
        """Find the manufacturer in the library."""
        return self.manufacturer_lookup.get(search.lower(), set())

    @async_cache
    async def get_model_listing(
        self,
        manufacturer: str,
        device_types: set[DeviceType] | None,
        discovery_by: DiscoveryBy | None = None,
    ) -> set[tuple[str, str]]:
        """Get listing of available models and display names for a given manufacturer."""
        models = self.manufacturer_models.get(manufacturer)
        if not models:
            return set()

        return {
            (model["id"], str(model.get("name") or model["id"]))
            for model in self.manufacturer_models.get(manufacturer, [])
            if self._model_matches_filters(model, device_types, discovery_by)
        }

    @staticmethod
    def _model_matches_filters(
        model: LibraryModel,
        device_types: set[DeviceType] | None,
        discovery_by: DiscoveryBy | None,
    ) -> bool:
        """Check whether an indexed model passes the requested filters.

        Device types and discovery modes this Powercalc version does not know about are treated
        as a non match, so profiles using a newly introduced value never break the listings.
        """
        try:
            model_device_type = DeviceType(model.get("device_type", DeviceType.LIGHT))
            model_discovery_by = DiscoveryBy(model.get("discovery_by", DiscoveryBy.ENTITY))
        except ValueError:
            return False

        if device_types and model_device_type not in device_types:
            return False

        return not discovery_by or model_discovery_by == discovery_by

    @async_cache
    async def find_model(self, manufacturer: str, search: set[str]) -> list[str]:
        """Find matching model IDs in the library."""
        models = self.model_lookup.get(manufacturer, {})
        return [
            model["id"]
            for phrase in search
            if (phrase_lower := phrase.lower()) in models
            for model in models[phrase_lower]
        ]

    @async_cache
    async def find_model_migration(self, manufacturer: str, model: str) -> str | None:
        """Find the canonical model id for a legacy profile id."""
        model_lower = model.lower()
        matches = {
            str(model_data.get("id"))
            for manufacturer_data in self.library_contents.get("manufacturers", [])
            if str(manufacturer_data.get("dir_name", "")).lower() == manufacturer
            for model_data in manufacturer_data.get("models", []) or []
            if model_lower in {str(legacy_id).lower() for legacy_id in model_data.get("legacy_ids", []) or []}
        }

        if len(matches) != 1:
            return None

        return next(iter(matches))

    async def get_model_metadata(self, manufacturer: str, model: str) -> ModelMetadata | None:
        """Return discovery metadata straight from the library index, without downloading the profile."""
        model_info = self.model_infos.get(f"{manufacturer}/{model}")
        if not model_info:
            return None

        try:
            device_type = DeviceType(model_info.get("device_type", DeviceType.LIGHT))
            discovery_by = DiscoveryBy(model_info.get("discovery_by", DiscoveryBy.ENTITY))
        except ValueError:
            return None

        return ModelMetadata(device_type=device_type, discovery_by=discovery_by)

    @async_cache
    async def load_model(
        self,
        manufacturer: str,
        model: str,
        force_update: bool = False,
        retry_count: int = 0,
    ) -> tuple[dict[str, Any], str] | None:
        """Load a model, downloading it if necessary, with retry logic."""
        model_info = self._get_library_model(manufacturer, model)
        storage_path = self.get_storage_path(manufacturer, model)
        model_path = os.path.join(storage_path, "model.json")

        if await self._needs_update(model_info, manufacturer, model, model_path, force_update):
            await self._download_profile_with_retry(manufacturer, model, storage_path, model_path)

        try:
            json_data = await self._load_model_json(model_path)
        except JSONDecodeError as e:
            return await self._handle_json_decode_error(e, manufacturer, model, retry_count)

        return json_data, storage_path

    def _get_library_model(self, manufacturer: str, model: str) -> LibraryModel:
        """Retrieve model info, or raise an error if not found."""
        model_info = self.model_infos.get(f"{manufacturer}/{model}")
        if not model_info:
            raise LibraryLoadingError(f"Model not found in library: {manufacturer}/{model}")
        return model_info

    async def _needs_update(
        self,
        model_info: LibraryModel,
        manufacturer: str,
        model: str,
        model_path: str,
        force_update: bool,
    ) -> bool:
        """Check if the model needs to be updated."""
        if force_update:
            return True

        path_exists = await self.hass.async_add_executor_job(os.path.exists, model_path)
        if not path_exists:
            return True

        existing_hash = self.profile_hashes.get(f"{manufacturer}/{model}")
        new_hash = model_info.get("hash")
        return existing_hash != new_hash

    async def _download_profile_with_retry(
        self,
        manufacturer: str,
        model: str,
        storage_path: str,
        model_path: str,
    ) -> None:
        """Attempt to download the profile, with retry logic and error handling."""
        try:
            model_info = self._get_library_model(manufacturer, model)
            model_hash = str(model_info.get("hash"))
            callback = partial(self.download_profile, manufacturer, model, storage_path, model_hash)
            await self.download_with_retry(callback)
            self.profile_hashes[f"{manufacturer}/{model}"] = model_hash
            await self.hass.async_add_executor_job(self._write_profile_hashes, self.profile_hashes)
        except ProfileDownloadError as e:
            path_exists, storage_path_exists = await self.hass.async_add_executor_job(
                self._profile_paths_exist,
                model_path,
                storage_path,
            )
            if not path_exists:
                if storage_path_exists:
                    await self.hass.async_add_executor_job(shutil.rmtree, storage_path)  # pragma: no cover
                raise e
            _LOGGER.debug("Failed to download profile, falling back to local profile")

    @staticmethod
    def _profile_paths_exist(model_path: str, storage_path: str) -> tuple[bool, bool]:
        """Check profile paths from the executor."""
        return os.path.exists(model_path), os.path.exists(storage_path)

    async def _load_model_json(self, model_path: str) -> dict[str, Any]:
        """Load the JSON data from the model file."""

        def _load_json() -> dict[str, Any]:
            with open(model_path) as f:
                return cast(dict[str, Any], json.load(f))

        return await self.hass.async_add_executor_job(_load_json)

    async def _handle_json_decode_error(
        self,
        error: JSONDecodeError,
        manufacturer: str,
        model: str,
        retry_count: int,
    ) -> tuple[dict[str, Any], str] | None:
        """Handle JSON decode errors with retry logic."""
        _LOGGER.error("model.json file is not valid JSON for manufacturer: %s, model: %s", manufacturer, model)
        if retry_count < 2:
            _LOGGER.debug("Retrying to load model.json file")
            return await self.load_model(manufacturer, model, True, retry_count + 1)
        raise LibraryLoadingError("Failed to load model.json file") from error

    def get_storage_path(self, manufacturer: str, model: str) -> str:
        """Retrieve the storage path for a given manufacturer and model."""
        return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, manufacturer, model))

    async def download_with_retry(
        self,
        callback: Callable[[], Coroutine[Any, Any, dict[str, Any] | None]],
    ) -> dict[str, Any] | None:
        """Download a file from a remote endpoint with retries"""
        max_retries = 3
        retry_count = 0

        while retry_count < max_retries:
            try:
                return await callback()
            except (ClientError, TimeoutError, ProfileDownloadError) as e:
                _LOGGER.debug(e)
                retry_count += 1
                if retry_count == max_retries:
                    raise ProfileDownloadError(
                        f"Failed to download even after {max_retries} retries, falling back to local copy",
                    ) from e

                await asyncio.sleep(self.retry_timeout)
                _LOGGER.warning("Failed to download, retrying... (Attempt %d of %d)", retry_count + 1, max_retries)
        return None  # pragma: no cover

    async def download_profile(self, manufacturer: str, model: str, storage_path: str, model_hash: str) -> None:
        """
        Download the profile from Github using the Powercalc download API
        Saves the profile to manufacturer/model directory in .storage/powercalc_profiles folder
        """

        _LOGGER.debug("Downloading profile: %s/%s from github", manufacturer, model)

        endpoint = f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}"

        def _save_file(data: bytes, directory: str) -> None:
            """Save file from Github to local storage directory"""
            path = os.path.join(storage_path, directory)
            os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(path, "wb") as f:
                f.write(data)

        session = async_get_clientsession(self.hass)

        try:
            async with asyncio.timeout(TIMEOUT_SECONDS):
                async with session.get(endpoint, params={"hash": model_hash}) as resp:
                    if resp.status != 200:
                        raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}")
                    resources = await resp.json()

                await self.hass.async_add_executor_job(lambda: os.makedirs(storage_path, exist_ok=True))

                # Download the files
                for resource in resources:
                    url = resource.get("url")
                    async with session.get(url) as resp:
                        if resp.status != 200:
                            raise ProfileDownloadError(f"Failed to download github URL: {url}")

                        contents = await resp.read()
                        await self.hass.async_add_executor_job(_save_file, contents, resource.get("path"))
        except (TimeoutError, aiohttp.ClientError) as e:
            raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}") from e

    def _get_profile_hashes_path(self) -> str:
        """Retrieve the local storage path for the profile hashes file."""
        return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes"))

    def _load_profile_hashes(self) -> dict[str, str]:
        """Load profile hashes from local storage"""

        path = self._get_profile_hashes_path()
        if not os.path.exists(path):
            return {}

        with open(path) as f:
            return json.load(f)  # type: ignore[no-any-return]

    def _write_profile_hashes(self, hashes: dict[str, str]) -> None:
        """Write profile hashes to local storage"""

        path = self._get_profile_hashes_path()
        with open(path, "w") as json_file:
            json.dump(hashes, json_file, indent=4)
