import logging
from typing import Any, Optional

import pandas as pd
from pydantic import BaseModel, ConfigDict, Field, field_validator

from rpmeta.constants import ALL_FEATURES
from rpmeta.helpers import to_minutes_rounded

logger = logging.getLogger(__name__)


class HwInfo(BaseModel):
    """
    Hardware information of the build system.

    This model represents the hardware specifications of the machine used for building packages.
    It includes CPU details, architecture, and memory information.
    """

    cpu_model_name: str = Field(
        description=(
            "The full name/model of the CPU (e.g., 'Intel(R) Core(TM) i7-8550U CPU @ 1.80GHz')"
        ),
    )
    cpu_arch: str = Field(
        description="CPU architecture (e.g., 'x86_64', 'aarch64')",
    )
    cpu_model: str = Field(
        description="CPU model number (e.g., '142', '49')",
        examples=["85", "142", "49"],
    )
    cpu_cores: int = Field(
        description="Number of CPU cores",
        gt=0,
    )
    ram: float = Field(
        description="Total RAM in Gi (gibibytes)",
        gt=0,
    )
    swap: float = Field(
        description="Total swap space in Gi (gibibytes)",
        ge=0,
    )

    model_config = ConfigDict(
        validate_assignment=True,
    )

    @staticmethod
    def _convert_human_readable_to_gi(value: str) -> float:
        """
        Convert human-readable memory units to Gi (gibibytes).

        Supports the following units:
        - B (bytes)
        - Ki (kibibyte)
        - Mi (mebibyte)
        - Gi (gibibyte)
        - Ti (tebibyte)
        - Pi (pebibyte)

        Args:
            value: Memory value with unit (e.g., '54Gi', '154Mi') or plain number
             in KB (/bin/free default)

        Returns:
            Memory value in Gi (gibibytes)
        """
        value = value.strip().replace(",", ".")

        unit_to_gi = {
            "B": 1 / (1024**3),
            "Ki": 1 / (1024**2),
            "Mi": 1 / 1024,
            "Gi": 1,
            "Ti": 1024,
            "Pi": 1024 * 1024,
        }

        for unit, factor in unit_to_gi.items():
            if value.endswith(unit):
                numeric_value = float(value.rstrip(unit))
                return round(numeric_value * factor, 1)

        # If no unit found, assume it's in KB (/bin/free default behavior)
        return round(float(value) / (1024**2), 1)

    @classmethod
    def parse_from_lscpu(cls, content: str) -> "HwInfo":
        logger.debug("lscpu output: %s", content)
        hw_info: dict[str, float | int | str] = {}
        for line in content.splitlines():
            if line.startswith("Model name:"):
                hw_info["cpu_model_name"] = line.split(":")[1].strip()
            elif line.startswith("Architecture:"):
                hw_info["cpu_arch"] = line.split(":")[1].strip()
            elif line.startswith("Model:"):
                hw_info["cpu_model"] = line.split(":")[1].strip()
            elif line.startswith("CPU(s):"):
                hw_info["cpu_cores"] = int(line.split(":")[1].strip())
            elif line.startswith("Mem:"):
                hw_info["ram"] = cls._convert_human_readable_to_gi(line.split()[1])
            elif line.startswith("Swap:"):
                hw_info["swap"] = cls._convert_human_readable_to_gi(line.split()[1])

        if hw_info.get("cpu_model") is None:
            hw_info["cpu_model"] = "unknown"

        logger.debug("Extracted hardware info: %s", hw_info)
        return cls(**hw_info)

    @field_validator("cpu_model", mode="before")
    @classmethod
    def validate_cpu_model(cls, v):
        if v is None:
            # setting to unknown for the model to be able to handle missing cpu models
            return "unknown"

        return v


