#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ctypes
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.parse
import urllib.request
import uuid
import zipfile
from pathlib import Path

SKILL_NAME = "MML解析工具（知一 Zhiyisoft）"
SKILL_ID = "mml-parser-zhiyisoft"
SKILL_VERSION = "1.3.8"
MIN_EXE_VERSION = (1, 3, 8)
OFFICIAL_DOWNLOAD_PAGE = "https://zhiyisoft.cn/product/MMLjiexi/download"
OFFICIAL_MANIFEST_URL = "https://zhiyisoft.cn/api/products/MMLjiexi/download-manifest"
MAX_DOWNLOAD_BYTES = 512 * 1024 * 1024

_AI_BLOCKED_KEYS = {
    "activation_code", "license_token", "public_key_pem", "server_public_key_pem",
    "device_credential", "credential_id", "key_sha256", "public_key_sha256",
    "device_fingerprint", "device_summary", "private_key_pem", "request_ip",
    "machine_guid", "disk_serial", "volume_serial", "system_drive",
    "customer_contact", "customer_email", "email", "order_id",
    "traceback", "logs", "input_files", "input_paths", "exe_path", "found_exe",
    "integration_authorization_token", "integration_component_sha256", "integration_request_sha256",
}


def sanitize_ai_payload(value):
    """Defense-in-depth filter for every payload crossing into an AI client."""
    if isinstance(value, dict):
        return {
            str(key): sanitize_ai_payload(item)
            for key, item in value.items()
            if str(key).lower() not in _AI_BLOCKED_KEYS
        }
    if isinstance(value, list):
        return [sanitize_ai_payload(item) for item in value]
    return value


def emit(payload: dict, exit_code: int = 0) -> int:
    print(json.dumps(sanitize_ai_payload(payload), ensure_ascii=False, indent=2))
    return exit_code


def version_tuple(value: str) -> tuple[int, ...]:
    return tuple(int(item) for item in re.findall(r"\d+", str(value or ""))[:4])


def candidate_version(path: Path) -> tuple[int, ...]:
    match = re.search(r"(?:^|[_-])[Vv]?(\d+\.\d+\.\d+(?:\.\d+)?)", str(path))
    if match:
        return version_tuple(match.group(1))
    if os.name != "nt":
        return ()
    try:
        size = ctypes.windll.version.GetFileVersionInfoSizeW(str(path), None)
        if not size:
            return ()
        buffer = ctypes.create_string_buffer(size)
        if not ctypes.windll.version.GetFileVersionInfoW(str(path), 0, size, buffer):
            return ()
        value = ctypes.c_void_p()
        length = ctypes.c_uint32()
        if not ctypes.windll.version.VerQueryValueW(buffer, "\\", ctypes.byref(value), ctypes.byref(length)):
            return ()
        words = ctypes.cast(value, ctypes.POINTER(ctypes.c_uint32))
        return (words[2] >> 16, words[2] & 0xFFFF, words[3] >> 16, words[3] & 0xFFFF)
    except Exception:
        return ()


def common_roots() -> list[Path]:
    values = [
        Path.cwd(), Path.home() / "Downloads", Path.home() / "Desktop", Path.home() / "Documents",
        Path(os.environ.get("LOCALAPPDATA", "")) / "Programs",
        Path(os.environ.get("ProgramFiles", "")), Path(os.environ.get("ProgramFiles(x86)", "")),
        Path("E:/知一自研软件/MML解析工具"),
    ]
    result: list[Path] = []
    seen: set[str] = set()
    for path in values:
        if not str(path) or not path.exists():
            continue
        key = os.path.normcase(str(path.resolve()))
        if key not in seen:
            seen.add(key)
            result.append(path)
    return result


