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


from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import sqlite3
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional, Tuple

import requests

try:
    # 可选：用于从 .env 文件加载环境变量（UPLOAD_JSON_SUBDIRS / UPLOAD_JSON_SUBDIR 等）
    from dotenv import load_dotenv  # type: ignore
except Exception:  # pragma: no cover
    load_dotenv = None  # type: ignore

DEFAULT_ENDPOINT = (
    os.environ.get("UPLOAD_ENDPOINT")
    or "http://127.0.0.1:8317/v0/management/auth-files"
).strip()
DEFAULT_DB_NAME = (os.environ.get("UPLOAD_DB") or "upload_state.sqlite3").strip()

# 环境变量：指定 JSON 所在子目录（相对 --dir / 当前目录），可配置多个
# 例：
#   Windows:  set UPLOAD_JSON_SUBDIRS=data;data2
#   bash:     export UPLOAD_JSON_SUBDIRS="data:data2"   (注意：在 Windows 上 os.pathsep 是 ;)
# 同时也支持逗号分隔：data,data2
ENV_JSON_SUBDIRS = "UPLOAD_JSON_SUBDIRS"
ENV_JSON_SUBDIR = "UPLOAD_JSON_SUBDIR"  # 单个子目录（兼容/简化用）

# 多个上传站点（接口地址）列表，建议配置在 .env
# 例：UPLOAD_ENDPOINTS=http://a/v0/management/auth-files;http://b/v0/management/auth-files
ENV_ENDPOINTS = "UPLOAD_ENDPOINTS"
ENV_TOKEN = "UPLOAD_TOKEN"


@dataclass(frozen=True)
class UploadResult:
    ok: bool
    status_code: Optional[int]
    text: str


def setup_logger() -> logging.Logger:
    logger = logging.getLogger("uploader")
    logger.setLevel(logging.INFO)

    class _CnFormatter(logging.Formatter):
        LEVEL_MAP = {
            "DEBUG": "调试",
            "INFO": "信息",
            "WARNING": "警告",
            "ERROR": "错误",
            "CRITICAL": "严重",
        }

        def format(self, record: logging.LogRecord) -> str:
            record.levelname_cn = self.LEVEL_MAP.get(record.levelname, record.levelname)
            return super().format(record)

    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(logging.INFO)
    formatter = _CnFormatter(
        fmt="[%(asctime)s] [%(levelname_cn)s] %(message)s",
        datefmt="%H:%M:%S",
    )
    handler.setFormatter(formatter)

    # 防止重复添加 handler（例如某些 IDE 运行方式）
    if not logger.handlers:
        logger.addHandler(handler)

    return logger


