from __future__ import annotations

from k4s.core.executor import ExecutorError
from k4s.core.downloader import AssetDownloader
from k4s.core.executor import Executor
from k4s.core.products import Step
from k4s.recipes.common.docker import build_docker_steps
from k4s.recipes.common.linux import (
    apt_install,
    detect_ubuntu_like,
    ensure_apt_updated,
    ensure_dir_owned,
    ensure_group,
    ensure_user,
    ensure_user_in_group,
    install_kubectl,
)
from k4s.recipes.common.run import check, q, run
from k4s.recipes.dataiku.model import DataikuInstallPlan
from k4s.recipes.dataiku.r_integration import build_r_steps
from k4s.ui.ui import Ui, Verbosity


ALLOWED_NODE_TYPES = {"design", "automation", "govern"}

GOVERN_DIP_PROPERTIES = "config/dip.properties"

def _extract_encrypted_password(output: str) -> str | None:
    """Extract an e:AES:... token from govern-admin output."""
    import re

    # The encrypted token can include additional ':'-separated segments.
    # Match "e:AES:" followed by any non-whitespace characters.
    m = re.search(r"(e:AES:\S+)", output or "")
    return m.group(1) if m else None


def _encrypt_govern_password(ex: Executor, plan: DataikuInstallPlan, *, ui: Ui | None = None) -> str | None:
    """Return encrypted govern password (e:AES:...) or None.

    Passes the password as a direct command-line argument to govern-admin.
    The command runs with silent=True so the password never appears in
    debug output.
    """
    if not plan.govern_jdbc_password:
        return None
    if not plan.encrypt_govern_password:
        return None
    try:
        cmd = (
            f"sudo -n -u {q(plan.os_user)} -i -- "
            f"{q(plan.data_dir)}/bin/govern-admin encrypt-password {q(plan.govern_jdbc_password)}"
        )
        rc, out, err = run(ex, cmd, silent=True)
        if rc != 0:
            if ui and ui.options.verbosity >= Verbosity.debug:
                ui.debug(f"govern-admin encrypt-password failed (rc={rc})")
                for line in (out or "").strip().splitlines()[:10]:
                    ui.debug(f"  stdout: {line}")
                for line in (err or "").strip().splitlines()[:10]:
                    ui.debug(f"  stderr: {line}")
            return None
        combined = (out or "") + "\n" + (err or "")
        token = _extract_encrypted_password(combined)
        if not token and ui and ui.options.verbosity >= Verbosity.debug:
            ui.debug("govern-admin encrypt-password succeeded but no e:AES token found in output")
            for line in combined.strip().splitlines()[:10]:
                ui.debug(f"  output: {line}")
        return token
    except Exception:
        return None


def _set_properties(content: str, values: dict[str, str]) -> str:
    """Set (or add) key=value pairs in a .properties file content."""
    lines = content.splitlines()
    existing_keys: set[str] = set()
    out: list[str] = []
    for line in lines:
        raw = line
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or stripped.startswith("!"):
            out.append(raw)
            continue
        if "=" not in line:
            out.append(raw)
            continue
        k, _ = line.split("=", 1)
        key = k.strip()
        if key in values:
            out.append(f"{key}={values[key]}")
            existing_keys.add(key)
        else:
            out.append(raw)
    for k, v in values.items():
        if k not in existing_keys:
            out.append(f"{k}={v}")
    return "\n".join(out) + ("\n" if out else "")


def _validate_govern_db_config(plan: DataikuInstallPlan) -> None:
    """Validate Govern DB config inputs."""
    if plan.node_type != "govern":
        return
    missing: list[str] = []
    if not plan.govern_jdbc_url:
        missing.append("--govern-jdbc-url")
    if not plan.govern_jdbc_user:
        missing.append("--govern-jdbc-user")
    if not plan.govern_jdbc_password:
        missing.append("--govern-jdbc-password (or use --ask-govern-jdbc-password)")
    if missing:
        raise ValueError(
            "Govern requires a PostgreSQL connection configuration.\n"
            f"Missing: {', '.join(missing)}"
        )


def build_installer_args(plan: DataikuInstallPlan) -> str:
    """Build installer.sh arguments for the selected DSS node type."""
    installer_args = f"-d {plan.data_dir} -p {plan.dss_port} -y"
    if plan.node_type != "design":
        installer_args = f"-t {plan.node_type} {installer_args}"
    return installer_args


