import base64
import datetime
import importlib
import io
import os
import subprocess
import tarfile
from dataclasses import dataclass
from importlib.metadata import PackageNotFoundError, distribution
from pathlib import Path
from typing import Optional
from uuid import UUID

from packaging.utils import parse_sdist_filename, parse_wheel_filename
from radicli import Arg
from wasabi import msg

from prodigy_teams_broker_sdk.client.envs import BuildImageResponse
from prodigy_teams_pam_sdk import Client
from prodigy_teams_pam_sdk.models import (
    PackageCreating,
    PackageReadingLatest,
    RecipeCreating,
)

from ..cli import cli
from ..errors import CLIError
from ..messages import Messages
from ..ui import loading
from ._state import get_auth_state


@dataclass
class DistPackage:
    """Package to publish, from a wheel, src directory or sdist path"""

    path: Path
    name: str
    version: str

    def __init__(self, path: Path) -> None:
        if path.suffix == ".whl":
            name, version, _, _2 = parse_wheel_filename(path.name)
        else:
            name, version = parse_sdist_filename(path.name)

        self.name = name.replace("-", "_").lower()
        self.version = str(version)
        self.path = path

    @property
    def versioned_name(self) -> str:
        return f"{self.name}-{self.version}"

    @property
    def filename(self) -> str:
        return self.path.name

    def get_bytes(self) -> bytes:
        with self.path.open("rb") as file_:
            data = file_.read()
        return data


@dataclass
class ModulePackage:
    """Package to publish, from an importable module."""

    name: str
    version: str

    def __init__(self, name: str, version_arg: str | None = None) -> None:
        self.name = name
        if version_arg is None:
            utc_now = datetime.datetime.utcnow()
            self.version = f"{utc_now.year}.{utc_now.month}.{utc_now.day}-{utc_now.second + 60 * utc_now.minute + 60 * 60 * utc_now.hour}"
        else:
            self.version = version_arg

    @property
    def versioned_name(self) -> str:
        return f"{self.name}-{self.version}"

    @property
    def filename(self) -> str:
        return f"{self.versioned_name}.tar"

    @property
    def is_path(self) -> bool:
        return False

    def get_bytes(self) -> bytes:
        return _export_module_files(self.name)


@dataclass
class FilesPackage:
    """Package to publish, from loose files"""

    path: Path
    name: str
    version: str

    def __init__(self, path: Path, name: str, version_arg: str | None = None) -> None:
        self.name = name
        self.path = path

        if version_arg is None:
            utc_now = datetime.datetime.utcnow()
            self.version = f"{utc_now.year}.{utc_now.month}.{utc_now.day}-{utc_now.second + 60 * utc_now.minute + 60 * 60 * utc_now.hour}"
        else:
            self.version = version_arg

    @property
    def versioned_name(self) -> str:
        return f"{self.name}-{self.version}"

    @property
    def filename(self) -> str:
        return f"{self.versioned_name}.tar"

    def get_bytes(self) -> bytes:
        # Archive the files
        buf = io.BytesIO()
        parent = Path(self.path)
        if parent.is_file():
            with tarfile.open(f"{self.name}.tar", "w", fileobj=buf) as tar:
                tar.add(parent, arcname=f"{self.name}/{parent.name}")
        else:
            with tarfile.open(f"{self.name}.tar", "w", fileobj=buf) as tar:
                for subpath in self._walk_path(self.path):
                    relative_path = subpath.relative_to(self.path.parent)
                    tar.add(subpath, arcname=relative_path)
        output_data = buf.getvalue()
        return output_data

    def _walk_path(self, path: Path) -> list[Path]:
        if path.is_file():
            return [self.path]
        output = []
        queue = [self.path]
        for item in queue:
            if item.is_file():
                if not item.name.startswith("."):
                    output.append(item)
            else:
                queue.extend(list(item.iterdir()))
        return output


