# Copyright 2024 Pex project contributors.
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import base64
import json
import pkgutil
from collections import Counter
from dataclasses import dataclass
from typing import Any

from pex import toml
from pex.sysconfig import SysPlatform


@dataclass(frozen=True)
class PlatformConfig:
    @classmethod
    def load(
        cls,
        *,
        name: str,
        platform_data: dict[str, Any],
        default_pbs_release: str,
        default_python_version: str
    ) -> PlatformConfig:
        return cls(
            platform=SysPlatform.parse(name),
            pbs_release=platform_data.get("pbs-release", default_pbs_release),
            python_version=platform_data.get("python-version", default_python_version),
            required=platform_data.get("required", True),
        )

    platform: SysPlatform.Value
    pbs_release: str
    python_version: str
    required: bool = True

    @property
    def name(self):
        # type:() -> str
        return str(self.platform)


@dataclass(frozen=True)
class ScieConfig:
    @classmethod
    def load(
        cls,
        *,
        pbs_release: str | None = None,
        python_version: str | None = None,
        encoded_config: str | None = None
    ) -> ScieConfig:
        if encoded_config:
            scie_config = json.loads(base64.urlsafe_b64decode(encoded_config))
        else:
            data = pkgutil.get_data(__name__, "package.toml")
            assert data is not None, f"Expected to find a sibling package.toml file to {__file__}."
            scie_config = toml.loads(data.decode())["scie"]
        default_pbs_release = pbs_release or scie_config["pbs-release"]
        default_python_version = python_version or scie_config["python-version"]
        return cls(
            platforms=tuple(
                PlatformConfig.load(
                    name=platform_name,
                    platform_data=platform_data,
                    default_pbs_release=default_pbs_release,
                    default_python_version=default_python_version,
                )
                for platform_name, platform_data in scie_config["platforms"].items()
            ),
            pex_extras=tuple(scie_config.get("pex-extras", ())),
            extra_lock_args=tuple(scie_config.get("extra-lock-args", ())),
        )

    platforms: tuple[PlatformConfig, ...]
    pex_extras: tuple[str, ...] = ()
    extra_lock_args: tuple[str, ...] = ()

    def current_platform(self) -> PlatformConfig:
        for platform_config in self.platforms:
            if platform_config.platform is SysPlatform.CURRENT:
                return platform_config
        raise KeyError(
            f"This scie configuration does not contain an entry for platform "
            f"{SysPlatform.CURRENT!r}, only the following platforms are defined: "
            f"{', '.join(platform_config.name for platform_config in self.platforms)}"
        )

    def encode(self) -> str:
        pbs_releases: Counter[str] = Counter()
        python_versions: Counter[str] = Counter()
        for platform_config in self.platforms:
            pbs_releases[platform_config.pbs_release] += 1
            python_versions[platform_config.python_version] += 1
        default_pbs_release, _count = pbs_releases.most_common(n=1)[0]
        default_python_version, _count = python_versions.most_common(n=1)[0]

        platforms = {}
        for platform_config in self.platforms:
            data: dict[str, Any] = {}
            if platform_config.pbs_release != default_pbs_release:
                data["pbs-release"] = platform_config.pbs_release
            if platform_config.python_version != default_python_version:
                data["python-version"] = platform_config.python_version
            if not platform_config.required:
                data["required"] = False
            platforms[platform_config.name] = data

        return base64.urlsafe_b64encode(
            json.dumps(
                {
                    "pbs-release": default_pbs_release,
                    "python-version": default_python_version,
                    "pex-extras": self.pex_extras,
                    "extra-lock-args": self.extra_lock_args,
                    "platforms": platforms,
                }
            ).encode()
        ).decode("ascii")