class InputRecord(BaseModel):
    """
    Input record for the model training and prediction.

    This record is used to pass the input data to the model for training or prediction.
    Exactly the same data structure is provided by user for the API or CLI.

    ## Example
    ```json
    {
      "package_name": "rust-winit",
      "epoch": 0,
      "version": "0.30.8",
      "mock_chroot": "fedora-41-x86_64",
      "hw_info": {
        "cpu_model_name": "Intel Xeon Processor (Cascadelake)",
        "cpu_arch": "x86_64",
        "cpu_model": "85",
        "cpu_cores": 6,
        "ram": 14.6,
        "swap": 8.0
      }
    }
    ```
    """

    package_name: str = Field(
        description="Name of the RPM package to build",
        examples=["rust-winit", "python-numpy"],
    )
    epoch: int = Field(
        description="Package epoch (usually 0)",
        ge=0,
        examples=[0, 1],
    )
    version: str = Field(
        description="Package version string",
        examples=["0.30.8", "1.24.3"],
    )
    hw_info: HwInfo = Field(
        description="Hardware information of the build system",
    )
    mock_chroot: str = Field(
        description="Mock chroot in format '<distro>-<version>-<arch>' (e.g., 'fedora-41-x86_64')",
        examples=["fedora-41-x86_64", "centos-stream-9-x86_64"],
    )

    @property
    def neva(self) -> str:
        """
        Name, Epoch, Version, Architecture; Release is (intentionally) missing in data.
        """
        return f"{self.package_name}-{self.epoch}:{self.version}-{self.os_arch}"

    @property
    def os(self) -> str:
        return self.mock_chroot.rsplit("-", 2)[0]

    @property
    def os_family(self) -> str:
        return self.os.rsplit("-")[0]

    @property
    def os_version(self) -> Optional[str]:
        if self.mock_chroot is None:
            return None

        return self.mock_chroot.rsplit("-", 2)[1]

    @property
    def os_arch(self) -> str:
        return self.mock_chroot.rsplit("-", 2)[2]

    def to_model_dict(self) -> dict[str, Any]:
        """
        Convert the record to dictionary with data that the _trained model_ expects.
        """
        return {
            "package_name": self.package_name,
            "epoch": self.epoch,
            "version": self.version,
            "os": self.os,
            "os_family": self.os_family,
            "os_version": self.os_version,
            "os_arch": self.os_arch,
            "hw_info": self.hw_info.model_dump(),
        }

    def to_data_frame(
        self,
        category_maps: dict[str, list[str]],
        category_dtypes: Optional[dict[str, pd.CategoricalDtype]] = None,
    ) -> pd.DataFrame:
        """
        Convert the record to a pandas DataFrame that the model understands.
        This is used for prediction.

        Args:
            category_maps: Mapping of column name to list of categories.
            category_dtypes: Pre-created CategoricalDtype objects.
                Falls back to building from category_maps if not provided.
        """
        df = pd.json_normalize(self.model_dump())
        df["os"] = self.os
        df["os_family"] = self.os_family
        df["os_version"] = self.os_version
        df["os_arch"] = self.os_arch

        for col, cat_list in category_maps.items():
            if category_dtypes is None:
                dtype = pd.CategoricalDtype(categories=cat_list, ordered=False)
            else:
                dtype = category_dtypes[col]

            df[col] = df[col].astype(dtype)

        df["hw_info.ram"] = df["hw_info.ram"].astype("float32")
        df["hw_info.swap"] = df["hw_info.swap"].astype("float32")

        # ensure all features are in the expected order for the model
        return df[ALL_FEATURES]


class Record(InputRecord):
    """
    A record of a successful build in build system in dataset.

    This extends InputRecord to include the actual build duration,
    which is used as the target variable for model training.
    """

    build_duration: int = Field(
        description="Actual build duration in seconds",
        gt=0,
        examples=[5, 60, 720],
    )

    def to_model_dict(self) -> dict[str, Any]:
        """
        Convert the record to dictionary that the model _to be trained_ understands + has the
         target feature.
        """
        return {
            **super().to_model_dict(),
            # the fetched data is in seconds, but the model gives better results in minutes
            "build_duration": to_minutes_rounded(self.build_duration),
        }