def _run_create_meta(
    package_name: str, *, extra_python_path: Path | None = None
) -> str:
    """Run create-meta locally to produce recipe metadata JSON.

    If *extra_python_path* is given it is prepended to ``PYTHONPATH`` so
    that ``create-meta`` can import loose ``.py`` files that aren't
    installed as packages.
    """
    env = None
    if extra_python_path is not None:
        env = dict(os.environ)
        existing = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = (
            f"{extra_python_path}{os.pathsep}{existing}"
            if existing
            else str(extra_python_path)
        )
    result = subprocess.run(
        [
            "python",
            "-m",
            "prodigy_teams_recipes_sdk",
            "create-meta",
            package_name,
        ],
        capture_output=True,
        text=True,
        env=env,
    )
    if result.returncode != 0:
        raise CLIError(
            f"create-meta failed (exit {result.returncode}):\n{result.stderr}"
        )
    return result.stdout.strip()


@cli.subcommand(
    "publish",
    "code",
    name_or_path=Arg(help="Path or name of importable module with the recipes"),
    is_module=Arg(
        "--is-module", short="-m", help="Code to publish is an importable module"
    ),
    is_path=Arg("--is-path", short="-p", help="Code to publish is a path"),
    package_version=Arg("--package-version", help="Version identifier for the package"),
    # fmt: on
)
def publish_code(
    name_or_path: str,
    is_module: bool = False,
    is_path: bool = False,
    package_version: Optional[str] = None,
) -> UUID | None:
    """
    Upload and advertise a recipes package from your Python environment.
    """
    if not is_path and not is_module:
        is_path = _is_path(name_or_path)
    auth = get_auth_state()
    if auth.broker_host is None:
        raise CLIError(Messages.E035)
    if auth.broker_id is None:
        raise CLIError(Messages.E036)

    if is_path:
        path = Path(name_or_path)
        if path.is_dir() and (path / "setup.py").exists():
            msg.info(f"Building sdist from {name_or_path}")
            package = DistPackage(_make_sdist_if_dir(path))
        elif not path.is_file():
            msg.info(f"Interpreting {name_or_path} as directory of code")
            package = FilesPackage(name=path.name, path=path)
        elif path.name.endswith(".py"):
            msg.info(f"Interpreting {name_or_path} as Python source file")
            package = FilesPackage(name=path.name.replace(".py", ""), path=path)
        else:
            package = DistPackage(_make_sdist_if_dir(Path(name_or_path)))
        contents = package.get_bytes()
    else:
        msg.info(f"Interpreting {name_or_path} as module")
        package = ModulePackage(name_or_path, package_version)
        msg.info("Getting package contents")
        contents = package.get_bytes()
        if package_version is None:
            msg.info(f"No version specified. Defaulting to {package.version}")

    # Generate recipe metadata locally
    # For FilesPackage we need to tell the subprocess where to find the source
    # so that ``create-meta`` can actually import the module.
    extra_python_path: Path | None = None
    if isinstance(package, FilesPackage):
        src = package.path.resolve()
        extra_python_path = src.parent if src.is_file() else src.parent
    with loading("Generating recipe metadata"):
        meta_json = _run_create_meta(package.name, extra_python_path=extra_python_path)

    import json

    meta_list = json.loads(meta_json)
    if len(meta_list) == 0:
        msg.fail("No recipes found in package", exits=1)

    # Look up the base container from the latest prodigy-teams-recipes package
    with loading("Resolving base container image"):
        base_package = auth.client.package.latest(
            body=PackageReadingLatest(
                name="prodigy-teams-recipes",
                broker_id=auth.broker_id,
            )
        )
    base_image = base_package.container
    msg.info(f"Base image: {base_image}")

    # Build container image via broker
    tar_b64 = base64.b64encode(contents).decode("ascii")
    with loading("Building container image"):
        result: BuildImageResponse = auth.broker_client.envs.build_image(
            package_name=package.name,
            package_version=package.version,
            base_image=base_image,
            tar_data_b64=tar_b64,
            meta_json=meta_json,
        )

    msg.info(f"Built image: {result.image_uri}")

    # Register package+recipes in PAM with image URI as container
    msg.info("Advertising package to PAM server")

    class _MetaCompat:
        """Shim to match the old RecipesMeta interface for _publish_package_to_pam."""

        def __init__(self, recipes: list) -> None:
            self.recipes = recipes

    package_id = _publish_package_to_pam(
        auth.client,
        auth.broker_id,
        name=package.name,
        meta=_MetaCompat(result.meta),
        container=result.image_uri,
        filename=package.filename,
        version=package.version,
    )
    msg.good("Publication complete")
    return package_id


