#!/usr/bin/env python3
"""Clone Daylily analysis repositories into the shared FSx workspace."""

import argparse
from contextlib import contextmanager
import os
import re
import shlex
import subprocess
import sys
import tempfile
from typing import Any, Dict, Iterable, List, Optional, Tuple

import yaml


CONFIG_DIR = os.path.expanduser("~/.config/daylily")
GLOBAL_CONFIG_PATH = os.path.join(CONFIG_DIR, "daylily_cli_global.yaml")
AVAILABLE_REPOS_PATH = os.path.join(CONFIG_DIR, "daylily_pipeline_command_catalog.yaml")
DEPLOY_KEYS_PATH = os.path.join(CONFIG_DIR, "github_deploy_keys.yaml")
GITHUB_KNOWN_HOSTS_PATH = os.path.join(CONFIG_DIR, "github_known_hosts")
SAFE_SEGMENT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
FULL_COMMIT_SHA_RE = re.compile(r"^[0-9a-fA-F]{40}$")
CLUSTER_NAME_ENV_KEYS = (
    "DAYLILY_CLUSTER_NAME",
    "DAY_EC_CLUSTER_NAME",
    "PCLUSTER_CLUSTER_NAME",
    "PARALLELCLUSTER_CLUSTER_NAME",
    "stack_name",
    "STACK_NAME",
    "cfn_cluster_name",
    "CFN_CLUSTER_NAME",
)
CLUSTER_NAME_CONFIG_PATHS = ("/etc/parallelcluster/cfnconfig",)
CLUSTER_NAME_CONFIG_KEYS = {
    "stack_name",
    "cluster_name",
    "clusterName",
    "ClusterName",
    "cfn_cluster_name",
}


class ConfigError(RuntimeError):
    """Raised when configuration files are missing or malformed."""


def _load_yaml_mapping(path: str) -> Dict[str, Any]:
    """Load a YAML file that must contain a mapping at the top level."""
    if not os.path.exists(path):
        raise ConfigError(f"Configuration file not found: {path}")

    try:
        with open(path, "r", encoding="utf-8") as handle:
            payload = yaml.safe_load(handle)
    except yaml.YAMLError as exc:
        raise ConfigError(f"Invalid YAML in {path}: {exc}") from exc

    if not isinstance(payload, dict):
        raise ConfigError(f"Configuration file must contain a YAML mapping: {path}")
    return payload


def load_global_config() -> Dict[str, Any]:
    config = _load_yaml_mapping(GLOBAL_CONFIG_PATH)
    if "daylily" not in config:
        raise ConfigError(f"Missing 'daylily' section in {GLOBAL_CONFIG_PATH}.")
    return config["daylily"]


def load_available_repos() -> Tuple[str, Dict[str, Dict[str, Any]]]:
    config = _load_yaml_mapping(AVAILABLE_REPOS_PATH)
    repositories = config.get("repositories")
    if not isinstance(repositories, dict) or not repositories:
        raise ConfigError(f"No repositories defined in {AVAILABLE_REPOS_PATH}.")
    default_repo = config.get("default_repository")
    if not default_repo:
        raise ConfigError(f"Missing default_repository in {AVAILABLE_REPOS_PATH}.")
    return default_repo, repositories


