#!/usr/bin/env python3

from __future__ import annotations

import argparse
import base64
import json
import os
import shlex
import site
import sys
import time
from pathlib import Path
from typing import Callable, TypeVar


# 项目根目录，后面需要从这里读取 .env 和 targets.json。
ROOT = Path(__file__).resolve().parents[1]
# 默认 .env 路径。
ENV_FILE = Path(os.environ.get("APP_ENV_FILE", ROOT / ".env")).expanduser()
# 泛型只用于 invoke 的返回值标注。
T = TypeVar("T")


def load_project_venv() -> None:
    # 如果当前本来就在项目虚拟环境里运行，就不需要重复处理。
    if str(ROOT / ".venv") in sys.prefix:
        return

    # 遍历项目 .venv 下可能存在的 site-packages 目录。
    for candidate in (ROOT / ".venv" / "lib").glob("python*/site-packages"):
        # 把依赖目录加入解释器搜索路径。
        site.addsitedir(str(candidate))


# 先注入项目虚拟环境，再导入腾讯云 SDK。
load_project_venv()

# 腾讯云凭据对象。
from tencentcloud.common import credential
# 腾讯云 SDK 统一异常类型。
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
# 客户端公共配置。
from tencentcloud.common.profile.client_profile import ClientProfile
# HTTP endpoint 配置。
from tencentcloud.common.profile.http_profile import HttpProfile
# CVM 用来解析实例私网 IP 和安全组。
from tencentcloud.cvm.v20170312 import cvm_client, models as cvm_models
# TAT 用来做一次性带外引导。
from tencentcloud.tat.v20201028 import models as tat_models, tat_client
# VPC 用来写安全组规则。
from tencentcloud.vpc.v20170312 import models as vpc_models, vpc_client


def load_env_file(path: Path) -> None:
    # .env 文件不存在时允许继续使用外部环境变量。
    if not path.exists():
        return

    # 逐行解析 KEY=VALUE。
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip()
        if value.startswith('"') and value.endswith('"'):
            value = value[1:-1]
        if value.startswith("'") and value.endswith("'"):
            value = value[1:-1]
        os.environ.setdefault(key, value)


