#!/usr/bin/env python3

# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

from __future__ import annotations

import sys
import shutil
import os
import getpass
import sysconfig
import site
import tempfile
import zipfile
from pathlib import Path
from urllib.request import urlopen, Request

import subprocess
import argparse
from typing import Optional
import importlib.resources as ir
import datatailr
from datatailr.logging import CYAN, RED, YELLOW, GREEN
from datatailr import User, is_dt_installed

DEFAULT_TARGET = "datatailr_demo_project"
TEMPLATES_REPO_OWNER = "Datatailr"
TEMPLATES_REPO_NAME = "datatailr-examples"
TEMPLATES_BRANCH = "main"
VENV_NAME = ".venv"


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="datatailr",
        description="Datatailr command line utilities. Use subcommands to access features.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command", metavar="<command>")

    # init-demo subcommand
    p_init = sub.add_parser(
        "init-demo",
        help="Copy the bundled demo project into a target directory.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_init.add_argument(
        "-d",
        "--destination",
        metavar="PATH",
        default=str(Path.cwd() / DEFAULT_TARGET),
        help="Destination directory to create (will be created if missing).",
    )
    p_init.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite existing files (by default existing files are skipped).",
    )
    p_init.add_argument(
        "-n",
        "--no-open",
        action="store_true",
        help="Do not open code-server in the demo project folder.",
    )

    # version subcommand
    sub.add_parser("version", help="Show python and datatailr versions.")

    # setup-cli subcommand to configure remote cli
    p_setup_cli = sub.add_parser(
        "setup-cli",
        help="Configure remote dt CLI and add 'dt' into user or venv PATH.\
            Remote CLI can be configured only outside of Datatailr.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_setup_cli.add_argument(
        "--url", help="Datatailr base URL (skips interactive prompt)."
    )
    p_setup_ssh = sub.add_parser(
        "setup-ssh",
        help="Configure SSH access to a workstation running in Datatailr.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_setup_ssh.add_argument(
        "workstation",
        nargs="?",
        help="Workstation name. If omitted, you'll be prompted to choose from available workstations.",
    )
    p_setup_ssh.add_argument(
        "--url", help="Datatailr base URL (skips interactive prompt)."
    )
    p_setup_skills = sub.add_parser(
        "setup-skills",
        help="Copy Datatailr skills to an agent skills directory.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_setup_skills.add_argument(
        "-g",
        "--global",
        dest="global_dir",
        action="store_true",
        help="Install to ~/.agents/skills instead of ./.agents/skills.",
    )
    p_setup_skills.add_argument(
        "-f",
        "--force",
        dest="force",
        action="store_true",
        help="Overwrite existing datatailr skills without prompting.",
    )
    p_login = sub.add_parser(
        "login",
        help="Login to your Datatailr installation.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_login.add_argument("--url", help="Datatailr base URL (skips interactive prompt).")

    # init-template subcommand
    p_init_template = sub.add_parser(
        "init-template",
        help="Clone a template folder from the Datatailr examples repository.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_init_template.add_argument(
        "template",
        metavar="TEMPLATE",
        help="Template folder name (e.g. getting-started, excel-addin-demo).",
    )
    p_init_template.add_argument(
        "-d",
        "--destination",
        metavar="PATH",
        default=str(Path.cwd()),
        help="Destination directory. Template folder will be created inside it.",
    )
    p_init_template.add_argument(
        "-v",
        "--venv",
        action="store_true",
        help="Create a virtual environment inside the template folder and print activation instructions.",
    )
    p_init_template.add_argument(
        "-r",
        "--install-deps",
        action="store_true",
        help="Install dependencies from requirements.txt into the virtual environment. Requires --venv.",
    )

    return parser


def copy_tree(src: Path, dst: Path, username: str, overwrite: bool = False) -> None:
    file_types = (
        ".py",
        ".yaml",
        ".yml",
        ".md",
        ".txt",
        ".csv",
        ".json",
        ".png",
        ".jpg",
        ".jpeg",
        ".ico",
        ".html",
        ".css",
        ".js",
    )
    for root, dirs, files in os.walk(src):
        rel = Path(root).relative_to(src)
        target_root = dst / rel
        target_root.mkdir(parents=True, exist_ok=True)
        for d in dirs:
            (target_root / d).mkdir(exist_ok=True)
        for f in files:
            if not f.lower().endswith(file_types):
                continue  # skip non-relevant files
            src_file = Path(root) / f
            dst_file = target_root / f
            if dst_file.exists() and not overwrite:
                continue  # skip existing unless force enabled
            # replace <>USERNAME<> placeholders in the files:

            with open(src_file, "r") as sf, open(dst_file, "w") as df:
                content = sf.read()
                content = content.replace("<>USERNAME<>", username)
                df.write(content)


def git_init_repo(path: Path) -> None:
    """Initialize a git repository at the given path if not already a git repo."""
    if (path / ".git").exists():
        return  # already a git repo

    if not shutil.which("git"):
        print(
            YELLOW(
                "Git is not installed or not found in PATH; skipping git initialization."
            )
        )
        return

    for command in [
        ["git", "init"],
        ["git", "add", "."],
        ["git", "commit", "-m", "Initial commit"],
    ]:
        subprocess.run(
            command,
            cwd=str(path),
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )


def install_code_extensions() -> None:
    """Install recommended code-server extensions for Datatailr."""
    if not shutil.which("code-server"):
        print(
            YELLOW(
                "Code server is not installed or not found in PATH; skipping extension installation."
            )
        )
        return
    extensions = ["ms-python.python", "ms-python.debugpy", "astral-sh.ty"]
    print(CYAN("Installing recommended code-server extensions..."))
    for ext in extensions:
        subprocess.run(["code-server", "--install-extension", ext], check=True)
    print(GREEN("Extensions installed."))


def open_folder(path: Path) -> None:
    """Launch code server in the target path."""
    if not shutil.which("code-server"):
        print(
            YELLOW(
                "Code server is not installed or not found in PATH; skipping code server launch."
            )
        )
        return
    subprocess.run(["code-server", "--reuse-window", str(path)], check=True)


def default_user_bin() -> Path:
    """Get the default user (or venv) scripts/bin directory."""
    scripts_dir = sysconfig.get_path("scripts")
    if scripts_dir:
        return Path(scripts_dir)
    base = Path(site.getuserbase())
    return base / ("Scripts" if os.name == "nt" else "bin")


def path_contains(path: Path) -> bool:
    entries = os.environ.get("PATH", "").split(os.pathsep)
    return any(Path(p).resolve() == path.resolve() for p in entries if p)


def write_shim(script_path: Path) -> None:
    python_exe = sys.executable
    if os.name == "nt":
        content = (
            f'@echo off\r\n"{python_exe}" -m datatailr._remote_client.remote_cli %*\r\n'
        )
        script_path.write_text(content, encoding="utf-8")
    else:
        content = f'#!/usr/bin/env sh\nexec "{python_exe}" -m datatailr._remote_client.remote_cli "$@"\n'
        script_path.write_text(content, encoding="utf-8")
        os.chmod(script_path, 0o755)


def setup_cli(target_dir: Path) -> int:
    if is_dt_installed():
        print(
            YELLOW(
                "Native dt CLI detected on this system. Skipping install - remote CLI can be installed only outside of Datatailr."
            )
        )
        return 0
    target_dir.mkdir(parents=True, exist_ok=True)
    script_name = "dt.cmd" if os.name == "nt" else "dt"
    script_path = target_dir / script_name
    write_shim(script_path)
    print(GREEN(f"Installed remote dt CLI to: {script_path}"))
    if not path_contains(target_dir):
        print(
            YELLOW(
                f'{target_dir} is not on PATH. Add it to your shell profile to enable "dt":\n'
                f'  - Linux/macOS (bash): export PATH="{target_dir}:$PATH"\n'
                f'  - Windows (PowerShell): $env:Path += ";{target_dir}"\n'
                f"Otherwise, use the full path: {script_path} to run the 'dt' command."
            )
        )
    return 0


def setup_ssh(workstation: Optional[str] = None, url: Optional[str] = None) -> int:
    try:
        from datatailr._remote_client.remote_cli import run_setup_ssh
    except Exception as exc:
        print(RED(f"Could not load remote SSH setup client: {exc}"))
        return 2

    return int(run_setup_ssh(workstation=workstation, base_url=url))


def login(url: Optional[str] = None) -> int:
    try:
        from datatailr._remote_client.remote_cli import run_login
    except Exception as exc:
        print(RED(f"Could not load remote login client: {exc}"))
        return 2
    return int(run_login(base_url=url))


def _get_skills_target_dir(global_dir: bool) -> Path:
    """Return target directory for skills: ./.agents/skills or ~/.agents/skills."""
    if global_dir:
        return (Path.home() / ".agents" / "skills").resolve()
    return (Path.cwd() / ".agents" / "skills").resolve()


def _collect_existing_skill_paths(source_dir: Path, target_dir: Path) -> list[Path]:
    """Return list of target paths that already exist and would be overwritten."""
    existing = []
    for root, dirs, files in os.walk(source_dir):
        rel = Path(root).relative_to(source_dir)
        target_root = target_dir / rel
        for file_name in files:
            dst_file = target_root / file_name
            if dst_file.exists():
                existing.append(dst_file)
    return existing


def _copy_agent_skills(
    source_dir: Path, target_dir: Path, overwrite: bool = False
) -> None:
    for root, dirs, files in os.walk(source_dir):
        rel = Path(root).relative_to(source_dir)
        target_root = target_dir / rel
        target_root.mkdir(parents=True, exist_ok=True)
        for directory in dirs:
            (target_root / directory).mkdir(parents=True, exist_ok=True)
        for file_name in files:
            src_file = Path(root) / file_name
            dst_file = target_root / file_name
            if dst_file.exists():
                if overwrite:
                    dst_file.unlink()
                else:
                    continue
            shutil.copy2(src_file, dst_file)


def setup_skills(global_dir: bool = False, force: bool = False) -> int:
    source_dir = Path(datatailr.__file__).resolve().parent / "agent_skills"
    target_dir = _get_skills_target_dir(global_dir)

    if target_dir.exists() and not target_dir.is_dir():
        print(RED(f"Target path exists and is not a directory: {target_dir}"))
        return 1

    target_dir.mkdir(parents=True, exist_ok=True)
    existing = _collect_existing_skill_paths(source_dir, target_dir)
    if existing and not force:
        print(YELLOW("The following skills would be overwritten:"))
        for p in existing:
            print(f"  {p.relative_to(target_dir)}")
        try:
            reply = input("Overwrite? (y/n): ").strip().lower()
        except EOFError:
            return 1
        if reply not in ("y", "yes"):
            print("Aborted.")
            return 1

    overwrite = bool(existing)
    _copy_agent_skills(source_dir, target_dir, overwrite=overwrite)

    print(GREEN(f"Installed Datatailr skills to: {target_dir}"))
    return 0


def _download_template_from_github(template_path: str) -> Path:
    """Download the template folder from GitHub. Returns path to extracted template dir."""
    archive_url = (
        f"https://github.com/{TEMPLATES_REPO_OWNER}/{TEMPLATES_REPO_NAME}"
        f"/archive/refs/heads/{TEMPLATES_BRANCH}.zip"
    )
    archive_root = f"{TEMPLATES_REPO_NAME}-{TEMPLATES_BRANCH}"
    inner_path = f"{archive_root}/{template_path}"

    print(CYAN(f"Downloading template '{template_path}' from GitHub..."))
    req = Request(archive_url, headers={"User-Agent": "Datatailr-CLI"})
    with urlopen(req) as resp:
        zip_data = resp.read()

    tmpdir = Path(tempfile.mkdtemp(prefix="datatailr_template_"))
    try:
        zip_path = tmpdir / "repo.zip"
        zip_path.write_bytes(zip_data)
        with zipfile.ZipFile(zip_path, "r") as zf:
            members = [
                m
                for m in zf.namelist()
                if m.startswith(inner_path + "/") or m == inner_path
            ]
            if not members:
                raise ValueError(
                    f"Template '{template_path}' not found in {TEMPLATES_REPO_OWNER}/{TEMPLATES_REPO_NAME}. "
                    f"Check the template name."
                )
            for m in members:
                zf.extract(m, tmpdir)
        return tmpdir / inner_path
    except Exception:
        shutil.rmtree(tmpdir, ignore_errors=True)
        raise


def _create_venv(project_dir: Path) -> Path:
    """Create a virtual environment in project_dir. Returns path to venv."""
    venv_path = project_dir / VENV_NAME
    if venv_path.exists():
        print(YELLOW(f"Virtual environment already exists at {venv_path}"))
        return venv_path
    print(CYAN(f"Creating virtual environment in {venv_path}..."))
    subprocess.run(
        [sys.executable, "-m", "venv", "--system-site-packages", str(venv_path)],
        check=True,
        capture_output=True,
    )
    print(GREEN("Virtual environment created."))
    return venv_path


def _get_pip_exe(venv_path: Path) -> Path:
    """Return path to pip inside the venv."""
    if os.name == "nt":
        return venv_path / "Scripts" / "pip.exe"
    return venv_path / "bin" / "pip"


def _install_requirements(project_dir: Path, venv_path: Path) -> None:
    """Install requirements.txt using the venv's pip."""
    requirements = project_dir / "requirements.txt"
    if not requirements.exists():
        print(YELLOW(f"No requirements.txt found in {project_dir}; skipping install."))
        return
    pip_exe = _get_pip_exe(venv_path)
    if not pip_exe.exists():
        print(RED(f"pip not found at {pip_exe}"))
        return
    print(CYAN("Installing dependencies from requirements.txt..."))
    subprocess.run(
        [str(pip_exe), "install", "-r", str(requirements)],
        cwd=str(project_dir),
        check=True,
    )
    print(GREEN("Dependencies installed."))


def _activation_instructions(venv_path: Path) -> str:
    if os.name == "nt":
        return f"{venv_path}\\Scripts\\activate"
    return f"source {venv_path}/bin/activate"


def _sanitize_template_path(template: str) -> str:
    """Reject path traversal and ensure a valid template name."""
    normalized = Path(template).as_posix().strip("/")
    if not normalized:
        raise ValueError("Template name cannot be empty.")
    if ".." in normalized:
        raise ValueError("Template path cannot contain '..'.")
    return normalized


def init_template(
    template: str,
    destination: str,
    venv: bool,
    install_deps: bool,
) -> int:
    """Clone a template from GitHub and optionally set up venv/deps."""
    if install_deps and not venv:
        print(
            RED(
                "Error: --install-deps requires --venv. "
                "Dependencies are installed into the virtual environment."
            )
        )
        print(YELLOW("Usage: datatailr init-template TEMPLATE [options]"))
        print(
            "  -d, --destination PATH  Destination directory (default: current directory)"
        )
        print(
            "  -v, --venv              Create a virtual environment in the template folder"
        )
        print("  -r, --install-deps      Install requirements.txt (requires --venv)")
        print("\nRun 'datatailr init-template --help' for full usage.")
        return 1

    try:
        template = _sanitize_template_path(template)
    except ValueError as e:
        print(RED(str(e)))
        return 1

    dest_dir = Path(destination).resolve()
    dest_dir.mkdir(parents=True, exist_ok=True)
    target_dir = dest_dir / template

    if target_dir.exists() and any(target_dir.iterdir()):
        print(
            RED(
                f"Target '{target_dir}' already exists and is not empty. "
                "Choose a different destination or remove the existing folder."
            )
        )
        return 1

    try:
        src_path = _download_template_from_github(template)
    except ValueError as e:
        print(RED(str(e)))
        return 1
    except Exception as e:
        print(RED(f"Failed to download template: {e}"))
        return 1

    try:
        shutil.copytree(src_path, target_dir)
    finally:
        shutil.rmtree(src_path.parent.parent, ignore_errors=True)

    print(GREEN(f"Template cloned to {target_dir}"))

    venv_path: Path | None = None
    if venv:
        venv_path = _create_venv(target_dir)
        if install_deps:
            _install_requirements(target_dir, venv_path)

    if venv_path is not None:
        activ = _activation_instructions(venv_path)
        print(CYAN(f"To activate the virtual environment, run:\n  {activ}"))

    return 0


def main(argv: list[str] | None = None) -> int:
    argv = argv or sys.argv[1:]
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "init-demo":
        try:
            template_root = ir.files("datatailr_demo")
        except Exception as e:  # pragma: no cover
            print(RED(f"Could not access demo template package: {e}"))
            return 1
        src_path = Path(str(template_root))

        target = Path(args.destination)
        if target.exists() and any(target.iterdir()) and not args.force:
            print(
                YELLOW(
                    f"Target '{target}' already exists and is not empty\nUse --force to overwrite or choose a different path."
                )
            )
            return 0
        print(CYAN(f"Copying Datatailr demo template to: {target}"))
        if is_dt_installed():
            username = User.signed_user().name
        else:
            username = getpass.getuser()
            print(
                YELLOW(f"The dt CLI is not installed. Using OS login name: {username}.")
            )
        if username is None:
            username = "dtuser"
        copy_tree(src_path, target, username, overwrite=args.force)
        git_init_repo(target)
        print(GREEN("Done. Explore the README.md inside the demo project."))
        install_code_extensions()
        if not args.no_open:
            open_folder(target)
        return 0

    elif args.command == "version":
        # Lazy import to avoid unnecessary startup cost
        import importlib.metadata as im

        pkg_version = (
            im.version("datatailr")
            if "datatailr" in {d.metadata["Name"] for d in im.distributions()}
            else "unknown"
        )
        # print nicely colored version info
        print(GREEN(f"Datatailr {pkg_version}"))
        print(CYAN(f"Python {sys.version.split()[0]}"))
        return 0
    elif args.command == "setup-cli":
        target_dir = default_user_bin()
        setup_rc = setup_cli(target_dir)
        if setup_rc != 0:
            return setup_rc
        return login(url=args.url)
    elif args.command == "login":
        return login(url=args.url)
    elif args.command == "setup-skills":
        return setup_skills(global_dir=args.global_dir, force=args.force)
    elif args.command == "setup-ssh":
        return setup_ssh(workstation=args.workstation, url=args.url)
    elif args.command == "init-template":
        return init_template(
            template=args.template,
            destination=args.destination,
            venv=args.venv,
            install_deps=args.install_deps,
        )

    # No subcommand: show help and exit
    parser.print_help()
    return 0


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())
