"""
This module provides utility functions for the atlasopenmagic package.
It includes functions to install packages from an environment file
and to build datasets from sample definitions.
"""

import io
import sys
import re
import warnings
import subprocess
from pathlib import Path
import yaml
import requests
from atlasopenmagic.metadata import get_urls, warn_with_color

warnings.showwarning = warn_with_color
warnings.simplefilter('always', DeprecationWarning)

def install_from_environment(*packages, environment_file=None):
    """
    Install specific packages listed in an environment.yml file via pip.

    Args:
        *packages: Package names to install (e.g., 'coffea', 'dask'). 
        if empty, all packages in the environment.yml will be installed.
        
        environment_file: Path to the environment.yml file. 
        If None, defaults to the environment.yml file contained in our notebooks repository.
    """
    if environment_file is None:
        environment_file = "https://raw.githubusercontent.com/atlas-outreach-data-tools/notebooks-collection-opendata/refs/heads/master/binder/environment.yml"

    is_url = str(environment_file).startswith("http")
    environment_file = Path(
        environment_file) if not is_url else environment_file

    if not is_url:
        if not environment_file.exists():
            raise FileNotFoundError(
                f"Environment file not found at {environment_file}")
        with environment_file.open('r') as file:
            environment_data = yaml.safe_load(file)
    else:
        response = requests.get(environment_file, timeout=100)
        if response.status_code != 200:
            raise ValueError(
                f"Failed to fetch environment file from URL: {environment_file}")
        environment_data = yaml.safe_load(io.StringIO(response.text))

    dependencies = environment_data.get('dependencies', None)

    if dependencies is None:
        raise ValueError(
            f"The environment.yml at {environment_file} is missing a 'dependencies:' section.\n\n"
            "Expected structure:\n"
            "---------------------\n"
            "name: myenv\n"
            "channels:\n"
            "  - conda-forge\n"
            "dependencies:\n"
            "  - python=3.11\n"
            "  - pip:\n"
            "    - mypippackage>=1.0.0\n"
            "---------------------\n")

    conda_packages = []
    pip_packages = []

    if not packages:
        for dep in dependencies:
            if isinstance(dep, str):
                conda_packages.append(dep)
            elif isinstance(dep, dict) and 'pip' in dep:
                pip_list = dep['pip']
                if not isinstance(pip_list, list):
                    raise ValueError(
                        f"Malformed 'pip:' section in {environment_file}.\n\n"
                        "Expected structure:\n"
                        "---------------------\n"
                        "dependencies:\n"
                        "  - pip:\n"
                        "    - package1>=1.0\n"
                        "    - package2>=2.0\n"
                        "---------------------\n"
                    )
                pip_packages.extend(pip_list)
    else:
        for dep in dependencies:
            if isinstance(dep, str):
                for pkg in packages:
                    # Match the package name at the beginning of the string;
                    # this avoids to match two different packages with the same
                    # initial name (e.g. torch, tochvision)
                    base_dep = re.split(r'[=<>]', dep, 1)[0]
                    if base_dep == pkg:
                        conda_packages.append(dep)
            elif isinstance(dep, dict):
                if 'pip' in dep:
                    pip_list = dep['pip']
                    if not isinstance(pip_list, list):
                        raise ValueError(
                            f"Malformed 'pip:' section in {environment_file}.\n\n"
                            "Expected structure:\n"
                            "---------------------\n"
                            "dependencies:\n"
                            "  - pip:\n"
                            "    - package1>=1.0\n"
                            "    - package2>=2.0\n"
                            "---------------------\n")
                    for pip_dep in pip_list:
                        for pkg in packages:
                            if pip_dep.startswith(pkg):
                                pip_packages.append(pip_dep)

    # all_packages = conda_packages + pip_packages
    # Temporarily only install pip packages, to be decided later whether to
    # install conda packages and how
    all_packages = pip_packages

    if all_packages:
        print(f"Installing packages: {all_packages}")

        # Detect if inside a virtualenv and remove the --user flag if so
        in_venv = (
            hasattr(sys, 'real_prefix') or
            (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)
        )

        pip_command = [sys.executable, "-m", "pip", "install", "--upgrade"]
        if not in_venv:
            pip_command.append("--user")

        pip_command += all_packages

        subprocess.run(pip_command, check=True)
        print("Installation complete. " \
        "You may need to restart your Python environment for changes to take effect.")
    else:
        raise ValueError(
            f"No matching packages found for {packages} in {environment_file}.\n\n"
            "Make sure the package names exactly match the beginning of the package entries in the file.\n")


def build_dataset(samples_defs, skim='noskim', protocol='https', cache=False):
    """
    Build a dict of MC samples URLs.

    Args:
        samples_defs (dict[str,dict]: The datasets to be built up with their definitions and
                                      colors. See examples for more info.
        skim (str, optional): The desired skim type. Defaults to 'noskim' for the base,
                              unfiltered dataset. Other examples: 'exactly4lep', '3lep'.
        protocol (str, optional): The desired URL protocol. Can be 'root', 'https', or 'eos'.
                                  Defaults to 'root'.
        cache (bool, optional): use the simplecache mechanism of fsspec to locally cache
                                files instead of streaming them
    """
    out = {}
    for name, info in samples_defs.items():
        urls = []
        for did in info['dids']:
            urls.extend(get_urls(str(did), skim=skim, protocol=protocol, cache=cache))
        sample = {'list': urls}
        if 'color' in info:
            sample['color'] = info['color']
        out[name] = sample
    return out

def build_data_dataset(data_keys, name="Data", color=None, protocol="https", cache=False):
    warnings.warn(
        "The build_data_dataset function is deprecated. "
        "Use build_dataset with the appropriate data definitions instead.",
        DeprecationWarning
    )
    return build_dataset(
        {name: {'dids': ["data"], 'color': color}},
        skim=data_keys,
        protocol=protocol,
        cache=cache
    )

def build_mc_dataset(mc_defs, skim='noskim', protocol='https', cache=False):
    warnings.warn(
        "The build_mc_dataset function is deprecated. "
        "Use build_dataset with the appropriate MC definitions instead.",
        DeprecationWarning
    )
    return build_dataset(mc_defs, skim=skim, protocol=protocol, cache=cache)