def env_truthy(name: str, default: bool = False) -> bool:
    # 兼容 .env 里常见的布尔字符串写法。
    raw = os.environ.get(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


# 先加载项目根目录的 .env。
load_env_file(ENV_FILE)

# 默认 deploy 实例 ID。
DEFAULT_DEPLOY_INSTANCE_ID = os.environ.get("DEPLOY_INSTANCE_ID", "").strip()
# 默认区域。
DEFAULT_REGION = os.environ.get("DEPLOY_REGION", "").strip()
# 默认监控目标文件。
TARGETS_FILE = Path(os.environ.get("TARGETS_FILE", ROOT / "config" / "targets.json")).expanduser()
# deploy 机上生成的专用 SSH 私钥路径。
DEFAULT_REMOTE_SSH_KEY_PATH = os.environ.get("SSH_PRIVATE_KEY_PATH", "/opt/hy-app-monitor/run/ssh/id_ed25519").strip()
# 默认远端 SSH 端口。
DEFAULT_SSH_PORT = int(os.environ.get("SSH_PORT", "22"))
# 默认远端 SSH 用户。
DEFAULT_SSH_USER = os.environ.get("SSH_USER", "root").strip() or "root"
# 独立 admin 前端目标机配置；配置齐备时会一并加入 SSH 引导。
ADMIN_FRONTEND_ENABLED = env_truthy("ADMIN_FRONTEND_ENABLED", False)
ADMIN_FRONTEND_HOST = os.environ.get("ADMIN_FRONTEND_HOST", "").strip()
ADMIN_FRONTEND_IP = os.environ.get("ADMIN_FRONTEND_IP", "").strip()
ADMIN_FRONTEND_INSTANCE_ID = os.environ.get("ADMIN_FRONTEND_INSTANCE_ID", "").strip()
# 独立 chatapp-cron 目标机配置；配置齐备时也会加入 SSH 引导。
CHATAPP_CRON_ENABLED = env_truthy("CHATAPP_CRON_ENABLED", False)
CHATAPP_CRON_HOST = os.environ.get("CHATAPP_CRON_HOST", "").strip()
CHATAPP_CRON_IP = os.environ.get("CHATAPP_CRON_IP", "").strip()
CHATAPP_CRON_INSTANCE_ID = os.environ.get("CHATAPP_CRON_INSTANCE_ID", "").strip()


def parse_args() -> argparse.Namespace:
    # 默认对 targets.json 里的所有实例做 SSH 引导。
    parser = argparse.ArgumentParser(description="bootstrap deploy->target ssh connectivity via TAT")
    parser.add_argument("--hosts", nargs="*", default=[], help="target host names, default all hosts from targets.json")
    parser.add_argument("--deploy-instance-id", default=DEFAULT_DEPLOY_INSTANCE_ID, help="deploy machine instance id")
    parser.add_argument("--region", default=DEFAULT_REGION, help="tencent cloud region")
    parser.add_argument("--ssh-user", default=DEFAULT_SSH_USER, help="remote ssh login user, default root")
    parser.add_argument("--ssh-port", type=int, default=DEFAULT_SSH_PORT, help="remote ssh port")
    parser.add_argument(
        "--remote-key-path",
        default=DEFAULT_REMOTE_SSH_KEY_PATH,
        help="private key path generated on deploy machine",
    )
    parser.add_argument(
        "--deploy-ip",
        default="",
        help="override deploy private ip used in security group and firewall rules",
    )
    return parser.parse_args()


def env_required(name: str) -> str:
    # 读取环境变量并去掉两端空白。
    value = os.environ.get(name, "").strip()
    if not value:
        raise SystemExit(f"missing env: {name}")
    return value


def build_cvm(secret_id: str, secret_key: str, region: str) -> cvm_client.CvmClient:
    # 构造腾讯云凭据对象。
    cred = credential.Credential(secret_id, secret_key)
    return cvm_client.CvmClient(
        cred,
        region,
        ClientProfile(httpProfile=HttpProfile(endpoint="cvm.tencentcloudapi.com")),
    )


def build_tat(secret_id: str, secret_key: str, region: str) -> tat_client.TatClient:
    # 构造腾讯云凭据对象。
    cred = credential.Credential(secret_id, secret_key)
    return tat_client.TatClient(
        cred,
        region,
        ClientProfile(httpProfile=HttpProfile(endpoint="tat.tencentcloudapi.com")),
    )


def build_vpc(secret_id: str, secret_key: str, region: str) -> vpc_client.VpcClient:
    # 构造腾讯云凭据对象。
    cred = credential.Credential(secret_id, secret_key)
    return vpc_client.VpcClient(
        cred,
        region,
        ClientProfile(httpProfile=HttpProfile(endpoint="vpc.tencentcloudapi.com")),
    )


def invoke(action: Callable[[], T], label: str, retries: int = 4, delay_seconds: int = 2) -> T:
    # 保存最后一次异常，方便重试结束后抛出。
    last_error: Exception | None = None
    for attempt in range(1, retries + 1):
        try:
            return action()
        except TencentCloudSDKException as exc:
            last_error = exc
            if attempt == retries:
                break
            time.sleep(delay_seconds * attempt)

    if last_error is not None:
        raise last_error
    raise RuntimeError(f"{label} failed without error")


def describe_instances(client: cvm_client.CvmClient, instance_ids: list[str]) -> dict[str, dict[str, object]]:
    # 一次性把 deploy 和目标机实例详情取回来，减少 API 往返。
    req = cvm_models.DescribeInstancesRequest()
    req.from_json_string(json.dumps({"InstanceIds": instance_ids}))
    resp = json.loads(invoke(lambda: client.DescribeInstances(req), "DescribeInstances").to_json_string())
    items = resp.get("InstanceSet") or []
    return {
        str(item.get("InstanceId") or ""): {
            "privateIp": str((item.get("PrivateIpAddresses") or [""])[0] or "").strip(),
            "securityGroupIds": [str(group_id) for group_id in (item.get("SecurityGroupIds") or []) if str(group_id)],
            "instanceName": str(item.get("InstanceName") or "").strip(),
        }
        for item in items
    }


def load_target_hosts(host_names: list[str]) -> list[dict[str, str]]:
    # 读取 targets.json，按 host 名过滤后返回唯一实例列表。
    if not TARGETS_FILE.exists():
        raise SystemExit(f"targets file not found: {TARGETS_FILE}")
    payload = json.loads(TARGETS_FILE.read_text(encoding="utf-8"))
    configured = payload.get("hosts") or []
    selected = {str(name).strip() for name in host_names if str(name).strip()}
    items: list[dict[str, str]] = []
    seen_instance_ids: set[str] = set()

    for host in configured:
        host_name = str(host.get("name") or "").strip()
        instance_id = str(host.get("instanceId") or "").strip()
        if not host_name or not instance_id:
            continue
        if selected and host_name not in selected:
            continue
        if instance_id in seen_instance_ids:
            continue
        seen_instance_ids.add(instance_id)
        items.append(
            {
                "host": host_name,
                "ip": str(host.get("ip") or "").strip(),
                "instanceId": instance_id,
            }
        )

    # Admin 前端机器不在 targets.json 时，也允许通过 .env 把它并进同一轮 SSH 引导。
    if ADMIN_FRONTEND_ENABLED and ADMIN_FRONTEND_HOST and ADMIN_FRONTEND_INSTANCE_ID:
        if (not selected) or ADMIN_FRONTEND_HOST in selected:
            if ADMIN_FRONTEND_INSTANCE_ID not in seen_instance_ids:
                seen_instance_ids.add(ADMIN_FRONTEND_INSTANCE_ID)
                items.append(
                    {
                        "host": ADMIN_FRONTEND_HOST,
                        "ip": ADMIN_FRONTEND_IP,
                        "instanceId": ADMIN_FRONTEND_INSTANCE_ID,
                    }
                )
    # 定时任务机器也不在 targets.json 时，通过 .env 并入同一轮 SSH 引导。
    if CHATAPP_CRON_ENABLED and CHATAPP_CRON_HOST and CHATAPP_CRON_INSTANCE_ID:
        if (not selected) or CHATAPP_CRON_HOST in selected:
            if CHATAPP_CRON_INSTANCE_ID not in seen_instance_ids:
                seen_instance_ids.add(CHATAPP_CRON_INSTANCE_ID)
                items.append(
                    {
                        "host": CHATAPP_CRON_HOST,
                        "ip": CHATAPP_CRON_IP,
                        "instanceId": CHATAPP_CRON_INSTANCE_ID,
                    }
                )
    return items


def render_keygen_script(remote_key_path: str) -> str:
    # deploy 机上只生成一次专用 SSH key，后续运行链路直接复用。
    return f"""set -euo pipefail
KEY_PATH={json.dumps(remote_key_path)}
KEY_DIR="$(dirname "$KEY_PATH")"
mkdir -p "$KEY_DIR"
chmod 700 "$KEY_DIR"
if [ ! -f "$KEY_PATH" ]; then
  ssh-keygen -t ed25519 -N '' -C 'hy-app-monitor-deploy' -f "$KEY_PATH" >/dev/null
fi
touch "$KEY_DIR/known_hosts"
chmod 600 "$KEY_PATH" "$KEY_PATH.pub" "$KEY_DIR/known_hosts"
cat "$KEY_PATH.pub"
"""


def render_target_bootstrap_script(public_key: str, deploy_ip: str, ssh_user: str, ssh_port: int) -> str:
    # 目标机侧只做 SSH 基础设施准备，不夹带业务运维逻辑。
    quoted_key = shlex.quote(public_key)
    quoted_ip = shlex.quote(deploy_ip)
    quoted_user = shlex.quote(ssh_user)
    root_login = "yes" if ssh_user == "root" else "no"

    if ssh_user == "root":
        principal_block = """
HOME_DIR=/root
install -d -m 700 "$HOME_DIR/.ssh"
touch "$HOME_DIR/.ssh/authorized_keys"
chmod 600 "$HOME_DIR/.ssh/authorized_keys"
grep -qxF "$PUBKEY" "$HOME_DIR/.ssh/authorized_keys" || printf '%s\\n' "$PUBKEY" >> "$HOME_DIR/.ssh/authorized_keys"
"""
    else:
        principal_block = """
if ! id -u "$SSH_USER" >/dev/null 2>&1; then
  useradd -m -s /bin/bash "$SSH_USER"
fi
if getent group docker >/dev/null 2>&1; then
  usermod -aG docker "$SSH_USER" >/dev/null 2>&1 || true
fi
HOME_DIR="$(getent passwd "$SSH_USER" | cut -d: -f6)"
install -d -o "$SSH_USER" -g "$SSH_USER" -m 700 "$HOME_DIR/.ssh"
touch "$HOME_DIR/.ssh/authorized_keys"
chown "$SSH_USER":"$SSH_USER" "$HOME_DIR/.ssh/authorized_keys"
chmod 600 "$HOME_DIR/.ssh/authorized_keys"
grep -qxF "$PUBKEY" "$HOME_DIR/.ssh/authorized_keys" || printf '%s\\n' "$PUBKEY" >> "$HOME_DIR/.ssh/authorized_keys"
chown "$SSH_USER":"$SSH_USER" "$HOME_DIR/.ssh/authorized_keys"
printf '%s\\n' "$SSH_USER ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/hy-app-monitor-ssh-user
chmod 440 /etc/sudoers.d/hy-app-monitor-ssh-user
"""

    return f"""set -euo pipefail
PUBKEY={quoted_key}
DEPLOY_IP={quoted_ip}
SSH_USER={quoted_user}
SSH_PORT={int(ssh_port)}
DEPLOY_CIDR="${{DEPLOY_IP}}/32"

if ! command -v sshd >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then
    export DEBIAN_FRONTEND=noninteractive
    apt-get update -y >/dev/null 2>&1 || true
    apt-get install -y openssh-server >/dev/null
  elif command -v dnf >/dev/null 2>&1; then
    dnf install -y openssh-server >/dev/null
  elif command -v yum >/dev/null 2>&1; then
    yum install -y openssh-server >/dev/null
  else
    echo "openssh-server unavailable" >&2
    exit 1
  fi
fi

{principal_block}

mkdir -p /etc/ssh/sshd_config.d
if ! grep -Eq '^[[:space:]]*Include[[:space:]]+/etc/ssh/sshd_config\\.d/\\*\\.conf' /etc/ssh/sshd_config; then
  printf '\\nInclude /etc/ssh/sshd_config.d/*.conf\\n' >> /etc/ssh/sshd_config
fi

cat > /etc/ssh/sshd_config.d/hy-app-monitor.conf <<EOF
PubkeyAuthentication yes
AuthorizedKeysFile .ssh/authorized_keys
PermitRootLogin {root_login}
Port {int(ssh_port)}
EOF

if command -v firewall-cmd >/dev/null 2>&1; then
  firewall-cmd --permanent --add-rich-rule="rule family=ipv4 source address=${{DEPLOY_CIDR}} port protocol=tcp port=${{SSH_PORT}} accept" >/dev/null 2>&1 || true
  firewall-cmd --reload >/dev/null 2>&1 || true
fi

if command -v ufw >/dev/null 2>&1; then
  ufw allow from "${{DEPLOY_IP}}" to any port "${{SSH_PORT}}" proto tcp >/dev/null 2>&1 || true
fi

if command -v iptables >/dev/null 2>&1; then
  iptables -C INPUT -s "${{DEPLOY_CIDR}}" -p tcp --dport "${{SSH_PORT}}" -j ACCEPT >/dev/null 2>&1 || iptables -I INPUT -s "${{DEPLOY_CIDR}}" -p tcp --dport "${{SSH_PORT}}" -j ACCEPT || true
fi

if command -v sshd >/dev/null 2>&1; then
  sshd -t
fi

if systemctl list-unit-files sshd.service >/dev/null 2>&1; then
  systemctl enable sshd >/dev/null 2>&1 || true
  systemctl restart sshd
elif systemctl list-unit-files ssh.service >/dev/null 2>&1; then
  systemctl enable ssh >/dev/null 2>&1 || true
  systemctl restart ssh
else
  echo "ssh systemd service not found" >&2
  exit 1
fi

echo "ssh_bootstrap_ok user=${{SSH_USER}} port=${{SSH_PORT}}"
"""


def run_tat_script(client: tat_client.TatClient, instance_id: str, script: str, timeout_seconds: int = 600) -> str:
    # 创建执行命令请求。
    req = tat_models.RunCommandRequest()
    req.from_json_string(
        json.dumps(
            {
                "CommandName": "hy-app-monitor-bootstrap-ssh",
                "CommandType": "SHELL",
                "Content": base64.b64encode(script.encode("utf-8")).decode("ascii"),
                "InstanceIds": [instance_id],
                "Timeout": timeout_seconds,
                "WorkingDirectory": "/root",
            }
        )
    )

    # 发起命令执行。
    resp = json.loads(invoke(lambda: client.RunCommand(req), "RunCommand").to_json_string())
    invocation_id = resp["InvocationId"]
    deadline = time.time() + timeout_seconds + 120

    # 轮询远端执行结果。
    while time.time() < deadline:
        time.sleep(3)
        query = tat_models.DescribeInvocationTasksRequest()
        query.from_json_string(
            json.dumps(
                {
                    "HideOutput": False,
                    "Limit": 20,
                    "Filters": [{"Name": "invocation-id", "Values": [invocation_id]}],
                }
            )
        )
        data = json.loads(invoke(lambda: client.DescribeInvocationTasks(query), "DescribeInvocationTasks").to_json_string())
        tasks = data.get("InvocationTaskSet") or []
        if not tasks:
            continue
        task = tasks[0]
        status = task.get("TaskStatus") or ""
        if status in {"SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "START_FAILED"}:
            output = base64.b64decode(task.get("TaskResult", {}).get("Output", "")).decode("utf-8", errors="replace")
            if status != "SUCCESS":
                raise SystemExit(output or f"TAT task failed: {status}")
            return output

    raise SystemExit("TAT task timed out")


def ensure_security_group_rule(
    client: vpc_client.VpcClient,
    security_group_id: str,
    *,
    deploy_ip: str,
    ssh_port: int,
) -> str:
    # 安全组只允许 deploy 私网 IP 访问 22，避免直接对整段网开放 SSH。
    req = vpc_models.CreateSecurityGroupPoliciesRequest()
    req.from_json_string(
        json.dumps(
            {
                "SecurityGroupId": security_group_id,
                "SecurityGroupPolicySet": {
                    "Ingress": [
                        {
                            "Protocol": "TCP",
                            "Port": str(ssh_port),
                            "CidrBlock": f"{deploy_ip}/32",
                            "Action": "ACCEPT",
                            "PolicyDescription": f"hy-app-monitor ssh {deploy_ip}",
                        }
                    ]
                },
            }
        )
    )

    try:
        invoke(lambda: client.CreateSecurityGroupPolicies(req), "CreateSecurityGroupPolicies")
        return "created"
    except TencentCloudSDKException as exc:
        if "Duplicate" in str(exc) or "already exists" in str(exc):
            return "exists"
        raise


def prepare_deploy_public_key(
    client: tat_client.TatClient,
    deploy_instance_id: str,
    remote_key_path: str,
) -> str:
    # deploy 机上生成完专用 key 后，把公钥文本拿回来继续给目标机铺设 authorized_keys。
    output = run_tat_script(client, deploy_instance_id, render_keygen_script(remote_key_path), timeout_seconds=300)
    public_key = output.strip().splitlines()[-1].strip()
    if not public_key.startswith("ssh-"):
        raise SystemExit(f"invalid public key output: {public_key or output.strip()}")
    return public_key


def main() -> None:
    args = parse_args()
    if not args.deploy_instance_id:
        raise SystemExit("deploy instance id is required; set DEPLOY_INSTANCE_ID or pass --deploy-instance-id")
    if not args.region:
        raise SystemExit("deploy region is required; set DEPLOY_REGION or pass --region")

    secret_id = env_required("TENCENT_SECRET_ID")
    secret_key = env_required("TENCENT_SECRET_KEY")
    targets = load_target_hosts(args.hosts)
    if not targets:
        raise SystemExit("no target hosts found in config/targets.json")

    cvm = build_cvm(secret_id, secret_key, args.region)
    tat = build_tat(secret_id, secret_key, args.region)
    vpc = build_vpc(secret_id, secret_key, args.region)

    instance_ids = [args.deploy_instance_id, *[target["instanceId"] for target in targets]]
    instance_map = describe_instances(cvm, instance_ids)
    deploy_detail = instance_map.get(args.deploy_instance_id)
    if not deploy_detail:
        raise SystemExit(f"deploy instance not found: {args.deploy_instance_id}")

    deploy_ip = str(args.deploy_ip or deploy_detail.get("privateIp") or "").strip()
    if not deploy_ip:
        raise SystemExit(f"missing deploy private ip: {args.deploy_instance_id}")

    public_key = prepare_deploy_public_key(tat, args.deploy_instance_id, args.remote_key_path)
    results: list[dict[str, object]] = []

    for target in targets:
        detail = instance_map.get(target["instanceId"])
        if not detail:
            raise SystemExit(f"instance not found: {target['instanceId']}")

        output = run_tat_script(
            tat,
            target["instanceId"],
            render_target_bootstrap_script(public_key, deploy_ip, args.ssh_user, args.ssh_port),
            timeout_seconds=600,
        ).strip()

        security_groups = list(detail.get("securityGroupIds") or [])
        sg_results = [
            {
                "securityGroupId": security_group_id,
                "status": ensure_security_group_rule(
                    vpc,
                    security_group_id,
                    deploy_ip=deploy_ip,
                    ssh_port=args.ssh_port,
                ),
            }
            for security_group_id in security_groups
        ]

        results.append(
            {
                "host": target["host"],
                "instanceId": target["instanceId"],
                "ip": target["ip"] or detail.get("privateIp") or "",
                "sshUser": args.ssh_user,
                "sshPort": args.ssh_port,
                "bootstrap": output,
                "securityGroups": sg_results,
            }
        )

    print(
        json.dumps(
            {
                "ok": True,
                "deployInstanceId": args.deploy_instance_id,
                "deployIp": deploy_ip,
                "remoteKeyPath": args.remote_key_path,
                "sshUser": args.ssh_user,
                "sshPort": args.ssh_port,
                "targets": results,
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()