#!/usr/bin/env python3

from __future__ import annotations

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


# 项目根目录，后面要从这里读 .env。
ROOT = Path(__file__).resolve().parents[1]
# 默认 .env 放在项目根目录，并且不入库。
ENV_FILE = Path(os.environ.get("APP_ENV_FILE", ROOT / ".env")).expanduser()
# 默认 deploy 机器实例 ID。
DEFAULT_INSTANCE_ID = os.environ.get("DEPLOY_INSTANCE_ID", "ins-ntmpcd3a")
# 默认区域。
DEFAULT_REGION = os.environ.get("DEPLOY_REGION", "me-saudi-arabia")
# 默认开放的端口。
DEFAULT_PORT = int(os.environ.get("OPEN_PORT", "2026"))
# 默认允许所有来源访问该端口。
DEFAULT_CIDR = os.environ.get("OPEN_CIDR", "0.0.0.0/0")
# 泛型只用于 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 用来查实例绑定的安全组。
from tencentcloud.cvm.v20170312 import cvm_client, models as cvm_models
# TAT 用来远程执行放行脚本。
from tencentcloud.tat.v20201028 import models as tat_models, tat_client


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

    # 逐行解析 shell 风格的 KEY=VALUE。
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        # 去掉空白方便判断。
        line = raw_line.strip()
        # 空行直接跳过。
        if not line:
            continue
        # 注释行直接跳过。
        if line.startswith("#"):
            continue
        # 无等号行直接跳过。
        if "=" not in line:
            continue

        # 只按第一个等号拆分。
        key, value = line.split("=", 1)
        # 清理 key 空白。
        key = key.strip()
        # 清理 value 空白。
        value = value.strip()

        # 去掉成对双引号。
        if value.startswith('"') and value.endswith('"'):
            value = value[1:-1]

        # 去掉成对单引号。
        if value.startswith("'") and value.endswith("'"):
            value = value[1:-1]

        # 只有环境里还没有这个 key 时才写入。
        os.environ.setdefault(key, value)


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


