import subprocess
import os
import shutil
from pathlib import Path
import math
import typing as T
import hashlib
import urllib.request
import urllib.error
import socket
import zipfile
import tarfile

try:
    import psutil
except ImportError:
    psutil = None
    # pip install psutil will improve CPU utilization.

from .config import read_config
from .base import get_simsize

git = shutil.which("git")

Pathlike = T.Union[str, Path]

__all__ = ["gitrev", "get_cpu_count", "get_mpi_count"]


def gitrev() -> str:
    if not git:
        return ""

    return subprocess.check_output([git, "rev-parse", "--short", "HEAD"], universal_newlines=True).strip()


def get_cpu_count(force: int = None) -> int:
    if force:
        max_cpu = force
        extradiv = 1
    else:
        max_cpu = None
        # without psutil, hyperthreaded CPU may overestimate physical count by factor of 2 (or more)
        if psutil is not None:
            max_cpu = psutil.cpu_count(logical=False)
            extradiv = 1
            if max_cpu is None:
                max_cpu = psutil.cpu_count()
                extradiv = 2
        if max_cpu is None:
            max_cpu = os.cpu_count()
            extradiv = 2

    return max_cpu // extradiv


def get_mpi_count(path: Pathlike, force: int = None) -> int:

    path = Path(path).expanduser()

    max_cpu = get_cpu_count(force)

    # %% config.nml file or directory or simsize.h5?
    if path.is_dir():
        size = get_simsize(path)
    elif path.is_file():
        if path.suffix in (".h5", ".nc", ".dat"):
            size = get_simsize(path)
        elif path.suffix in (".ini", ".nml"):
            params = read_config(path)
            # OK to use indat_size because we're going to run a sim on this machine
            size = get_simsize(params["indat_size"])
    else:
        raise FileNotFoundError(f"{path} is not a file or directory")

    mpi_count = 1
    if size[2] == 1:
        # 2D sim
        for i in range(max_cpu, 2, -1):
            mpi_count = max(math.gcd(size[1], i), mpi_count)
            if i < mpi_count:
                break
    else:
        # 3D sim
        for i in range(max_cpu, 2, -1):
            mpi_count = max(math.gcd(size[2], i), mpi_count)
            if i < mpi_count:
                break

    return mpi_count


def url_retrieve(url: str, outfile: Pathlike, filehash: T.Sequence[str] = None, overwrite: bool = False):
    """
    Parameters
    ----------
    url: str
        URL to download from
    outfile: pathlib.Path
        output filepath (including name)
    filehash: tuple of str, str
        hash type (md5, sha1, etc.) and hash
    overwrite: bool
        overwrite if file exists
    """
    outfile = Path(outfile).expanduser().resolve()
    if outfile.is_dir():
        raise ValueError("Please specify full filepath, including filename")
    # need .resolve() in case intermediate relative dir doesn't exist
    if overwrite or not outfile.is_file():
        outfile.parent.mkdir(parents=True, exist_ok=True)
        try:
            urllib.request.urlretrieve(url, str(outfile))
        except (socket.gaierror, urllib.error.URLError) as err:
            raise ConnectionError(f"could not download {url} due to {err}")

    if filehash:
        if not file_checksum(outfile, filehash[0], filehash[1]):
            raise ValueError(f"Hash mismatch: {outfile}")


def file_checksum(fn: Path, mode: str, filehash: str) -> bool:
    h = hashlib.new(mode)
    h.update(fn.read_bytes())
    return h.hexdigest() == filehash


def extract_zip(fn: Pathlike, outpath: Pathlike, overwrite: bool = False):
    outpath = Path(outpath).expanduser().resolve()
    # need .resolve() in case intermediate relative dir doesn't exist
    if outpath.is_dir() and not overwrite:
        return

    fn = Path(fn).expanduser().resolve()
    with zipfile.ZipFile(fn) as z:
        z.extractall(str(outpath.parent))


def extract_tar(fn: Pathlike, outpath: Pathlike, overwrite: bool = False):
    outpath = Path(outpath).expanduser().resolve()
    # need .resolve() in case intermediate relative dir doesn't exist
    if outpath.is_dir() and not overwrite:
        return

    fn = Path(fn).expanduser().resolve()
    if not fn.is_file():
        # tarfile gives confusing error on missing file
        raise FileNotFoundError(fn)
    with tarfile.open(fn) as z:
        z.extractall(str(outpath.parent))