import os
import re
import json
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional

from PyExpUtils.utils.cmdline import flagString

"""doc
Takes an integer number of hours and returns a well-formatted time string.
```python
time = hours(3)
print(time) # -> '2:59:59
```
"""
def hours(n: int):
    return f'{n-1}:59:59'

"""doc
Takes an integer number of gigabytes and returns a well-formatted memory string.
```python
memory = gb(4)
print(memory) # -> '4G'
```
"""
def gb(n: int):
    return f'{n}G'


@dataclass
class SingleNodeOptions:
    account: str
    time: str
    cores: int
    mem_per_core: str | float

    # task management args
    sequential: int = 1
    threads_per_task: int = 1

    # job reporting args
    log_path: str = "$SCRATCH/job_output_%j.txt"

@dataclass
class MultiNodeOptions:
    account: str
    time: str
    cores: int
    mem_per_core: str | float

    # task management args
    sequential: int = 1

    # job reporting args
    log_path: str = "$SCRATCH/job_output_%j.txt"

# ----------------
# -- Validation --
# ----------------
def check_account(account: str):
    assert account.startswith('rrg-') or account.startswith('def-')
    assert not account.endswith('_cpu') and not account.endswith('_gpu')

def check_time(time: str):
    assert isinstance(time, str)

    # while technically slurm is more permissive, I find being more explicit removes
    # some common footguns. Example the "int:int" format is oft misunderstood as "hours:minutes"

    # "hour:minute:second"
    h_m_s = re.match(r'^\d+:\d+:\d+$', time)

    # "days-hours"
    d_h = re.match(r'^\d+-\d+$', time)

    # "days-hours:minutes:seconds"
    d_h_m_s = re.match(r'^\d+-\d+:\d+:\d+$', time)

    assert h_m_s or d_h or d_h_m_s

def normalize_memory(memory: float | str) -> str:
    if isinstance(memory, (float, int)):
        mbs = int(memory * 1024)
        memory = f'{mbs}M'

    assert isinstance(memory, str)
    assert re.match(r'^\d+[G|M|K]$', memory)
    return memory

def shared_validation(options: SingleNodeOptions | MultiNodeOptions):
    check_account(options.account)
    check_time(options.time)
    options.mem_per_core = normalize_memory(options.mem_per_core)


def single_validation(options: SingleNodeOptions):
    shared_validation(options)
    # TODO: validate that the current cluster has nodes that can handle the specified request

def multi_validation(options: MultiNodeOptions):
    shared_validation(options)

def validate(options: SingleNodeOptions | MultiNodeOptions):
    if isinstance(options, SingleNodeOptions): single_validation(options)
    elif isinstance(options, MultiNodeOptions): multi_validation(options)

# ------------------
# -- External API --
# ------------------

def memory_in_mb(memory: str | float) -> float:
    memory = normalize_memory(memory)

    if memory.endswith('M'):
        return int(memory[:-1])

    if memory.endswith('G'):
        return int(memory[:-1]) * 1024

    if memory.endswith('K'):
        return int(memory[:-1]) / 1024

    raise Exception('Unknown memory unit')

def to_cmdline_flags(
    options: SingleNodeOptions | MultiNodeOptions,
    skip_validation: bool = False,
) -> str:
    if not skip_validation:
        validate(options)

    args = [
        ('--account', options.account),
        ('--time', options.time),
        ('--mem-per-cpu', options.mem_per_core),
        ('--output', options.log_path),
    ]

    if isinstance(options, SingleNodeOptions):
        args += [
            ('--ntasks', options.cores),
            ('--nodes', 1),
            ('--cpus-per-task', 1),
        ]

    elif isinstance(options, MultiNodeOptions):
        args += [
            ('--ntasks', options.cores),
            ('--cpus-per-task', 1),
        ]

    return flagString(args)

def fromFile(path: str):
    with open(path, 'r') as f:
        d = json.load(f)

    assert 'type' in d, 'Need to specify scheduling strategy.'
    t = d['type']
    del d['type']

    if t == 'single_node':
        return SingleNodeOptions(**d)

    elif t == 'multi_node':
        return MultiNodeOptions(**d)

    raise Exception('Unknown scheduling strategy')

def buildParallel(executable: str, tasks: Iterable[Any], opts: SingleNodeOptions | MultiNodeOptions, parallelOpts: Dict[str, Any] = {}):
    threads = 1
    if isinstance(opts, SingleNodeOptions):
        threads = opts.threads_per_task

    cores = int(opts.cores / threads)

    parallel_exec = f'srun -N1 -n{threads} --exclusive {executable}'
    if isinstance(opts, SingleNodeOptions):
        parallel_exec = executable

    task_str = ' '.join(map(str, tasks))
    return f'run-parallel --parallel {cores} --exec "{parallel_exec}" --tasks {task_str}'

def schedule(
    script: str,
    opts: Optional[SingleNodeOptions | MultiNodeOptions] = None,
    script_name: str = 'auto_slurm.sh',
    cleanup: bool = True,
    skip_validation: bool = False,
) -> None:
    with open(script_name, 'w') as f:
        f.write(script)

    cmdArgs = ''
    if opts is not None:
        cmdArgs = to_cmdline_flags(opts, skip_validation=skip_validation)

    os.system(f'sbatch {cmdArgs} {script_name}')

    if cleanup:
        os.remove(script_name)
