#!/usr/bin/env python
from concurrent.futures import ThreadPoolExecutor
import subprocess
import textwrap
import argparse
import toml
import json
import sys

from typing import Dict, Any, Tuple, List


class Package:
    def __init__(self, attrs: Dict[str, Any]) -> None:
        self.attrs = attrs
        self.name = attrs["name"]
        self.source = self.attrs["source"]

    def fetch(self) -> Tuple["Package", subprocess.CompletedProcess]:
        raise NotImplementedError()

    def expression(self, output: str) -> str:
        raise NotImplementedError()


class UrlPackage(Package):
    def fetch(self) -> Tuple[Package, subprocess.CompletedProcess]:
        return (
            self,
            subprocess.run(
                [
                    "nix-prefetch-url",
                    "--unpack",
                    self.source["url"],
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            ),
        )

    def expression(self, output: str) -> str:
        sha256 = output.rstrip()
        return textwrap.dedent("""
            %s = super.%s.overridePythonAttrs (
              _: {
                src = pkgs.fetchzip {
                  url = "%s";
                  sha256 = "%s";
                };
              }
            );""" % (self.name, self.name, self.source["url"], sha256))


class GitPackage(Package):
    def fetch(self) -> Tuple[Package, subprocess.CompletedProcess]:
        reference = self.source.get("resolved_reference", self.source["reference"])

        return (
            self,
            subprocess.run(
                [
                    "nix-prefetch-git",
                    "--fetch-submodules",
                    "--url",
                    self.source["url"],
                    "--rev",
                    reference,
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            ),
        )

    def expression(self, output: str) -> str:
        meta = json.loads(output)
        return textwrap.dedent("""
            %s = super.%s.overridePythonAttrs (
              _: {
                src = pkgs.fetchgit {
                  url = "%s";
                  rev = "%s";
                  sha256 = "%s";
                };
              }
            );""" % (self.name, self.name, meta["url"], meta["rev"], meta["sha256"]))


def parse_args() -> argparse.Namespace:
    argparser = argparse.ArgumentParser(description="Poetry2nix CLI")

    subparsers = argparser.add_subparsers(dest="subcommand")
    subparsers.required = True

    parser_lock = subparsers.add_parser("lock", help="Generate overrides for git hashes",)
    parser_lock.add_argument(
        "--lock", default="poetry.lock", help="Path to input poetry.lock",
    )
    parser_lock.add_argument(
        "--out", default="poetry-git-overlay.nix", help="Output file",
    )
    return argparser.parse_args()


def indent(expr: str, spaces: int = 2) -> str:
    i = " " * spaces
    return "\n".join([(i if l != "" else "") + l for l in expr.split("\n")])


def main() -> None:
    args = parse_args()

    with open(args.lock) as lockf:
        lock = toml.load(lockf)

    pkgs: List[Package] = []
    for pkg in lock["package"]:
        if "source" in pkg:
            source_type = pkg["source"]["type"]
            if source_type == "git":
                pkgs.append(GitPackage(pkg))
            elif source_type == "url":
                pkgs.append(UrlPackage(pkg))

    with ThreadPoolExecutor() as e:
        futures = []

        for pkg in pkgs:
            futures.append(e.submit(pkg.fetch))

        lines = [
            "{ pkgs }:",
            "self: super: {",
        ]

        for f in futures:
            package, p = f.result()
            if p.returncode != 0:
                sys.stderr.write(p.stderr)
                sys.stderr.flush()
                exit(p.returncode)
            expr = package.expression(p.stdout)
            lines.append(indent(expr))

        lines.extend(["", "}", ""])

        expr = "\n".join(lines)

    with open(args.out, "w") as fout:
        fout.write(expr)

    print(f"Wrote {args.out}")


if __name__ == "__main__":
    main()