#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Incrementally sync PCAPdroid files from a Windows host directory to a local directory."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
import time
import traceback
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, TextIO, Tuple


DEFAULT_SOURCE_ROOT = os.path.join(
    os.path.expanduser("~"), "Documents", "MuMu共享文件夹", "Download", "PCAPdroid"
)
DEFAULT_LOCAL_ROOT = Path("./Flowdata")
MANIFEST_FILE_NAME = ".pcapdroid_sync_manifest.json"
RUN_LOG_FILE_NAME = "pull_pcapdroid_files.log"


@dataclass
class ChecksumInfo:
    algorithm: str
    digest: str


@dataclass
class SourceFileInfo:
    path: Path
    size: int
    mtime: int


def emit_progress(enabled: bool, step: str, message: str, **extra: Any) -> None:
    if not enabled:
        return
    print(json.dumps({"progress": True, "step": step, "message": message, **extra}, ensure_ascii=True), flush=True)


def emit_result(enabled: bool, ok: bool, message: str, **extra: Any) -> int:
    payload = {"ok": ok, "message": message, **extra}
    if enabled:
        print(json.dumps(payload, ensure_ascii=True), flush=True)
    return 0 if ok else 1


class TeeStream:
    def __init__(self, *streams: TextIO) -> None:
        self.streams = streams
        self.encoding = getattr(streams[0], "encoding", "utf-8") if streams else "utf-8"

    def write(self, data: str) -> int:
        for stream in self.streams:
            stream.write(data)
        return len(data)

    def flush(self) -> None:
        for stream in self.streams:
            stream.flush()

    def isatty(self) -> bool:
        return False


@contextmanager
def mirror_output_to_local_log(local_root: Path) -> Iterator[Path]:
    local_root.mkdir(parents=True, exist_ok=True)
    log_path = local_root / RUN_LOG_FILE_NAME
    original_stdout = sys.stdout
    original_stderr = sys.stderr
    with log_path.open("w", encoding="utf-8", buffering=1) as log_file:
        sys.stdout = TeeStream(original_stdout, log_file)
        sys.stderr = TeeStream(original_stderr, log_file)
        print(f"Log file: {log_path}", flush=True)
        print(f"Run started: {time.strftime('%Y-%m-%d %H:%M:%S')}", flush=True)
        try:
            yield log_path
        finally:
            print(f"Run finished: {time.strftime('%Y-%m-%d %H:%M:%S')}", flush=True)
            sys.stdout = original_stdout
            sys.stderr = original_stderr


def resolve_windows_host_path(path_text: str) -> Path:
    raw_path = str(path_text or "").strip().strip('"')
    if not raw_path:
        raise ValueError("源目录不能为空")
    return Path(raw_path).expanduser()


def scan_source_files(source_root: Path) -> List[SourceFileInfo]:
    files: List[SourceFileInfo] = []
    for path in source_root.rglob("*"):
        if not path.is_file():
            continue
        stat_result = path.stat()
        files.append(
            SourceFileInfo(
                path=path,
                size=int(stat_result.st_size),
                mtime=int(stat_result.st_mtime),
            )
        )
    files.sort(key=lambda item: str(item.path))
    return files


def get_file_checksum(path: Path, algorithm: str) -> str:
    hasher = hashlib.new(algorithm)
    with path.open("rb") as file_obj:
        for chunk in iter(lambda: file_obj.read(1024 * 1024), b""):
            hasher.update(chunk)
    return hasher.hexdigest().lower()


def get_source_checksum(source_path: Path) -> ChecksumInfo:
    algorithm = "sha256"
    return ChecksumInfo(algorithm=algorithm, digest=get_file_checksum(source_path, algorithm))


def load_manifest(manifest_path: Path) -> Dict[str, Dict[str, int]]:
    if not manifest_path.exists():
        return {}

    try:
        return json.loads(manifest_path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}


def save_manifest(manifest_path: Path, manifest: Dict[str, Dict[str, int]]) -> None:
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = manifest_path.with_suffix(f"{manifest_path.suffix}.tmp")
    temp_path.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    os.replace(temp_path, manifest_path)


def local_matches_source(local_path: Path, source_info: SourceFileInfo) -> bool:
    if not local_path.exists():
        return False

    local_stat = local_path.stat()
    if local_stat.st_size != source_info.size:
        return False

    return int(local_stat.st_mtime) == source_info.mtime


def needs_sync(
    source_info: SourceFileInfo,
    local_path: Path,
    manifest_entry: Optional[Dict[str, int]],
) -> Tuple[bool, str]:
    if not local_path.exists():
        return True, "本地不存在"

    if local_matches_source(local_path, source_info):
        return False, "已同步"

    if manifest_entry is not None:
        manifest_size = manifest_entry.get("size")
        manifest_mtime = manifest_entry.get("mtime")
        if manifest_size == source_info.size and manifest_mtime == source_info.mtime:
            return True, "本地文件与同步记录不一致"

    source_checksum = get_source_checksum(source_info.path)
    local_checksum = get_file_checksum(local_path, source_checksum.algorithm)
    if local_checksum != source_checksum.digest:
        return True, "校验和不同"

    return False, "内容相同"