def _is_path(name_or_path: str) -> bool:
    try:
        path = Path(name_or_path)
    except ValueError:
        return False
    else:
        return path.exists()


def _make_sdist_if_dir(path: Path) -> Path:
    if not path.is_dir():
        return path
    env = dict(os.environ)
    env["SKIP_CYTHON"] = "1"
    r = subprocess.run(["python", "setup.py", "sdist"], cwd=path, text=True, env=env)
    r.check_returncode()
    dists = list((path / "dist").iterdir())
    if len(dists) == 0:
        raise ValueError("sdist not created after subprocess in {path / 'dist'}")
    # Return last modified file
    return max(dists, key=lambda p: p.stat().st_mtime)


def _publish_package_to_pam(
    client: Client,
    broker_id: UUID,
    *,
    name: str,
    filename: str,
    version: str,
    container: str,
    meta,
) -> UUID:
    msg.info(f"Creating package with {len(meta.recipes)} recipes")
    # Items are dicts at runtime but the generated model types them loosely
    recipes: list[dict] = [r for r in meta.recipes if isinstance(r, dict)]
    package = client.package.create(
        PackageCreating(
            name=name.replace("_", "-"),
            filename=filename,
            version=version,
            broker_id=broker_id,
            container=container,
            meta={r["name"]: r for r in recipes},
        )
    )

    msg.good(
        Messages.T002.format(noun="package", name=f"{package.name} ({package.id})")
    )
    for recipe_data in recipes:
        # The recipe create-meta command is supposed to output
        # entries that match the RecipeCreating body, except
        # for missing a package_id
        body = RecipeCreating(**recipe_data, package_id=package.id)
        r = client.recipe.create(body)
        recipe_type = "action" if body.job_type == "action" else "task"
        msg.good(Messages.T002.format(noun=f"{recipe_type} recipe", name=r.name))
    return package.id


def _export_module_files(module_name: str) -> bytes:
    module = importlib.import_module(module_name)
    assert not isinstance(module, str)
    assert module is not None
    if module.__file__ is not None:
        path = Path(module.__file__)
    else:
        path = Path(module.__path__[0])

    dist_path = _get_dist_info_path(module_name)
    if dist_path is None:
        # Package isn't a distribution -- it's just on the path
        record = _get_module_sources(path)
    elif dist_path.name.endswith("egg-info"):
        # Package is an 'egg'. This happens most often from pip install -e
        record = _read_sources(dist_path / "SOURCES.txt")
    elif not (dist_path / "RECORD").exists():
        # I think ths means it's an editable install?
        record = _get_module_sources(path)
    else:
        # Package is a distribution. This happens when you've installed from a
        # wheel or sdist
        record = _read_record(dist_path / "RECORD")
    # Archive the files
    buf = io.BytesIO()
    with tarfile.open(f"{module_name}.tar", "w", fileobj=buf) as tar:
        for subpath in record:
            tar.add(path.parent / Path(subpath), arcname=subpath)
    output_data = buf.getvalue()
    return output_data


def _get_module_sources(parent: Path) -> list[str]:
    if not parent.is_dir() and parent.name != "__init__.py":
        return [parent.name]
    elif not parent.is_dir():
        parent = parent.parent
    queue = [parent]
    output = []
    for path in queue:
        if path.name == "__pycache__":
            continue
        elif path.is_dir():
            queue.extend(path.iterdir())
        elif path.suffix == ".py":
            output.append(str(path.relative_to(parent.parent)))
    return output


def _get_dist_info_path(package_name: str) -> Path | None:
    try:
        dist = distribution(package_name)
    except PackageNotFoundError:
        return None
    dist_path = getattr(dist, "_path", None)
    if dist_path is None:
        return None
    return Path(dist_path)


def _read_sources(path: Path) -> list[str]:
    with path.open("r", encoding="utf8") as file_:
        output = file_.read().strip().split("\n")
    return output


def _read_record(path: Path) -> list[str]:
    output = []
    with path.open("r", encoding="utf8") as file_:
        for line in file_.read().strip().split("\n"):
            subpath, sha, size = line.split(",")
            output.append(subpath)
    return output