def find_exe(explicit: str = "") -> tuple[Path | None, Path | None]:
    candidates: list[Path] = []
    configured = explicit or os.environ.get("MML_PARSER_EXE", "")
    if configured:
        candidates.append(Path(configured).expanduser())
    for root in common_roots():
        for name in ("MML解析工具.exe", "MML解析工具_V*.exe"):
            candidates.extend(root.glob(name))
            candidates.extend(root.glob(f"*/{name}"))
            candidates.extend(root.glob(f"*/*/{name}"))
    seen: set[str] = set()
    newest_old: Path | None = None
    valid: list[tuple[tuple[int, ...], float, Path]] = []
    for path in candidates:
        try:
            path = path.resolve()
            key = os.path.normcase(str(path))
            if key in seen or not path.is_file():
                continue
            seen.add(key)
            version = candidate_version(path)
            if version >= MIN_EXE_VERSION:
                valid.append((version, path.stat().st_mtime, path))
            elif newest_old is None or path.stat().st_mtime > newest_old.stat().st_mtime:
                newest_old = path
        except OSError:
            continue
    valid.sort(reverse=True)
    return (valid[0][2] if valid else None), newest_old


def find_official_bridge() -> Path | None:
    candidates: list[Path] = []
    configured = str(os.environ.get("MML_PARSER_BRIDGE_EXE") or "").strip()
    if configured:
        candidates.append(Path(configured).expanduser())
    candidates.append(Path(sys.executable))
    module_root = Path(__file__).resolve().parent
    candidates.extend([
        module_root.parent / "assets" / f"MML解析工具_MCP_Server_V{SKILL_VERSION}.exe",
        module_root / "assets" / f"MML解析工具_MCP_Server_V{SKILL_VERSION}.exe",
        Path("E:/知一自研软件/MML解析工具") / f"MML解析工具_MCP_Server_V{SKILL_VERSION}.exe",
    ])
    for candidate in candidates:
        try:
            resolved = candidate.resolve(strict=True)
            if (
                resolved.is_file()
                and resolved.read_bytes()[:2] == b"MZ"
                and candidate_version(resolved) == MIN_EXE_VERSION
                and re.fullmatch(r"MML解析工具_MCP_Server_V1\.3\.8\.exe", resolved.name, re.IGNORECASE)
            ):
                return resolved
        except OSError:
            continue
    return None


