from typing import Any, Dict, List, Optional, Sequence, Union, cast

import launchpad as lp
import os
from launchpad.nodes.python.node import PyNode

from tlaunch.lp_k8s import client
from tlaunch.lp_k8s.config import Config
from tlaunch.lp_k8s.source import Source


def launch(programs: Union[lp.Program, Sequence[lp.Program]],
           source: Optional[Source] = None,
           namespace: Optional[str] = None,
           kube_config: Optional[str] = None,
           group_config: Optional[Dict[str, Config]] = None) -> Any:
    # Pretend to be launching with local multiprocess here.
    launch_type = lp.LaunchType.LOCAL_MULTI_PROCESSING

    if not isinstance(programs, Sequence):
        programs = cast(Sequence[lp.Program], [programs])

    if namespace is None:
        if 'TPOD_NAMESPACE' in os.environ:
            namespace = os.environ['TPOD_NAMESPACE']
        else:
            namespace = 'default'

    # NOTE: original launchpad is launching only one single program here
    program = programs[0]

    launch_kubernetes(program, launch_type, source, namespace, kube_config,
                      group_config)


def launch_kubernetes(
        program: lp.Program,
        launch_type: lp.LaunchType = lp.LaunchType.LOCAL_MULTI_PROCESSING,
        source: Optional[Source] = None,
        namespace: Optional[str] = None,
        kube_config: Optional[str] = None,
        group_config: Optional[Dict[str, Config]] = None):
    # Namespace and name of Lpjob should be stored in launch_config
    l_config = {
        "kube_namespace": namespace or "default",
        "lpjob_name": client.validate_name(program.name),
    }

    for label, nodes in program.groups.items():
        for node in nodes:
            node._initialize_context(launch_type, launch_config=l_config)

    for label, nodes in program.groups.items():
        for node in nodes:
            for handle in node._input_handles:
                handle.connect(node, label)

    groups = cast(Dict[str, List[PyNode]], program.groups)

    client.create_lp_job(groups, program.name, source, namespace, kube_config,
                         group_config)
