import asyncio
from logging import Logger, getLogger
from typing import Awaitable, Callable, Dict, Optional, List
from arcnagios.utils import Alpha, Result, ResultOk, ResultError, map_option
from arcnagios.arcutils import Arcstat, J_UNDEFINED, jobstate_of_str, ParseError

class ArcClientError:

    returncode: int
    prog: str
    stdout: str

    def __init__(self, returncode: int, prog: str, stdout: str, stderr: str):
        self.returncode = returncode
        self.prog = prog
        self.output = stdout  # reconsider usage
        self._stderr = stderr # logged by this module

_DEFAULT_LOG = getLogger(__name__)

def _parse_nothing(
        returncode: int, stdout: str, stderr: str, log: Logger
    ) -> Result[None, None]:
    # pylint: disable=W0613

    if returncode == 0:
        return ResultOk(None)
    else:
        return ResultError(None)

class ArcstatResponse:
    jobs: Dict[str, Arcstat]

    def __init__(self, jobs: Dict[str, Arcstat]):
        self.jobs = jobs

def _convert_arcstat(jobid, jobstat):
    if jobstat['State'] == 'Undefined':
        state = J_UNDEFINED
        specific_state = None
    elif 'Specific state' in jobstat:
        state = jobstate_of_str(jobstat['State'])
        specific_state = jobstat['Specific state']
    else:
        raise ParseError('Missing "State" or "Specific state" for %s.'
                         % jobid)
    return Arcstat(state = state, specific_state = specific_state,
                   exit_code = map_option(int, jobstat.get('Exit code')),
                   submitted = jobstat.get('Submitted'),
                   job_error = jobstat.get('Job Error'))

def _parse_arcstat_response(
        returncode: int, stdout: str, stderr: str, log: Logger
    ) -> Result[ArcstatResponse, None]:
    # pylint: disable=W0613

    # A return code 1 can mean "No jobs found, try later" or a real error, but
    # at least the following should trigger if the system is unable to launch
    # the command.
    if returncode > 1:
        return ResultError(None)

    jobstats = {}
    line_number = 0

    def parse_error(msg):
        log.warning('Unparsed line %s: %s', line_number, msg)

    jobid: Optional[str] = None
    jobfield: Optional[str] = None
    jobstat: Dict[str, str] = {}

    def flush():
        nonlocal jobid, jobstat, jobstats
        if jobid is None:
            return
        if not jobstat:
            log.warning('Not further information for %s', jobid)
            jobid = None
            return
        jobstats[jobid] = _convert_arcstat(jobid, jobstat)
        jobid = None
        jobstat = {}

    for line in stdout.split('\n'):
        line_number += 1

        if line.startswith('No jobs') or line.startswith('Status of '):
            break
        if line == '':
            continue

        if line.startswith('Job:'):
            flush()
            jobid = line[4:].strip()
            log.debug('Found %s', jobid)
            jobstat = {}
            jobfield = None
        elif line.startswith('Warning:'):
            log.warning('%s', line)
        elif line.startswith('  '):
            if jobfield is None:
                parse_error('Continuation line %r before job field.')
                continue
            jobstat[jobfield] += '\n' + line
        elif line.startswith(' '):
            kv = line.strip()
            try:
                jobfield, v = kv.split(':', 1)
                if jobid is None:
                    parse_error('Missing "Job: ..." header before %r' % kv)
                    continue
                jobstat[jobfield] = v.strip()
            except ValueError:
                parse_error('Expecting "<key>: <value>", got %r' % line)
                continue
        else:
            parse_error('Unrecognized output %r' % line)
    flush()

    return ResultOk(ArcstatResponse(jobstats))

async def _call_arc_process(
        prog: str,
        args: List[str],
        parse: Callable[[int, str, str, Logger], Result[Alpha, None]],
        log: Logger
    ) -> Result[Alpha, ArcClientError]:

    # FIXME: Is there a proper way to pass a modified environment to
    # asyncio.subprocess.create_subprocess_exec?
    prog_utc, args_utc = "/usr/bin/env", ["TZ=UTC", prog] + args

    log.debug('Calling %s %s', prog, ' '.join(args))
    # pylint: disable=E1101
    # asyncio.subprocess is not subprocess
    proc = await asyncio.subprocess.create_subprocess_exec(
            prog_utc, *args_utc,
            stdin=asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE)
    stdout_, stderr_ = await proc.communicate()
    stdout = stdout_.decode('utf-8')
    stderr = stderr_.decode('utf-8')
    returncode = proc.returncode
    assert returncode is not None

    result = parse(returncode, stdout, stderr, log)
    if isinstance(result, ResultOk):
        return result
    for line in stderr.split('\n'):
        if line:
            log.error('[%s/%d] %s', prog, proc.pid, line)
    return ResultError(ArcClientError(returncode, prog, stdout, stderr))

# Jobs

def arcget(
        job_id: str,
        top_output_dir: Optional[str] = None,
        timeout: Optional[float] = None,
        log: Logger = _DEFAULT_LOG,
    ) -> Awaitable[Result[None, ArcClientError]]:

    args = [job_id]
    if not timeout is None:
        args += ['-t', str(int(timeout + 0.5))]
    if not top_output_dir is None:
        args += ['-D', top_output_dir]
    return _call_arc_process('arcget', args, _parse_nothing, log)

def arcstat(
        jobids: Optional[List[str]] = None,
        timeout: float = 5.0,
        show_unavailable: bool = False,
        log: Logger = _DEFAULT_LOG,
    ) -> Awaitable[Result[ArcstatResponse, ArcClientError]]:

    args = ['-l', '-t', str(int(timeout + 0.5))]
    if jobids is None:
        args.append('-a')
    else:
        args.extend(jobids)
    if show_unavailable:
        args.append('-u')
    return _call_arc_process('arcstat', args, _parse_arcstat_response, log)

def arcclean(
        job_id: str, force: bool = False, timeout: Optional[float] = None,
        log: Logger = _DEFAULT_LOG,
    ) -> Awaitable[Result[None, ArcClientError]]:

    args = [job_id]
    if not timeout is None:
        args += ['-t', str(int(timeout + 0.5))]
    if force:
        args.append('-f')
    return _call_arc_process('arcclean', args, _parse_nothing, log)

def arckill(
        job_id: str, timeout: Optional[float] = None,
        log: Logger = _DEFAULT_LOG,
    ) -> Awaitable[Result[None, ArcClientError]]:

    args = [job_id]
    if not timeout is None:
        args += ['-t', str(int(timeout + 0.5))]
    return _call_arc_process('arckill', args, _parse_nothing, log)

# Storage

def arcrm(
        url: str, force: bool, timeout: float,
        log: Logger = _DEFAULT_LOG,
    ) -> Awaitable[Result[None, ArcClientError]]:

    args = [url]
    if not timeout is None:
        args += ['-t', str(int(timeout + 0.5))]
    if force:
        args.append('-f')
    return _call_arc_process('arcrm', args, _parse_nothing, log)