def invoke_exe(exe: Path, request: dict, timeout: int) -> dict:
    bridge = find_official_bridge()
    if bridge is None:
        raise RuntimeError("未找到与当前版本一致的官方 Skill Bridge/MCP Server EXE。")
    with tempfile.TemporaryDirectory(prefix="mml-parser-zhiyisoft-") as td:
        request_path = Path(td) / "request.json"
        response_path = Path(td) / "response.json"
        request_path.write_text(json.dumps(request, ensure_ascii=False, indent=2), encoding="utf-8")
        completed = subprocess.run(
            [
                str(bridge),
                "--bridge-invoke",
                "--bridge-client-exe", str(exe),
                "--bridge-request", str(request_path),
                "--bridge-response", str(response_path),
            ],
            capture_output=True,
            text=True,
            timeout=timeout,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        if not response_path.exists():
            raise RuntimeError(f"官方桥接程序未生成响应（退出码 {completed.returncode}）。{completed.stderr[-500:] if completed.stderr else ''}")
        payload = json.loads(response_path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            raise RuntimeError("EXE 返回的 Skill 响应不是 JSON 对象。")
        return sanitize_ai_payload(payload)


def is_official_https(url: str) -> bool:
    parsed = urllib.parse.urlparse(url)
    host = (parsed.hostname or "").lower()
    return parsed.scheme == "https" and (host == "zhiyisoft.cn" or host.endswith(".zhiyisoft.cn"))


def load_official_manifest() -> dict:
    request = urllib.request.Request(OFFICIAL_MANIFEST_URL, headers={"User-Agent": f"MMLParserSkill/{SKILL_VERSION}"})
    with urllib.request.urlopen(request, timeout=30) as response:
        if not is_official_https(response.geturl()):
            raise RuntimeError("下载清单发生到非知一官网域名的重定向，已拒绝。")
        raw = response.read(1024 * 1024 + 1)
    if len(raw) > 1024 * 1024:
        raise RuntimeError("官网下载清单超过 1 MB 安全上限。")
    payload = json.loads(raw.decode("utf-8"))
    if not isinstance(payload, dict) or str(payload.get("manifest_version") or "") != "1.0":
        raise RuntimeError("官网下载清单格式无效。")
    return payload


def _extract_standalone_zip(archive_path: Path, target_dir: Path, relative_launcher: str, release_version: str) -> Path:
    launcher_parts = [part for part in relative_launcher.replace("\\", "/").split("/") if part]
    if len(launcher_parts) < 2 or ".." in launcher_parts or any(":" in part for part in launcher_parts):
        raise RuntimeError("官网下载清单中的启动路径无效。")
    expected_top = launcher_parts[0]
    final_root = target_dir / expected_top
    final_launcher = target_dir.joinpath(*launcher_parts)
    if final_launcher.is_file() and candidate_version(final_launcher) == version_tuple(release_version):
        return final_launcher.resolve()
    if final_root.exists():
        raise RuntimeError(f"目标目录已存在但版本无效，请先人工检查：{final_root}")
    with zipfile.ZipFile(archive_path) as bundle:
        members = bundle.infolist()
        if not members or len(members) > 5000:
            raise RuntimeError("下载 ZIP 文件数量异常。")
        names = set()
        total_uncompressed = 0
        for member in members:
            normalized = member.filename.replace("\\", "/")
            parts = [part for part in normalized.split("/") if part not in ("", ".")]
            if normalized.startswith("/") or ".." in parts or any(":" in part for part in parts):
                raise RuntimeError("下载 ZIP 包含不安全路径。")
            if ((member.external_attr >> 16) & 0o170000) == 0o120000:
                raise RuntimeError("下载 ZIP 不允许符号链接。")
            total_uncompressed += max(0, int(member.file_size))
            if total_uncompressed > 1024 * 1024 * 1024:
                raise RuntimeError("下载 ZIP 解压体积超过 1 GB 安全上限。")
            names.add(normalized)
        normalized_launcher = "/".join(launcher_parts)
        if normalized_launcher not in names:
            raise RuntimeError("下载 ZIP 缺少清单指定的启动 EXE。")
        with tempfile.TemporaryDirectory(prefix=".mml-parser-extract-", dir=str(target_dir)) as td:
            staging = Path(td)
            bundle.extractall(staging)
            staged_root = staging / expected_top
            staged_launcher = staging.joinpath(*launcher_parts)
            if not staged_launcher.is_file():
                raise RuntimeError("下载 ZIP 中的启动文件不是有效 Windows EXE。")
            with staged_launcher.open("rb") as stream:
                if stream.read(2) != b"MZ":
                    raise RuntimeError("下载 ZIP 中的启动文件不是有效 Windows EXE。")
            if candidate_version(staged_launcher) != version_tuple(release_version):
                raise RuntimeError("下载 ZIP 中的 EXE 版本与官网清单不一致。")
            shutil.move(str(staged_root), str(final_root))
    return final_launcher.resolve()


def download_official(destination: str) -> dict:
    manifest = load_official_manifest()
    exe_release = manifest.get("exe") if isinstance(manifest.get("exe"), dict) else {}
    direct = str(exe_release.get("download_url") or "").strip()
    expected_sha256 = str(exe_release.get("sha256") or "").strip().lower()
    release_version = str(exe_release.get("version") or "").strip()
    filename = Path(str(exe_release.get("filename") or "")).name
    package_type = str(exe_release.get("package_type") or "exe").strip().lower()
    relative_launcher = str(exe_release.get("launch_relative_path") or "").strip()
    if not exe_release.get("available") or not direct:
        return {"ok": False, "code": "manual_download_required", "message": "官网当前未发布可由 Skill 下载的 EXE，请在官网下载页核实。", "download_page_url": OFFICIAL_DOWNLOAD_PAGE, "manifest_url": OFFICIAL_MANIFEST_URL}
    expected_suffix = ".zip" if package_type == "standalone_zip" else ".exe"
    if not is_official_https(direct) or not urllib.parse.urlparse(direct).path.lower().endswith(expected_suffix):
        raise RuntimeError("官网下载清单中的软件地址或封装类型无效。")
    if not re.fullmatch(r"[0-9a-f]{64}", expected_sha256):
        raise RuntimeError("官网下载清单缺少有效的 EXE SHA-256，已拒绝自动下载。")
    if version_tuple(release_version) < MIN_EXE_VERSION:
        raise RuntimeError(f"官网最新 EXE {release_version or '未知版本'} 仍低于 Skill 要求，已拒绝调用。")
    target_dir = Path(destination).expanduser() if destination else Path.home() / "Downloads"
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / (filename or Path(urllib.parse.urlparse(direct).path).name or "MML解析工具.exe")
    download_request = urllib.request.Request(direct, headers={"User-Agent": f"MMLParserSkill/{SKILL_VERSION}"})
    total = 0
    digest = hashlib.sha256()
    try:
        with urllib.request.urlopen(download_request, timeout=60) as response, target.open("wb") as output:
            if not is_official_https(response.geturl()):
                raise RuntimeError("下载发生到非知一官网域名的重定向，已拒绝。")
            while True:
                chunk = response.read(1024 * 1024)
                if not chunk:
                    break
                total += len(chunk)
                if total > MAX_DOWNLOAD_BYTES:
                    raise RuntimeError("下载文件超过 512 MB 安全上限。")
                digest.update(chunk)
                output.write(chunk)
        actual_sha256 = digest.hexdigest()
        if actual_sha256 != expected_sha256:
            raise RuntimeError(f"下载文件的 SHA-256 与官网清单不一致（实际 {actual_sha256}）。")
        if package_type == "standalone_zip":
            resolved_exe = _extract_standalone_zip(target, target_dir, relative_launcher, release_version)
        else:
            with target.open("rb") as handle:
                if handle.read(2) != b"MZ":
                    raise RuntimeError("下载内容不是有效的 Windows EXE。")
            resolved_exe = target.resolve()
        if candidate_version(resolved_exe) != version_tuple(release_version):
            raise RuntimeError("下载的 EXE 版本与云端最新版本不一致，已拒绝调用。")
    except Exception:
        target.unlink(missing_ok=True)
        raise
    return {
        "ok": True,
        "code": "downloaded",
        "exe_path": str(resolved_exe),
        "version": release_version,
        "sha256": expected_sha256,
        "bytes": total,
        "manifest_url": OFFICIAL_MANIFEST_URL,
    }


def require_exe(args) -> tuple[Path | None, dict | None]:
    exe, old = find_exe(args.exe)
    if exe:
        return exe, None
    if old:
        minimum = ".".join(str(item) for item in MIN_EXE_VERSION)
        return None, {"ok": False, "code": "software_upgrade_required", "message": f"已找到 MML解析工具，但版本低于 {minimum}，必须经用户授权重新下载官网最新 EXE，校验通过后才能调用。", "found_exe": str(old), "download_page_url": OFFICIAL_DOWNLOAD_PAGE, "manifest_url": OFFICIAL_MANIFEST_URL}
    return None, {"ok": False, "code": "software_not_found", "message": "本机未找到 MML解析工具 EXE。必须经用户授权从官网下载最新 EXE，校验通过后才能调用。", "download_page_url": OFFICIAL_DOWNLOAD_PAGE, "manifest_url": OFFICIAL_MANIFEST_URL}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=f"{SKILL_NAME} Skill v{SKILL_VERSION}")
    parser.add_argument("--exe", default="", help="显式指定 MML解析工具 EXE")
    sub = parser.add_subparsers(dest="command", required=True)
    status = sub.add_parser("status", help="查找 EXE 并查询设备身份及额度")
    status.add_argument("--model-name", default="unknown-ai-model")
    download = sub.add_parser("download", help="从知一官网下载软件")
    download.add_argument("--consent", action="store_true", help="确认用户已明确授权下载")
    download.add_argument("--destination", default="")
    trial = sub.add_parser("start-trial", help="经用户授权后登记试用")
    trial.add_argument("--consent", action="store_true", help="确认用户已明确授权进入试用期")
    trial.add_argument("--model-name", default="unknown-ai-model")
    preflight = sub.add_parser("preflight", help="可信计算文件规模并进行 AI 权益预检")
    preflight.add_argument("--model-name", required=True)
    preflight.add_argument("--input", action="append", required=True, dest="inputs")
    preflight.add_argument("--allow-overage", action="store_true", help="仅在用户明确同意并持有报价时兑换权益")
    preflight.add_argument("--offer-id", default="")
    preflight.add_argument("--request-id", default="")
    scan = sub.add_parser("scan", help="权益通过后扫描可导出报表")
    scan.add_argument("--model-name", required=True)
    scan.add_argument("--input", action="append", required=True, dest="inputs")
    scan.add_argument("--request-id", default="")
    parse = sub.add_parser("parse", help="安全解析 MML TXT 文件")
    parse.add_argument("--model-name", required=True)
    parse.add_argument("--input", action="append", required=True, dest="inputs")
    parse.add_argument("--output", required=True)
    parse.add_argument("--output-format", choices=("csv", "xlsx", "db"), default="csv")
    parse.add_argument("--split-mode", choices=("partial", "all", "none"), default="partial")
    parse.add_argument("--compress", action="store_true")
    parse.add_argument("--report", action="append", default=[])
    parse.add_argument("--all-reports", action="store_true", help="仅当用户明确选择全部报表时使用")
    parse.add_argument("--instruction", default="解析并导出 MML TXT 数据")
    parse.add_argument("--allow-overage", action="store_true", help="确认用户已授权消耗激活期追加额度")
    parse.add_argument("--request-id", default="")
    parse.add_argument("--timeout", type=int, default=3600)
    return parser


def main() -> int:
    args = build_parser().parse_args()
    if args.command == "download":
        if not args.consent:
            return emit({"ok": False, "code": "download_consent_required", "message": "下载前必须获得用户明确授权。"}, 2)
        try:
            result = download_official(args.destination)
            return emit(result, 0 if result.get("ok") else 2)
        except Exception as exc:
            return emit({"ok": False, "code": "download_failed", "message": str(exc), "download_page_url": OFFICIAL_DOWNLOAD_PAGE, "manifest_url": OFFICIAL_MANIFEST_URL}, 2)

    exe, error = require_exe(args)
    if error:
        return emit(error, 2)
    assert exe is not None
    try:
        if args.command == "status":
            request = {"action": "status", "ai_model_name": args.model_name, "skill_id": SKILL_ID, "skill_version": SKILL_VERSION}
            result = invoke_exe(exe, request, 60)
        elif args.command == "start-trial":
            if not args.consent:
                return emit({"ok": False, "code": "trial_consent_required", "message": "进入试用期前必须获得用户明确授权。"}, 2)
            request = {"action": "start_trial", "user_consent": True, "ai_model_name": args.model_name, "skill_id": SKILL_ID, "skill_version": SKILL_VERSION}
            result = invoke_exe(exe, request, 60)
        elif args.command in {"preflight", "scan"}:
            request_id = args.request_id or uuid.uuid4().hex
            request = {
                "action": args.command,
                "input_files": args.inputs,
                "ai_model_name": args.model_name,
                "request_id": request_id,
                "skill_version": SKILL_VERSION,
                "skill_id": SKILL_ID,
            }
            if args.command == "preflight":
                request["overage_consent"] = bool(args.allow_overage)
                request["offer_id"] = args.offer_id
            result = invoke_exe(exe, request, 120)
            result.setdefault("request_id", request_id)
        else:
            if bool(args.report) == bool(args.all_reports):
                return emit({
                    "ok": False,
                    "code": "report_selection_required",
                    "message": "必须提供一个或多个 --report，或在用户明确选择全部时提供 --all-reports。",
                    "next_action": "ask_user_for_report_selection",
                }, 2)
            request_id = args.request_id or uuid.uuid4().hex
            request = {
                "action": "parse", "input_files": args.inputs, "output_path": args.output,
                "output_format": args.output_format, "split_mode": args.split_mode,
                "compress_output": bool(args.compress), "reports": args.report,
                "select_all_reports": bool(args.all_reports),
                "ai_model_name": args.model_name, "instruction": args.instruction,
                "overage_consent": bool(args.allow_overage), "request_id": request_id,
                "skill_version": SKILL_VERSION,
                "skill_id": SKILL_ID,
            }
            result = invoke_exe(exe, request, max(60, args.timeout))
            result.setdefault("request_id", request_id)
        return emit(result, 0 if result.get("ok") else 2)
    except subprocess.TimeoutExpired:
        return emit({"ok": False, "code": "exe_timeout", "message": "MML解析工具执行超时。", "exe_path": str(exe)}, 2)
    except Exception as exc:
        return emit({"ok": False, "code": "exe_invoke_failed", "message": str(exc), "exe_path": str(exe)}, 2)


if __name__ == "__main__":
    raise SystemExit(main())