def run_preflight(ui: Ui, ex: Executor, plan: DataikuInstallPlan) -> None:
    """Run read-only preflight checks for Dataiku install."""
    ui.log("Checking OS basics and non-interactive sudo.")

    if plan.node_type not in ALLOWED_NODE_TYPES:
        raise ValueError(f"Unsupported Dataiku node type: {plan.node_type}")

    _validate_govern_db_config(plan)

    if not detect_ubuntu_like(ex):
        ui.warning("OS detection: this recipe currently targets Ubuntu/Debian (apt).")

    check(ex, "sudo -n true", error_hint="Configure passwordless sudo (NOPASSWD) for the SSH user.")

    _, os_release, _ = run(ex, "cat /etc/os-release 2>/dev/null || true")
    os_info: dict[str, str] = {}
    for line in (os_release or "").splitlines():
        if "=" not in line:
            continue
        k, v = line.split("=", 1)
        os_info[k.strip()] = v.strip().strip('"')

    os_id = (os_info.get("ID") or "").lower()
    ver = os_info.get("VERSION_ID") or ""
    # pretty = os_info.get("PRETTY_NAME") or f"{os_id} {ver}".strip()
    # ui.info(f"Target OS: {pretty}")

    if os_id == "ubuntu" and ver.startswith("24.04"):
        raise ExecutorError(
            "Dataiku dependency installer does not support Ubuntu 24.04 for this DSS version.\n"
            "Use a supported OS (commonly Ubuntu 22.04) or try a newer DSS version if Dataiku supports it."
        )