def ensure_directory(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def is_lock_initialized_workspace(path: str) -> bool:
    """Return true only for a new workspace containing lock metadata and no analysis data."""
    try:
        entries = set(os.listdir(path))
    except OSError:
        return False
    metadata_dir = os.path.join(path, ".dayoa_agent")
    write_lock = os.path.join(metadata_dir, "write.lock")
    return (
        entries == {".dayoa_agent"}
        and os.path.isdir(metadata_dir)
        and not os.path.islink(metadata_dir)
        and os.path.isdir(write_lock)
        and not os.path.islink(write_lock)
    )


def safe_path_segment(value: str, *, field_name: str) -> str:
    text = str(value or "").strip()
    if not text:
        raise ConfigError(f"{field_name} is required.")
    if not SAFE_SEGMENT_RE.fullmatch(text):
        raise ConfigError(
            f"{field_name} must be a single path-safe segment matching "
            f"{SAFE_SEGMENT_RE.pattern!r}: {text!r}"
        )
    if text in {".", ".."} or ".." in text or "/" in text or "%" in text:
        raise ConfigError(f"{field_name} must not contain path traversal.")
    return text


def _strip_shell_value(value: str) -> str:
    text = value.strip()
    if len(text) >= 2 and text[0] == text[-1] and text[0] in {"'", '"'}:
        return text[1:-1]
    return text


def _cluster_name_from_env(keys: Iterable[str] = CLUSTER_NAME_ENV_KEYS) -> Optional[str]:
    for key in keys:
        value = os.environ.get(key)
        if value:
            return safe_path_segment(value, field_name=key)
    return None


def _cluster_name_from_cfnconfig(path: str) -> Optional[str]:
    if not os.path.exists(path):
        return None
    try:
        with open(path, "r", encoding="utf-8") as handle:
            for line in handle:
                text = line.strip()
                if not text or text.startswith("#") or "=" not in text:
                    continue
                key, value = text.split("=", 1)
                if key.strip() in CLUSTER_NAME_CONFIG_KEYS:
                    return safe_path_segment(
                        _strip_shell_value(value),
                        field_name=f"{path}:{key.strip()}",
                    )
    except OSError as exc:
        raise ConfigError(f"Could not read ParallelCluster config {path}: {exc}") from exc
    return None


def resolve_executing_entity(explicit_value: Optional[str]) -> str:
    if explicit_value:
        return safe_path_segment(explicit_value, field_name="executing_entity")

    env_cluster_name = _cluster_name_from_env()
    if env_cluster_name:
        return env_cluster_name

    for path in CLUSTER_NAME_CONFIG_PATHS:
        cluster_name = _cluster_name_from_cfnconfig(path)
        if cluster_name:
            return cluster_name

    raise ConfigError(
        "executing_entity is required when ParallelCluster cluster identity is unavailable. "
        "Pass --executing-entity, export DAYLILY_CLUSTER_NAME, or run on a headnode "
        "with /etc/parallelcluster/cfnconfig containing stack_name."
    )


def clone_repository(
    git_url: str,
    destination: str,
    ref: Optional[str],
    *,
    git_env: Optional[Dict[str, str]] = None,
) -> None:
    detached_commit = ref if ref and FULL_COMMIT_SHA_RE.fullmatch(ref) else None
    cmd = ["git", "clone"]
    if detached_commit:
        cmd.append("--no-checkout")
    elif ref:
        named_ref = ref
        for prefix in ("refs/heads/", "refs/tags/"):
            if named_ref.startswith(prefix):
                named_ref = named_ref[len(prefix) :]
                break
        cmd.extend(["--branch", named_ref])
    cmd.extend([git_url, destination])
    try:
        kwargs: Dict[str, Any] = {"check": True}
        if git_env is not None:
            kwargs["env"] = git_env
        subprocess.run(cmd, **kwargs)
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"Git clone failed with exit code {exc.returncode}.") from exc
    if detached_commit:
        try:
            checkout_kwargs: Dict[str, Any] = {"check": True}
            if git_env is not None:
                checkout_kwargs["env"] = git_env
            subprocess.run(
                [
                    "git",
                    "-C",
                    destination,
                    "fetch",
                    "--depth=1",
                    "origin",
                    detached_commit,
                ],
                **checkout_kwargs,
            )
            resolved = subprocess.run(
                ["git", "-C", destination, "rev-parse", "FETCH_HEAD"],
                check=True,
                capture_output=True,
                text=True,
                env=git_env,
            ).stdout.strip()
            if resolved.lower() != detached_commit.lower():
                raise RuntimeError(
                    f"Fetched commit does not match requested SHA: {resolved or 'missing'}"
                )
            subprocess.run(
                ["git", "-C", destination, "checkout", "--detach", "FETCH_HEAD"],
                **checkout_kwargs,
            )
        except subprocess.CalledProcessError as exc:
            raise RuntimeError(f"Git checkout failed with exit code {exc.returncode}.") from exc