def copy_file(source_path: Path, local_path: Path) -> None:
    local_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source_path, local_path)


def sync_pcapdroid_files(
    source_root: str,
    local_root: Path,
    *,
    verbose: bool = False,
) -> int:
    source_root_path = resolve_windows_host_path(source_root)
    local_root.mkdir(parents=True, exist_ok=True)
    manifest_path = local_root / MANIFEST_FILE_NAME
    previous_manifest = load_manifest(manifest_path)
    next_manifest: Dict[str, Dict[str, int]] = {}

    if not source_root_path.is_dir():
        print(f"源目录不存在: {source_root_path}", file=sys.stderr)
        return 1

    source_files = scan_source_files(source_root_path)

    total = len(source_files)
    copied = 0
    skipped = 0
    failed = 0

    print(f"源文件总数: {total}")
    print(f"源目录: {source_root_path}")
    print(f"本地目标目录: {local_root}")
    print(f"清单文件: {manifest_path}")

    for index, source_file in enumerate(source_files, start=1):
        try:
            relative_path = source_file.path.relative_to(source_root_path)
            local_path = local_root.joinpath(*relative_path.parts)
            manifest_key = relative_path.as_posix()
            should_copy, reason = needs_sync(
                source_file,
                local_path,
                previous_manifest.get(manifest_key),
            )

            if not should_copy:
                skipped += 1
                if verbose:
                    print(f"[{index}/{total}] 跳过: {source_file.path} ({reason})")
                next_manifest[manifest_key] = {
                    "size": source_file.size,
                    "mtime": source_file.mtime,
                }
                continue

            copy_file(source_file.path, local_path)
            copied += 1
            print(f"[{index}/{total}] 已拷贝: {source_file.path} -> {local_path} ({reason})")
            next_manifest[manifest_key] = {
                "size": source_file.size,
                "mtime": source_file.mtime,
            }
        except Exception as exc:  # noqa: BLE001
            failed += 1
            print(f"[{index}/{total}] 失败: {source_file.path} ({exc})", file=sys.stderr)

        if not verbose and index % 200 == 0:
            print(
                f"进度: {index}/{total}，新增/更新 {copied}，跳过 {skipped}，失败 {failed}"
            )

    if failed == 0:
        save_manifest(manifest_path, next_manifest)
    else:
        try:
            save_manifest(manifest_path, next_manifest)
        except OSError:
            pass

    print()
    print("同步完成")
    print(f"新增/更新: {copied}")
    print(f"跳过: {skipped}")
    print(f"失败: {failed}")

    return 0 if failed == 0 else 2


def run_remote_pull(args: argparse.Namespace) -> int:
    source_root = str(resolve_windows_host_path(args.source_root))

    emit_progress(
        args.json_progress,
        "pull_pcap_files",
        "Starting incremental PCAPdroid copy from Windows host directory",
        status="running",
        source_root=source_root,
    )
    exit_code = sync_pcapdroid_files(
        source_root=source_root,
        local_root=Path(args.local_root),
        verbose=args.verbose,
    )
    if exit_code != 0:
        return emit_result(
            args.json_progress,
            False,
            f"pull_pcap_files failed: exit code {exit_code}",
            source_root=source_root,
            local_root=str(Path(args.local_root)),
            exit_code=exit_code,
        )
    return emit_result(
        args.json_progress,
        True,
        "pull_pcap_files completed",
        source_root=source_root,
        local_root=str(Path(args.local_root)),
        exit_code=exit_code,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="使用 Windows 主机共享目录增量同步 PCAPdroid 文件到本地目录。"
    )
    parser.add_argument(
        "--source-root",
        default=DEFAULT_SOURCE_ROOT,
        help=f"PCAPdroid 源目录，默认: {DEFAULT_SOURCE_ROOT}",
    )
    parser.add_argument(
        "--local-root",
        default=str(DEFAULT_LOCAL_ROOT),
        help=f"本地目录，默认: {DEFAULT_LOCAL_ROOT}",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="打印每个已跳过文件。",
    )
    parser.add_argument(
        "--json-progress",
        action="store_true",
        help="输出 JSON progress/result，供中控解析。",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    local_root = Path(args.local_root)
    with mirror_output_to_local_log(local_root):
        try:
            if args.json_progress:
                return run_remote_pull(args)
            return sync_pcapdroid_files(
                source_root=args.source_root,
                local_root=local_root,
                verbose=args.verbose,
            )
        except Exception:  # noqa: BLE001
            traceback.print_exc()
            return 1


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