import os
from collections.abc import Iterable
from logging import getLogger
from warnings import catch_warnings

from ubiquerg import size

from .. import Project, Sample
from .const import (
    ALL_INPUTS_KEY,
    INPUT_FILE_SIZE_KEY,
    MISSING_KEY,
    PROP_KEY,
    REQUIRED_INPUTS_KEY,
    SAMPLES_KEY,
    SIZING_KEY,
    TANGIBLE_KEY,
)
from .schema import read_schema
from .validation import _get_attr_values, _validate_sample_object

_LOGGER = getLogger(__name__)


def inspect_project(
    p: Project, sample_names: list[str] | None = None, max_attr: int = 10
) -> None:
    """
    Print inspection info: Project or, if sample_names argument is provided, matched samples.

    Args:
        p: Project to inspect
        sample_names: List of samples to inspect
        max_attr: Max number of sample attributes to display
    """
    if sample_names:
        samples = p.get_samples(sample_names)
        if not samples:
            print("No samples matched by names: {}".format(sample_names))
            return
        for s in samples:
            print(s.__str__(max_attr=max_attr))
            print("\n")
        return
    print(p)
    return


def get_input_files_size(
    sample: Sample, schema: str | list[dict]
) -> dict[str, list[str] | set[str] | float]:
    """
    Determine which of this Sample's required attributes/files are missing and calculate sizes.

    The names of the attributes that are required and/or deemed as inputs
    are sourced from the schema, more specifically from required_input_attrs
    and input_attrs sections in samples section. Note, this function does
    perform actual Sample object validation with jsonschema.

    Args:
        sample: Sample to investigate
        schema: Schema dict to validate against or a path to one

    Returns:
        Dictionary with validation data, i.e missing, required_inputs,
        all_inputs, input_file_size

    Raises:
        ValidationError: If any required sample attribute is missing
    """

    def _compute_input_file_size(inputs: Iterable[str]) -> float:
        """Compute total size of input files."""
        with catch_warnings(record=True) as w:
            total_bytes = sum(
                size(f, size_str=False) or 0.0
                for f in inputs
                if f != "" and f is not None
            )
            if w:
                _LOGGER.warning(
                    f"{len(w)} input files missing, job input size was "
                    f"not calculated accurately"
                )
        return total_bytes / (1024**3)

    if isinstance(schema, str):
        schema = read_schema(schema)

    # first, validate attrs existence using jsonschema
    _validate_sample_object(schemas=schema, sample=sample)

    all_inputs = set()
    required_inputs = set()
    schema = schema[-1]  # use only first schema, in case there are imports
    sample_schema_dict = schema[PROP_KEY][SAMPLES_KEY]["items"]
    if SIZING_KEY in sample_schema_dict:
        all_inputs.update(_get_attr_values(sample, sample_schema_dict[SIZING_KEY]))
    if TANGIBLE_KEY in sample_schema_dict:
        required_inputs = set(
            _get_attr_values(sample, sample_schema_dict[TANGIBLE_KEY])
        )
        all_inputs.update(required_inputs)
    with catch_warnings(record=True) as w:
        # input_file_size = sum(
        #     [
        #         size(f, size_str=False) or 0.0
        #         for f in all_inputs
        #         if f != "" and f is not None
        #     ]
        # ) / (1024**3)
        input_file_size = _compute_input_file_size(all_inputs)
        if w:
            _LOGGER.warning(
                f"{len(w)} input files missing, job input size was "
                f"not calculated accurately"
            )

    return {
        MISSING_KEY: [i for i in required_inputs if not os.path.exists(i)],
        REQUIRED_INPUTS_KEY: required_inputs,
        ALL_INPUTS_KEY: all_inputs,
        INPUT_FILE_SIZE_KEY: input_file_size,
    }