def load_deploy_key_reference(repo_key: str) -> Tuple[str, str]:
    config = _load_yaml_mapping(DEPLOY_KEYS_PATH)
    if config.get("config_version") != 1:
        raise ConfigError(f"Unsupported deploy-key config_version in {DEPLOY_KEYS_PATH}.")
    deploy_keys = config.get("deploy_keys")
    if not isinstance(deploy_keys, dict):
        raise ConfigError(f"Missing deploy_keys mapping in {DEPLOY_KEYS_PATH}.")
    reference = deploy_keys.get(repo_key)
    if not isinstance(reference, dict):
        raise ConfigError(
            f"No deploy-key reference is configured for repository '{repo_key}' in "
            f"{DEPLOY_KEYS_PATH}."
        )
    secret_arn = str(reference.get("secret_arn") or "").strip()
    region = str(reference.get("region") or "").strip()
    if not secret_arn or not region:
        raise ConfigError(
            f"Repository '{repo_key}' deploy-key reference requires secret_arn and region."
        )
    return secret_arn, region


def fetch_deploy_key(secret_arn: str, region: str) -> str:
    try:
        import boto3
    except ImportError as exc:
        raise RuntimeError("boto3 is required for AWS deploy-key authentication.") from exc

    try:
        response = boto3.client("secretsmanager", region_name=region).get_secret_value(
            SecretId=secret_arn
        )
    except Exception as exc:
        raise RuntimeError(
            "Unable to read the configured DayOA deploy-key secret from Secrets Manager."
        ) from exc
    key = response.get("SecretString")
    if not isinstance(key, str) or not key.strip():
        raise RuntimeError(
            "The configured deploy-key secret must contain a non-empty SecretString."
        )
    stripped = key.strip()
    if not (
        stripped.startswith("-----BEGIN OPENSSH PRIVATE KEY-----")
        and stripped.endswith("-----END OPENSSH PRIVATE KEY-----")
    ):
        raise RuntimeError("The configured deploy-key secret is not an OpenSSH private key.")
    return stripped + "\n"


@contextmanager
def deploy_key_git_environment(repo_key: str):
    secret_arn, region = load_deploy_key_reference(repo_key)
    private_key = fetch_deploy_key(secret_arn, region)
    if not os.path.isfile(GITHUB_KNOWN_HOSTS_PATH) or os.path.getsize(GITHUB_KNOWN_HOSTS_PATH) == 0:
        raise ConfigError(f"GitHub known-hosts file is missing or empty: {GITHUB_KNOWN_HOSTS_PATH}")

    with tempfile.TemporaryDirectory(prefix="day-clone-key-") as temp_dir:
        os.chmod(temp_dir, 0o700)
        key_path = os.path.join(temp_dir, "deploy_key")
        with open(key_path, "w", encoding="utf-8") as handle:
            handle.write(private_key)
        os.chmod(key_path, 0o600)
        env = os.environ.copy()
        env["GIT_TERMINAL_PROMPT"] = "0"
        env["GIT_SSH_COMMAND"] = " ".join(
            [
                "ssh",
                "-i",
                shlex.quote(key_path),
                "-o",
                "IdentitiesOnly=yes",
                "-o",
                "StrictHostKeyChecking=yes",
                "-o",
                f"UserKnownHostsFile={shlex.quote(GITHUB_KNOWN_HOSTS_PATH)}",
            ]
        )
        try:
            yield env
        finally:
            private_key = ""


