from contextlib import contextmanager
import json
import os
import subprocess
import sys
import xml.sax
import xml.sax.handler

from typing import Any, Dict, Iterator, List, Set


@contextmanager
def log(msg: str) -> Iterator[None]:
    print(msg, file=sys.stderr, end='', flush=True)
    try:
        yield
    finally:
        print('\r', file=sys.stderr, end='', flush=True)


class ParseNixStoreQueryGraphML(xml.sax.handler.ContentHandler):
    def __init__(self) -> None:
        super().__init__()
        self.roots: Set[str] = set()
        self.non_roots: Set[str] = set()
        self.deps: Dict[str, List[str]] = {}

    def startElement(self, name: str, attrs: Any) -> None:
        if name == "edge":
            source = attrs.getValue("source")
            target = attrs.getValue("target")
            self.non_roots.add(target)
            self.deps.setdefault(source, []).append(target)
            if target in self.roots:
                self.roots.remove(target)
            if source not in self.non_roots:
                self.roots.add(source)


def getDeps(drv: str) -> ParseNixStoreQueryGraphML:
    with log("Loading dependency tree..."):
        with subprocess.Popen(
                ["nix-store", "--query", "--graphml", drv], stdout=subprocess.PIPE) as process:
            parser = ParseNixStoreQueryGraphML()
            assert process.stdout
            xml.sax.parse(process.stdout, parser)
            assert process.wait() == 0
        return parser


def getDrvInfo(drv: str) -> Any:
    with log(f"Loading {drv}..."):
        with subprocess.Popen(
                ["nix", "--experimental-features", "nix-command",
                    "show-derivation", "/nix/store/" + drv],
                stdout=subprocess.PIPE) as process:
            assert process.stdout
            info = json.load(process.stdout)
            assert process.wait() == 0
        return info


def allBuilt(drv: str) -> bool:
    # TODO: How to pin outputs one at a time?
    # Currently, we only pin when all outputs are available.
    # It would be better to pin the outputs we've got.
    return all(os.path.exists(o['path']) for d in getDrvInfo(
        drv).values() for o in d['outputs'].values())


def isDrv(thing: str) -> bool:
    return thing.endswith(".drv")


def removesuffix(s: str, suf: str) -> str:
    return s[:-len(suf)] if s.endswith(suf) else s


def pin(drv: str) -> None:
    outPath = os.path.join(sys.argv[2], removesuffix(drv, ".drv"))
    if not os.path.exists(outPath):
        subprocess.run(["nix-store", "--realise", "--add-root",
                       outPath, "/nix/store/" + drv], check=True)


def pinBuiltThings(thing: str,
                   done: Set[str],
                   deps: Dict[str, List[str]]) -> None:
    if thing in done:
        return
    done.add(thing)
    if not isDrv(thing) or allBuilt(thing):
        pin(thing)
        return
    for dep in deps.get(thing, []):
        pinBuiltThings(dep, done, deps)


dep_graph = getDeps(sys.argv[1])
for root in dep_graph.roots:
    pinBuiltThings(root, set(), dep_graph.deps)