def env_required(name: str) -> str:
    # 读取环境变量并去掉两端空白。
    value = os.environ.get(name, "").strip()
    # 缺失时直接退出，避免后面报一串 SDK 错误。
    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)
    # 返回 CVM 客户端。
    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)
    # 返回 TAT 客户端。
    return tat_client.TatClient(
        cred,
        region,
        ClientProfile(httpProfile=HttpProfile(endpoint="tat.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
            # 最后一次重试不再 sleep，直接退出循环。
            if attempt == retries:
                break
            # 简单线性退避，降低偶发 EOF 的影响。
            time.sleep(delay_seconds * attempt)

    # 有异常就把最后一次异常抛出去。
    if last_error is not None:
        raise last_error

    # 理论上不会到这里，这里只是兜底。
    raise RuntimeError(f"{label} failed without error")


def resolve_security_group_id(client: cvm_client.CvmClient, instance_id: str) -> str:
    # 创建查询实例详情的请求对象。
    req = cvm_models.DescribeInstancesRequest()
    # 只查询目标实例。
    req.from_json_string(json.dumps({"InstanceIds": [instance_id]}))
    # 发起实例查询请求。
    resp = json.loads(invoke(lambda: client.DescribeInstances(req), "DescribeInstances").to_json_string())
    # 拿到实例数组。
    instances = resp.get("InstanceSet") or []

    # 实例不存在时直接退出。
    if not instances:
        raise SystemExit(f"instance not found: {instance_id}")

    # 读取第一个绑定安全组。
    security_group_ids = instances[0].get("SecurityGroupIds") or []

    # 没有绑定安全组时直接退出。
    if not security_group_ids:
        raise SystemExit(f"no security group attached: {instance_id}")

    # 返回第一个安全组 ID。
    return str(security_group_ids[0])


def remote_script(secret_id: str, secret_key: str, region: str, security_group_id: str, port: int, cidr: str) -> str:
    # 把必要参数序列化后再 base64，方便塞进远端 Python 脚本。
    payload = {
        "secret_id": secret_id,
        "secret_key": secret_key,
        "region": region,
        "security_group_id": security_group_id,
        "port": port,
        "cidr": cidr,
    }

    # 远端脚本里直接解码这个 blob。
    blob = base64.b64encode(json.dumps(payload).encode("utf-8")).decode("ascii")

    # 返回完整远端 shell 脚本。
    return f"""set -euo pipefail

if command -v firewall-cmd >/dev/null 2>&1; then
  firewall-cmd --add-port={port}/tcp --permanent >/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 {port}/tcp >/dev/null 2>&1 || true
fi

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

python3 -m pip install --quiet --disable-pip-version-check tencentcloud-sdk-python >/dev/null 2>&1 || true

python3 - <<'PY'
import base64
import json

from tencentcloud.common import credential
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.vpc.v20170312 import models, vpc_client

# 读取本次放行需要的参数。
payload = json.loads(base64.b64decode({blob!r}).decode('utf-8'))
# 构造凭据对象。
cred = credential.Credential(payload['secret_id'], payload['secret_key'])
# 初始化 VPC 客户端，用来写安全组规则。
client = vpc_client.VpcClient(
    cred,
    payload['region'],
    ClientProfile(httpProfile=HttpProfile(endpoint='vpc.tencentcloudapi.com')),
)

# 构造创建安全组入站规则的请求。
req = models.CreateSecurityGroupPoliciesRequest()
req.from_json_string(json.dumps({{
    'SecurityGroupId': payload['security_group_id'],
    'SecurityGroupPolicySet': {{
        'Ingress': [{{
            'Protocol': 'TCP',
            'Port': str(payload['port']),
            'CidrBlock': payload['cidr'],
            'Action': 'ACCEPT',
            'PolicyDescription': f"hy-app-monitor {{payload['port']}}",
        }}]
    }}
}}))

try:
    # 创建安全组规则。
    client.CreateSecurityGroupPolicies(req)
    print('security_group_rule=created')
except TencentCloudSDKException as exc:
    # 已存在时按成功处理，避免重复调用失败。
    if 'Duplicate' in str(exc) or 'already exists' in str(exc):
        print('security_group_rule=exists')
    else:
        raise
PY

echo "port_opened={port}"
"""


def run_tat_script(client: tat_client.TatClient, instance_id: str, script: str, timeout_seconds: int = 300) -> str:
    # 创建执行命令请求对象。
    req = tat_models.RunCommandRequest()
    # 把 shell 脚本下发到目标实例执行。
    req.from_json_string(
        json.dumps(
            {
                "CommandName": "hy-app-monitor-open-port",
                "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。
    invocation_id = resp["InvocationId"]
    # 设置轮询截止时间。
    deadline = time.time() + timeout_seconds + 120

    # 循环轮询任务结果。
    while time.time() < deadline:
        # 两次轮询之间稍微停一下。
        time.sleep(3)
        # 创建查询任务结果请求。
        query = tat_models.DescribeInvocationTasksRequest()
        # 按 invocation id 查询本次任务。
        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["TaskStatus"]

        # 进入终态时退出轮询。
        if status in {"SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "START_FAILED"}:
            # 解码远端标准输出。
            output = base64.b64decode(task.get("TaskResult", {}).get("Output", "")).decode("utf-8", errors="replace")
            # 非 SUCCESS 直接退出并带上远端输出。
            if status != "SUCCESS":
                raise SystemExit(output or f"TAT task failed: {status}")
            # SUCCESS 时返回远端输出。
            return output

    # 超过截止时间还没完成就直接报超时。
    raise SystemExit("TAT task timed out")


def main() -> None:
    # 从 .env 或外部环境中读取腾讯云密钥。
    secret_id = env_required("TENCENT_SECRET_ID")
    # 从 .env 或外部环境中读取腾讯云密钥。
    secret_key = env_required("TENCENT_SECRET_KEY")
    # 读取目标实例 ID。
    instance_id = os.environ.get("DEPLOY_INSTANCE_ID", DEFAULT_INSTANCE_ID)
    # 读取区域。
    region = os.environ.get("DEPLOY_REGION", DEFAULT_REGION)
    # 读取待放行端口。
    port = int(os.environ.get("OPEN_PORT", str(DEFAULT_PORT)))
    # 读取允许访问的网段。
    cidr = os.environ.get("OPEN_CIDR", DEFAULT_CIDR)

    # 初始化 CVM 客户端，用来查安全组。
    cvm = build_cvm(secret_id, secret_key, region)
    # 初始化 TAT 客户端，用来远程执行脚本。
    tat = build_tat(secret_id, secret_key, region)
    # 优先用显式指定的安全组，否则自动从实例上解析。
    security_group_id = os.environ.get("SECURITY_GROUP_ID", "").strip() or resolve_security_group_id(cvm, instance_id)

    # 远程执行端口放行脚本。
    output = run_tat_script(tat, instance_id, remote_script(secret_id, secret_key, region, security_group_id, port, cidr))
    # 打印实例 ID 方便确认目标机器。
    print(f"instance_id={instance_id}")
    # 打印实际使用的安全组 ID。
    print(f"security_group_id={security_group_id}")
    # 打印远端执行结果。
    print(output.strip())


if __name__ == "__main__":
    # 作为脚本执行时直接走主流程。
    main()