def check_repository_access(
    git_url: str,
    ref: Optional[str],
    *,
    git_env: Optional[Dict[str, str]] = None,
) -> None:
    kwargs: Dict[str, Any] = {"check": True, "stdout": subprocess.DEVNULL}
    if git_env is not None:
        kwargs["env"] = git_env
    try:
        if ref and FULL_COMMIT_SHA_RE.fullmatch(ref):
            subprocess.run(["git", "ls-remote", git_url], **kwargs)
            with tempfile.TemporaryDirectory(prefix="day-clone-ref-") as temp_dir:
                subprocess.run(["git", "init", "--bare", temp_dir], **kwargs)
                subprocess.run(
                    ["git", "-C", temp_dir, "fetch", "--depth=1", git_url, ref],
                    **kwargs,
                )
                resolved = subprocess.run(
                    ["git", "-C", temp_dir, "rev-parse", "FETCH_HEAD"],
                    check=True,
                    capture_output=True,
                    text=True,
                    env=git_env,
                ).stdout.strip()
                if resolved.lower() != ref.lower():
                    raise RuntimeError(
                        f"Fetched commit does not match requested SHA: {resolved or 'missing'}"
                    )
            return

        cmd = ["git", "ls-remote", "--exit-code", git_url]
        if ref:
            if ref.startswith(("refs/heads/", "refs/tags/")):
                cmd.append(ref)
            else:
                cmd.extend([f"refs/heads/{ref}", f"refs/tags/{ref}"])
        subprocess.run(cmd, **kwargs)
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(
            f"Git repository authentication or ref check failed with exit code {exc.returncode}."
        ) from exc


@contextmanager
def repository_git_environment(repo_key: str, auth_mode: str):
    if auth_mode == "none":
        yield None
        return
    if auth_mode != "aws_deploy_key":
        raise ConfigError(f"Unsupported auth_mode for repository '{repo_key}': {auth_mode!r}")
    with deploy_key_git_environment(repo_key) as env:
        yield env


def derive_default_path(repo_url: str) -> str:
    basename = os.path.basename(repo_url)
    if basename.endswith(".git"):
        basename = basename[:-4]
    return basename or "repository"


def print_available_repositories(repositories: Dict[str, Dict[str, Any]]) -> None:
    print("Clone syntax:")
    print("  day-clone --repository <repo-key> --destination <analysis-id> --git-tag <ref>")
    print("  day-clone -d <analysis-id> -t <ref>  # uses default_repository")
    print()
    print("Available repositories:\n")
    for key, repo in repositories.items():
        name = repo.get("display_name", key)
        print(f"- {key}: {name}")
        if description := repo.get("description"):
            print(f"    {description}")
        default_ref = repo.get("default_ref")
        if default_ref:
            print(f"    Default ref: {default_ref}")
        print()