def init_db(db_path: Path) -> sqlite3.Connection:
    conn = sqlite3.connect(str(db_path))

    # v2: 按 endpoint + file_path 去重
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS uploads_v2 (
          endpoint TEXT NOT NULL,
          file_path TEXT NOT NULL,
          file_name TEXT NOT NULL,
          sha256 TEXT NOT NULL,
          size_bytes INTEGER NOT NULL,
          mtime_ns INTEGER NOT NULL,
          status TEXT NOT NULL,          -- success | skipped | failed
          http_status INTEGER,
          response_text TEXT,
          updated_at INTEGER NOT NULL,
          PRIMARY KEY (endpoint, file_path)
        )
        """
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_uploads_v2_status ON uploads_v2(status)"
    )

    # 兼容迁移：如果旧表 uploads 存在且 v2 没数据，则把旧数据迁移到 v2（按 DEFAULT_ENDPOINT）
    try:
        old_exists = (
            conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='uploads'"
            ).fetchone()
            is not None
        )
        if old_exists:
            v2_count = conn.execute("SELECT COUNT(1) FROM uploads_v2").fetchone()[0]
            if v2_count == 0:
                conn.execute(
                    """
                    INSERT OR IGNORE INTO uploads_v2(
                      endpoint, file_path, file_name, sha256, size_bytes, mtime_ns,
                      status, http_status, response_text, updated_at
                    )
                    SELECT
                      ?, file_path, file_name, sha256, size_bytes, mtime_ns,
                      status, http_status, response_text, updated_at
                    FROM uploads
                    """,
                    (DEFAULT_ENDPOINT,),
                )
    except Exception:
        # 迁移失败不影响主流程
        pass

    conn.commit()
    return conn


def sha256_of_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()


def get_file_fingerprint(path: Path) -> Tuple[int, int, str]:
    st = path.stat()
    size = int(st.st_size)
    mtime_ns = int(getattr(st, "st_mtime_ns", int(st.st_mtime * 1e9)))
    digest = sha256_of_file(path)
    return size, mtime_ns, digest


def is_file_stable(path: Path, stable_seconds: float = 1.5) -> bool:
    """避免上传正在写入中的文件：短暂等待并检查 size/mtime 是否变化。"""
    try:
        st1 = path.stat()
    except FileNotFoundError:
        return False
    time.sleep(stable_seconds)
    try:
        st2 = path.stat()
    except FileNotFoundError:
        return False

    return (st1.st_size == st2.st_size) and (st1.st_mtime_ns == st2.st_mtime_ns)


def iter_json_files(folder: Path) -> Iterable[Path]:
    # 递归扫描：如果你只想扫描当前目录（不含子目录），把 rglob 改成 glob
    for p in folder.rglob("*.json"):
        if p.is_file():
            yield p


def parse_list_from_env(var_name: str) -> list[str]:
    raw = os.environ.get(var_name, "").strip()
    if not raw:
        return []

    parts: list[str] = []
    # 同时兼容 "a;b" / "a:b" / "a,b"
    for token in raw.replace(",", os.pathsep).split(os.pathsep):
        token = token.strip()
        if token:
            parts.append(token)

    # 去重并保持顺序
    seen: set[str] = set()
    uniq: list[str] = []
    for p in parts:
        if p not in seen:
            seen.add(p)
            uniq.append(p)
    return uniq


def parse_endpoints_from_env(fallback_endpoint: str) -> list[str]:
    """读取需要上传的 endpoint 列表。

    - 优先使用环境变量 UPLOAD_ENDPOINTS
    - endpoints 只支持用 ; 或 , 分隔（不要用 ':'），因为 URL 自身包含 'http://', 端口等 ':'
    - 如果未配置，则回退到命令行 --endpoint
    """

    raw = os.environ.get(ENV_ENDPOINTS, "").strip()
    if raw:
        parts = [p.strip() for p in raw.replace(",", ";").split(";")]
        endpoints = [p for p in parts if p]
        if endpoints:
            return endpoints

    return [fallback_endpoint]


def normalize_token(raw: str) -> str:
    t = (raw or "").strip()
    if not t:
        return ""
    if t.lower().startswith("bearer "):
        return t.split(" ", 1)[1].strip()
    return t


def parse_subdirs_from_env(base_dir: Path) -> list[Path]:
    """从环境变量读取 JSON 子目录列表。

    支持：
    - UPLOAD_JSON_SUBDIRS：多个子目录，用 os.pathsep 分隔（Windows 是 ;，Linux/mac 是 :）
    - 也支持逗号分隔
    - UPLOAD_JSON_SUBDIR：单个子目录（兼容/简化）

    返回：绝对路径列表（不存在的会被过滤掉）

    说明：
    - 如果传入的是绝对路径，则直接使用
    - 如果是相对路径，则拼接 base_dir
    """

    raw = os.environ.get(ENV_JSON_SUBDIRS, "").strip()
    if not raw:
        raw = os.environ.get(ENV_JSON_SUBDIR, "").strip()

    if not raw:
        return []

    parts = (
        parse_list_from_env(ENV_JSON_SUBDIRS)
        if os.environ.get(ENV_JSON_SUBDIRS)
        else []
    )
    if not parts:
        parts = [raw]

    dirs: list[Path] = []
    for sub in parts:
        candidate = Path(sub).expanduser()
        if candidate.is_absolute():
            p = candidate.resolve()
        else:
            p = (base_dir / sub).expanduser().resolve()

        if p.exists() and p.is_dir():
            dirs.append(p)

    # 去重并保持顺序
    seen: set[str] = set()
    uniq: list[Path] = []
    for d in dirs:
        s = str(d)
        if s not in seen:
            seen.add(s)
            uniq.append(d)

    return uniq


def db_has_success_or_skipped(
    conn: sqlite3.Connection, endpoint: str, file_path: str
) -> bool:
    row = conn.execute(
        "SELECT status FROM uploads_v2 WHERE endpoint = ? AND file_path = ?",
        (endpoint, file_path),
    ).fetchone()
    if not row:
        return False
    return row[0] in ("success", "skipped")


def db_upsert(
    conn: sqlite3.Connection,
    *,
    endpoint: str,
    file_path: str,
    file_name: str,
    sha256: str,
    size_bytes: int,
    mtime_ns: int,
    status: str,
    http_status: Optional[int],
    response_text: str,
) -> None:
    now = int(time.time())
    conn.execute(
        """
        INSERT INTO uploads_v2(
          endpoint, file_path, file_name, sha256, size_bytes, mtime_ns,
          status, http_status, response_text, updated_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(endpoint, file_path) DO UPDATE SET
          file_name=excluded.file_name,
          sha256=excluded.sha256,
          size_bytes=excluded.size_bytes,
          mtime_ns=excluded.mtime_ns,
          status=excluded.status,
          http_status=excluded.http_status,
          response_text=excluded.response_text,
          updated_at=excluded.updated_at
        """,
        (
            endpoint,
            file_path,
            file_name,
            sha256,
            size_bytes,
            mtime_ns,
            status,
            http_status,
            response_text[:4000],
            now,
        ),
    )
    conn.commit()


def should_treat_as_duplicate(resp: requests.Response) -> bool:
    """无法提前拿到服务器已有列表时，依赖响应判断“已存在”。

    你可以按实际接口返回格式在这里加强判断逻辑。
    """
    if resp.status_code in (409, 208):
        return True

    text = (resp.text or "").lower()
    # 常见提示关键字（按需增删）
    keywords = [
        "already",
        "exists",
        "duplicate",
        "重复",
        "已存在",
        "已上传",
    ]
    return any(k in text for k in keywords)


def upload_file(
    *,
    endpoint: str,
    token: str,
    path: Path,
    timeout_s: int,
    verify_tls: bool,
    extra_headers_json: Optional[str] = None,
) -> UploadResult:
    headers = {
        "Accept": "application/json, text/plain, */*",
        "Authorization": f"Bearer {token}",
    }

    if extra_headers_json:
        try:
            extra = json.loads(extra_headers_json)
            if not isinstance(extra, dict):
                raise ValueError("extra headers must be a JSON object")
            for k, v in extra.items():
                headers[str(k)] = str(v)
        except Exception as e:
            return UploadResult(False, None, f"Invalid --extra-headers-json: {e}")

    # requests 会自动生成 multipart boundary；不要手工设置 Content-Type（否则 boundary 不匹配）
    with path.open("rb") as f:
        files = {
            "file": (path.name, f, "application/json"),
        }
        try:
            resp = requests.post(
                endpoint,
                headers=headers,
                files=files,
                timeout=timeout_s,
                verify=verify_tls,
            )
        except Exception as e:
            return UploadResult(False, None, f"Request error: {e}")

    # 只要 2xx 就认为成功
    if 200 <= resp.status_code < 300:
        return UploadResult(True, resp.status_code, resp.text or "")

    # 如果服务器返回“已存在”，可当作跳过（ok=True 但由上层标记为 skipped）
    if should_treat_as_duplicate(resp):
        return UploadResult(True, resp.status_code, resp.text or "")

    return UploadResult(False, resp.status_code, resp.text or "")


def main() -> int:
    parser = argparse.ArgumentParser(description="Upload local JSON files to target site")
    parser.add_argument("--dir", default=".", help="工作目录（默认当前目录）")
    parser.add_argument("--endpoint", default=DEFAULT_ENDPOINT, help="上传接口 URL")
    parser.add_argument(
        "--token",
        default=None,
        help=f"上传鉴权 token（也可用环境变量 {ENV_TOKEN} 配置）",
    )
    parser.add_argument(
        "--interval",
        type=int,
        default=120,
        help="扫描间隔秒数（默认 120 秒）",
    )
    parser.add_argument(
        "--db",
        default=DEFAULT_DB_NAME,
        help="本地状态数据库文件名（默认 upload_state.sqlite3）",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=60,
        help="单次上传超时秒数（默认 60）",
    )
    parser.add_argument(
        "--verify-tls",
        action="store_true",
        help="开启 TLS 证书校验（仅 https 有意义；默认关闭）",
    )
    parser.add_argument(
        "--once",
        action="store_true",
        help="只扫描并上传一轮后退出（默认持续运行）",
    )
    parser.add_argument(
        "--extra-headers-json",
        default=None,
        help='附加请求头(JSON对象)，如 {"Origin":"...","Referer":"..."}',
    )

    args = parser.parse_args()

    logger = setup_logger()

    # 如果安装了 python-dotenv，则自动加载当前目录下的 .env
    if load_dotenv is not None:
        load_dotenv()

    base_dir = Path(args.dir).expanduser().resolve()
    if not base_dir.exists() or not base_dir.is_dir():
        logger.error(f"目录不存在或不是文件夹: {base_dir}")
        return 2

    # 从环境变量读取子目录；如果没配，则默认扫描 base_dir 本身
    subdirs = parse_subdirs_from_env(base_dir)
    scan_dirs = subdirs if subdirs else [base_dir]

    db_path = Path(args.db).expanduser().resolve()
    conn = init_db(db_path)

    # token：优先命令行，其次环境变量/.env
    token = normalize_token(args.token or os.environ.get(ENV_TOKEN, ""))
    if not token:
        logger.error(f"未配置 token：请使用 --token 或在 .env 设置 {ENV_TOKEN}")
        return 2

    endpoints = parse_endpoints_from_env(args.endpoint)

    logger.info(f"工作目录: {base_dir}")
    if subdirs:
        logger.info(
            f"扫描子目录: {', '.join(str(d) for d in scan_dirs)}（来自环境变量 {ENV_JSON_SUBDIRS}/{ENV_JSON_SUBDIR}）"
        )
    else:
        logger.info("未设置子目录环境变量：将扫描工作目录本身")
    logger.info(f"上传站点数: {len(endpoints)}")
    for i, ep in enumerate(endpoints, 1):
        logger.info(f"站点[{i}]: {ep}")
    logger.info(f"状态库: {db_path}")
    logger.info(f"扫描间隔: {args.interval} 秒")

    while True:
        any_new = False

        # 汇总所有扫描目录的文件
        all_files: list[Path] = []
        for d in scan_dirs:
            all_files.extend(iter_json_files(d))

        # 为了可控，按 mtime 升序上传（更接近“按生成时间”）
        files = sorted(all_files, key=lambda p: p.stat().st_mtime_ns)

        for path in files:
            file_path_str = str(path)

            for endpoint in endpoints:
                if db_has_success_or_skipped(conn, endpoint, file_path_str):
                    continue

                if not is_file_stable(path):
                    logger.info(f"跳过（正在写入）: {path.name}")
                    break

                any_new = True

                try:
                    size_bytes, mtime_ns, digest = get_file_fingerprint(path)
                except Exception as e:
                    logger.error(f"失败（读取/计算哈希）: {path.name}: {e}")
                    db_upsert(
                        conn,
                        endpoint=endpoint,
                        file_path=file_path_str,
                        file_name=path.name,
                        sha256="",
                        size_bytes=0,
                        mtime_ns=0,
                        status="failed",
                        http_status=None,
                        response_text=str(e),
                    )
                    continue
                result = upload_file(
                    endpoint=endpoint,
                    token=token,
                    path=path,
                    timeout_s=args.timeout,
                    verify_tls=args.verify_tls,
                    extra_headers_json=args.extra_headers_json,
                )

                if result.ok:
                    status = "success"
                    if result.status_code is not None and not (
                        200 <= result.status_code < 300
                    ):
                        status = "skipped"

                    if status == "success":
                        logger.info(
                            f"成功: {path.name} -> {endpoint} | HTTP={result.status_code}"
                        )
                    else:
                        logger.info(
                            f"跳过（已存在）: {path.name} -> {endpoint} | HTTP={result.status_code}"
                        )

                    db_upsert(
                        conn,
                        endpoint=endpoint,
                        file_path=file_path_str,
                        file_name=path.name,
                        sha256=digest,
                        size_bytes=size_bytes,
                        mtime_ns=mtime_ns,
                        status=status,
                        http_status=result.status_code,
                        response_text=result.text,
                    )
                else:
                    logger.error(
                        f"失败: {path.name} -> {endpoint} | HTTP={result.status_code} | {result.text[:300]}"
                    )
                    db_upsert(
                        conn,
                        endpoint=endpoint,
                        file_path=file_path_str,
                        file_name=path.name,
                        sha256=digest,
                        size_bytes=size_bytes,
                        mtime_ns=mtime_ns,
                        status="failed",
                        http_status=result.status_code,
                        response_text=result.text,
                    )

        if args.once:
            logger.info("已完成（单次模式）。")
            return 0

        wait_s = max(5, int(args.interval))

        if not any_new:
            # 倒计时读秒（同一行刷新），避免刷屏
            for remaining in range(wait_s, 0, -1):
                sys.stdout.write(
                    f"\r[{time.strftime('%H:%M:%S')}] [信息] 暂无新文件，等待中... {remaining} 秒"
                )
                sys.stdout.flush()
                time.sleep(1)
            sys.stdout.write("\n")
            sys.stdout.flush()
        else:
            time.sleep(wait_s)


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