def build_install_steps(ui: Ui, ex: Executor, plan: DataikuInstallPlan, *, force: bool = False) -> list[Step]:
    """Build Dataiku DSS install steps based on the plan."""

    def _preflight():
        run_preflight(ui, ex, plan)

    def _users_and_dirs():
        ui.log("Creating Dataiku OS user/group and required directories (idempotent).")
        ensure_group(ex, plan.os_group)
        ensure_user(ex, user=plan.os_user, group=plan.os_group)
        ensure_dir_owned(ex, plan.root_dir, owner=plan.os_user, group=plan.os_group)
        ensure_dir_owned(ex, plan.install_dir, owner=plan.os_user, group=plan.os_group)
        ensure_dir_owned(ex, plan.data_dir, owner=plan.os_user, group=plan.os_group)

    def _install_prereqs():
        ui.log("Installing base packages (apt): ca-certificates, curl, gnupg, tar, wget.")
        ensure_apt_updated(ex)
        apt_install(ex, ["ca-certificates", "curl", "gnupg", "tar", "wget"])

    def _ensure_jdbc_driver():
        """Install a JDBC driver for Dataiku (used by Govern and some connections).

        For Govern, PostgreSQL JDBC is required. We install it via apt and copy
        the jar into DATA_DIR/lib/jdbc (the location documented by Dataiku).
        """
        jdbc_dir = f"{plan.data_dir.rstrip('/')}/lib/jdbc"
        ui.log(f"Ensuring JDBC driver directory exists: {jdbc_dir}")
        check(ex, f"sudo -n mkdir -p {q(jdbc_dir)}")
        check(ex, f"sudo -n chown -R {q(plan.os_user)}:{q(plan.os_group)} {q(plan.data_dir)}")

        ui.log("Installing PostgreSQL JDBC driver (apt: libpostgresql-jdbc-java).")
        ensure_apt_updated(ex)
        apt_install(ex, ["libpostgresql-jdbc-java"])

        # Copy the jar into Dataiku's expected location.
        jar_src = "/usr/share/java/postgresql.jar"
        jar_dst = f"{jdbc_dir}/postgresql.jar"
        check(
            ex,
            f"sudo -n cp {q(jar_src)} {q(jar_dst)}",
            error_hint="PostgreSQL JDBC jar not found. Ensure libpostgresql-jdbc-java is installed.",
        )
        check(ex, f"sudo -n chown {q(plan.os_user)}:{q(plan.os_group)} {q(jar_dst)}")
        check(ex, f"sudo -n chmod 644 {q(jar_dst)}")

    def _configure_govern_db():
        """Write Govern PostgreSQL settings to config/dip.properties (idempotent)."""
        dip_path = f"{plan.data_dir.rstrip('/')}/{GOVERN_DIP_PROPERTIES}"
        ui.log(f"Configuring Govern PostgreSQL settings in {dip_path}.")

        # Encrypt password if requested and possible.
        encrypted_password = _encrypt_govern_password(ex, plan, ui=ui)

        password_value = encrypted_password or (plan.govern_jdbc_password or "")
        if encrypted_password:
            ui.log("Govern DB password encrypted successfully.")
        elif plan.encrypt_govern_password:
            ui.warning(
                "Could not encrypt the Govern DB password automatically. "
                "It will be written in cleartext to dip.properties. "
                "You can later encrypt it using: DATA_DIR/bin/govern-admin encrypt-password"
            )
        else:
            ui.log("Password encryption skipped (--no-encrypt-govern-password).")

        new_content = (
            "# Managed by k4s\n"
            f"psql.jdbc.url={plan.govern_jdbc_url or ''}\n"
            f"psql.jdbc.user={plan.govern_jdbc_user or ''}\n"
            f"psql.jdbc.password={password_value}\n"
        )

        # Write without leaking content to debug output.
        check(ex, f"sudo -n tee {q(dip_path)} >/dev/null", stdin_data=new_content, silent=True)
        check(ex, f"sudo -n chown {q(plan.os_user)}:{q(plan.os_group)} {q(dip_path)}", silent=True)

    def _govern_init_db():
        """Bootstrap Govern DB (first run only)."""
        if not plan.govern_init_db:
            ui.log("Skipping govern-admin init-db (disabled).")
            return
        ui.log("Bootstrapping Govern database (govern-admin init-db).")
        cmd = f"sudo -n -u {q(plan.os_user)} {q(plan.data_dir)}/bin/govern-admin init-db"
        # Best-effort: allow re-runs without failing the whole install.
        rc, out, err = run(ex, cmd)
        if rc != 0:
            msg = (err or out).strip()
            lowered = msg.lower()
            if "already" in lowered or "exists" in lowered or "schema version is not none" in lowered:
                ui.warning(f"govern-admin init-db did not complete cleanly, but may already be initialized:\n{msg}")
            else:
                raise ExecutorError(f"govern-admin init-db failed.\n{msg}")

    def _download_archive():
        if plan.archive_path:
            ui.log("Uploading local DSS archive to target machine (offline mode).")
            ex.upload_file(plan.archive_path, plan.archive_remote_path, owner_spec=f"{plan.os_user}:{plan.os_group}", use_sudo=True)
        else:
            ui.log("Downloading DSS archive from Dataiku CDN.")
            AssetDownloader(ex).download(plan.archive_url, plan.archive_remote_path, sudo=True)
            check(ex, f"sudo -n chown {q(plan.os_user)}:{q(plan.os_group)} {q(plan.archive_remote_path)}")

    def _extract():
        ui.log("Extracting DSS archive.")
        cmd = (
            f"sudo -n -u {q(plan.os_user)} -i -- sh -c "
            f"{q(f'tar xzf {plan.archive_remote_path} -C {plan.install_dir}')}"
        )
        check(ex, cmd)

    def _install_deps():
        ui.log("Running Dataiku dependency installer (install-deps.sh).")
        check(
            ex,
            f"sudo -n DEBIAN_FRONTEND=noninteractive {q(plan.extracted_dir)}/scripts/install/install-deps.sh -yes",
        )

    def _run_installer():
        installer_args = build_installer_args(plan)
        ui.log(f"Running DSS installer ({plan.node_type} node).")
        cmd = (
            f"sudo -n -u {q(plan.os_user)} -i -- sh -c "
            f"{q(f'cd {plan.extracted_dir} && ./installer.sh {installer_args}')}"
        )
        check(ex, cmd)

    def _service_name() -> str:
        """Return the systemd service name for this DSS node type."""
        return f"dataiku.{plan.node_type}"

    def _install_boot():
        """Register this DSS instance as a systemd service via install-boot.sh.

        Uses the -n flag to create a per-node service (e.g. dataiku.design,
        dataiku.automation) so multiple instances can coexist independently.
        """
        svc = _service_name()
        ui.log(f"Installing boot script (service: {svc}).")
        check(
            ex,
            f"sudo -n sh {q(plan.extracted_dir)}/scripts/install/install-boot.sh "
            f"-n {q(plan.node_type)} {q(plan.data_dir)} {q(plan.os_user)}",
        )

    def _start():
        svc = _service_name()
        ui.log(f"Enabling and starting {svc} via systemd.")
        check(ex, "sudo -n systemctl daemon-reload")
        check(ex, f"sudo -n systemctl enable {svc}")
        check(
            ex,
            f"sudo -n systemctl start {svc}",
            error_hint=f"DSS failed to start. Check port conflicts or run: journalctl -xeu {svc}.service",
        )
        ui.log("Waiting for DSS to become ready.")
        run(ex, "sleep 15")

    def _already_installed() -> bool:
        rc, _, _ = run(ex, f"test -x {q(plan.data_dir)}/bin/dss")
        return rc == 0

    def _wipe_existing():
        """Stop DSS and remove the data directory for a clean reinstall."""
        svc = _service_name()
        ui.log(f"Stopping {svc} service.")
        run(ex, f"sudo -n systemctl stop {svc} || true")
        run(ex, f"sudo -n systemctl disable {svc} || true")
        # Fallback: also try dss stop in case the service is not registered.
        run(ex, f"sudo -n -u {q(plan.os_user)} -i -- sh -c {q(f'{plan.data_dir}/bin/dss stop')} || true")
        ui.log(f"Removing existing data directory: {plan.data_dir}")
        check(ex, f"sudo -n rm -rf {q(plan.data_dir)}")
        ui.log(f"Recreating data directory: {plan.data_dir}")
        ensure_dir_owned(ex, plan.data_dir, owner=plan.os_user, group=plan.os_group)

    def _guarded_install():
        if _already_installed() and not force:
            raise ExecutorError(
                f"Dataiku already installed at {plan.data_dir}.\n"
                "  Use --force to reinstall."
            )
        if force and _already_installed():
            _wipe_existing()
        _download_archive()
        _extract()
        _install_deps()
        _run_installer()
        _install_boot()
        # Govern requires DB config + init-db before first start.
        if plan.node_type != "govern":
            _start()

    steps: list[Step] = [
        Step(title=f"Preflight (Dataiku) on {ex.host}", run=_preflight),
        Step(title="Create OS user/group and directories", run=_users_and_dirs),
        Step(title="Install system prerequisites", run=_install_prereqs),
    ]

    if plan.install_kubectl:
        def _install_kubectl():
            ui.log("Installing kubectl via official Kubernetes apt repository.")
            install_kubectl(ex)
        steps.append(Step(title="Install kubectl", run=_install_kubectl))

    if plan.install_docker:
        steps.extend(
            build_docker_steps(
                ui,
                ex,
                data_root=plan.docker_data_root,
            )
        )

    if plan.add_to_docker_group:
        def _add_docker_group():
            rc, _, _ = run(ex, "getent group docker >/dev/null 2>&1", silent=True)
            if rc != 0:
                ui.log("Docker group does not exist; skipping (install Docker first or use --with-docker).")
                return
            ui.log(f"Adding {plan.os_user} to the docker group.")
            ensure_user_in_group(ex, user=plan.os_user, group="docker")
        steps.append(Step(title=f"Add {plan.os_user} to docker group", run=_add_docker_group))

    node_label = {
        "design": "Design",
        "automation": "Automation",
        "govern": "Govern",
    }.get(plan.node_type, plan.node_type)
    steps.append(
        Step(title=f"Install Dataiku DSS ({node_label} node)", run=_guarded_install)
    )

    # JDBC driver must be installed AFTER installer.sh because the installer
    # recreates DATA_DIR and would wipe any files placed there earlier.
    if plan.with_jdbc or plan.node_type == "govern":
        steps.append(Step(title="Install JDBC driver", run=_ensure_jdbc_driver))

    if plan.node_type == "govern":
        steps.append(Step(title="Configure Govern database connection", run=_configure_govern_db))
        steps.append(Step(title="Initialize Govern database (init-db)", run=_govern_init_db))
        steps.append(Step(title="Start Dataiku Govern", run=_start))

    if plan.install_r:
        steps.extend(
            build_r_steps(
                ui,
                ex,
                extracted_dir=plan.extracted_dir,
                data_dir=plan.data_dir,
                os_user=plan.os_user,
                r_repo=plan.r_repo,
                r_pkg_dir=plan.r_pkg_dir,
            )
        )

    if plan.os_user_password:
        def _set_password():
            ui.log(f"Setting password for OS user '{plan.os_user}'.")
            # Use chpasswd which reads from stdin (password never appears in ps/logs).
            rc, out, err = ex.execute(
                "sudo -n chpasswd",
                stdin_data=f"{plan.os_user}:{plan.os_user_password}\n",
            )
            if rc != 0:
                raise ExecutorError(f"Failed to set OS user password (rc={rc}).\n{err or out}")

        steps.append(Step(title=f"Set password for '{plan.os_user}'", run=_set_password))

    return steps