def main(argv: List[str]) -> int:
    try:
        global_cfg = load_global_config()
        default_repo_key, available_repos = load_available_repos()
    except ConfigError as err:
        print(f"Error: {err}", file=sys.stderr)
        return 2

    parser = argparse.ArgumentParser(
        description="Clone Daylily analysis repositories into the FSx analysis workspace.",
    )
    parser.add_argument(
        "-d",
        "--destination",
        help="Name of the analysis workspace directory to create under the user-specific root.",
    )
    parser.add_argument(
        "-t",
        "--git-tag",
        help=(
            "Git branch, tag, or full commit SHA to clone. "
            "Defaults to the repository's configured default."
        ),
    )
    parser.add_argument(
        "-r",
        "--git-repo",
        help="Override the git repository URL to clone. Overrides --repository.",
    )
    parser.add_argument(
        "-w",
        "--which-one",
        choices=("https", "ssh"),
        default=None,
        help="Require the selected repository's configured transport to match this value.",
    )
    parser.add_argument(
        "-c",
        "--clone-root",
        help="Root directory where analysis workspaces are created. Defaults to the configured analysis_root.",
    )
    parser.add_argument(
        "--executing-entity",
        help=(
            "Entity directory to create within the clone root. Defaults to the cluster "
            "name from exported cluster env or /etc/parallelcluster/cfnconfig."
        ),
    )
    parser.add_argument(
        "--repository",
        help="Key of the repository defined in daylily_pipeline_command_catalog.yaml to clone.",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available repositories and exit.",
    )
    parser.add_argument(
        "--check-auth",
        action="store_true",
        help="Verify repository authentication and the selected ref without cloning.",
    )

    args = parser.parse_args(argv)

    if args.list:
        print_available_repositories(available_repos)
        return 0

    repo_config: Optional[Dict[str, Any]] = None
    if not args.git_repo:
        repo_key = args.repository or default_repo_key
        repo_config = available_repos.get(repo_key)
        if repo_config is None:
            print(
                f"Error: repository '{repo_key}' is not defined in daylily_pipeline_command_catalog.yaml.",
                file=sys.stderr,
            )
            return 1
    else:
        repo_key = "custom"

    if args.git_repo:
        git_url = args.git_repo
        relative_path = derive_default_path(git_url)
        default_ref = None
        auth_mode = "none"
    else:
        configured_transport = str(repo_config.get("clone_transport") or "").strip()
        auth_mode = str(repo_config.get("auth_mode") or "").strip()
        if configured_transport not in {"https", "ssh"}:
            print(
                f"Error: repository '{repo_key}' has invalid or missing clone_transport.",
                file=sys.stderr,
            )
            return 1
        if not auth_mode:
            print(
                f"Error: repository '{repo_key}' is missing auth_mode.",
                file=sys.stderr,
            )
            return 1
        if args.which_one and args.which_one != configured_transport:
            print(
                f"Error: repository '{repo_key}' requires {configured_transport} transport; "
                f"{args.which_one} override is not permitted.",
                file=sys.stderr,
            )
            return 1
        if auth_mode == "aws_deploy_key" and configured_transport != "ssh":
            print(
                f"Error: repository '{repo_key}' aws_deploy_key authentication requires ssh.",
                file=sys.stderr,
            )
            return 1
        url_key = f"{configured_transport}_url"
        git_url = repo_config.get(url_key, "")
        if not git_url:
            print(
                f"Error: repository '{repo_key}' does not define a {url_key}.",
                file=sys.stderr,
            )
            return 1
        relative_path = repo_config.get("relative_path") or repo_key
        default_ref = repo_config.get("default_ref")

    git_ref = args.git_tag or default_ref

    if args.check_auth:
        try:
            with repository_git_environment(repo_key, auth_mode) as git_env:
                check_repository_access(git_url, git_ref, git_env=git_env)
        except (ConfigError, RuntimeError) as err:
            print(f"Error: {err}", file=sys.stderr)
            return 1
        print(f"Repository authentication validated: {repo_key} @ {git_ref or 'HEAD'}")
        return 0

    if not args.destination:
        print(
            "Error: --destination (-d) is required when cloning a repository.",
            file=sys.stderr,
        )
        return 1

    clone_root = args.clone_root or global_cfg.get("analysis_root")
    if not clone_root:
        print(
            "Error: analysis_root is not configured in daylily_cli_global.yaml.",
            file=sys.stderr,
        )
        return 2

    clone_root = os.path.abspath(os.path.expanduser(clone_root))
    if not os.path.isdir(clone_root):
        print(
            f"Error: clone_root directory '{clone_root}' does not exist.",
            file=sys.stderr,
        )
        return 1

    try:
        executing_entity = resolve_executing_entity(args.executing_entity)
        destination = safe_path_segment(args.destination, field_name="destination")
    except ConfigError as err:
        print(f"Error: {err}", file=sys.stderr)
        return 1

    user_root = os.path.join(clone_root, executing_entity)
    ensure_directory(user_root)

    destination_root = os.path.join(user_root, destination)
    if os.path.exists(destination_root) and not is_lock_initialized_workspace(destination_root):
        print(
            f"Error: destination '{destination_root}' already exists.",
            file=sys.stderr,
        )
        return 1

    relative_path = os.path.basename(relative_path.rstrip("/"))
    target_repo_dir = os.path.join(destination_root, relative_path)
    ensure_directory(destination_root)

    print("Cloning repository...")
    try:
        with repository_git_environment(repo_key, auth_mode) as git_env:
            clone_repository(git_url, target_repo_dir, git_ref, git_env=git_env)
    except (ConfigError, RuntimeError) as err:
        print(f"Error: {err}", file=sys.stderr)
        return 1

    print()
    print("Great success! Daylily repository cloned.")
    print(f"Repository: {git_url}")
    if git_ref:
        print(f"Reference : {git_ref}")
    print(f"Location  : {target_repo_dir}")
    print()
    print("To get started:")
    print(f"  cd {target_repo_dir}")
    print("  # initialize and run the analysis repository per its documentation")

    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
