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

import argparse
import os
import shlex
import socket
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterable, Optional

import pymysql
import yaml


APP_DIR = Path(__file__).resolve().parents[1]
if str(APP_DIR) not in sys.path:
    sys.path.insert(0, str(APP_DIR))

from main import get_action_timeout  # noqa: E402
from src.core.config import load_config  # noqa: E402


STATUS_RUNNING = 0
STATUS_SUCCESS = 1
STATUS_FAILED = 2
STATUS_SKIPPED = 3
OUTPUT_LIMIT = 5000


@dataclass
class CrawlerTask:
    task_key: str
    task_name: str
    action: str
    cron_expression: str
    active: bool = True
    timeout: Optional[int] = None

    @classmethod
    def from_dict(cls, row: dict) -> "CrawlerTask":
        action = str(row.get("action") or row.get("params") or "").strip()
        return cls(
            task_key=str(row.get("key") or row.get("task_key") or action).strip(),
            task_name=str(row.get("name") or row.get("task_name") or action).strip(),
            action=action,
            cron_expression=str(row.get("cron") or row.get("cron_expression") or "* * * * *").strip(),
            active=bool(row.get("active", True)),
            timeout=int(row["timeout"]) if row.get("timeout") not in (None, "") else None,
        )

    def resolved_timeout(self) -> int:
        return int(self.timeout or get_action_timeout(self.action))


@dataclass
class CommandResult:
    exit_code: int
    output: str
    elapsed: float
    timed_out: bool = False


@dataclass
class TaskRunResult:
    log_id: int
    status: int
    output: str
    elapsed: float
    error_message: str = ""


class TaskLock:
    def __init__(self, action: str, lock_dir: Path):
        safe_action = "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in str(action or "all"))
        self.lock_dir = Path(lock_dir)
        self.path = self.lock_dir / f"{safe_action}.lock"
        self._fp = None
        self._locked = False

    def acquire(self) -> bool:
        self.lock_dir.mkdir(parents=True, exist_ok=True)
        self._fp = open(self.path, "a+b")
        try:
            self._fp.seek(0)
            if os.name == "nt":
                import msvcrt

                msvcrt.locking(self._fp.fileno(), msvcrt.LK_NBLCK, 1)
            else:
                import fcntl

                fcntl.flock(self._fp.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            self._locked = True
            self._fp.seek(0)
            self._fp.truncate()
            self._fp.write(f"pid={os.getpid()} started_at={int(time.time())}\n".encode("utf-8"))
            self._fp.flush()
            return True
        except (BlockingIOError, OSError):
            self.release()
            return False

    def release(self) -> None:
        if not self._fp:
            return
        try:
            if self._locked:
                if os.name == "nt":
                    import msvcrt

                    self._fp.seek(0)
                    msvcrt.locking(self._fp.fileno(), msvcrt.LK_UNLCK, 1)
                else:
                    import fcntl

                    fcntl.flock(self._fp.fileno(), fcntl.LOCK_UN)
        finally:
            self._locked = False
            self._fp.close()
            self._fp = None


class TaskLogStore:
    def __init__(self):
        cfg = load_config().database
        self.prefix = cfg.prefix
        self.conn = pymysql.connect(
            host=cfg.host,
            port=cfg.port,
            user=cfg.username,
            password=cfg.password,
            database=cfg.database,
            charset=cfg.charset,
            autocommit=True,
        )
        self.ensure_table()

    @property
    def table(self) -> str:
        return f"`{self.prefix}crawler_task_log`"

    def close(self) -> None:
        self.conn.close()

    def ensure_table(self) -> None:
        sql = f"""
        CREATE TABLE IF NOT EXISTS {self.table} (
          `id` BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
          `task_key` VARCHAR(100) NOT NULL DEFAULT '' COMMENT '任务标识',
          `task_name` VARCHAR(100) NOT NULL DEFAULT '' COMMENT '任务名称',
          `action` VARCHAR(64) NOT NULL DEFAULT '' COMMENT '爬虫动作',
          `cron_expression` VARCHAR(64) NOT NULL DEFAULT '' COMMENT 'cron规则',
          `source` VARCHAR(32) NOT NULL DEFAULT 'docker' COMMENT '来源',
          `container_id` VARCHAR(128) NOT NULL DEFAULT '' COMMENT '容器ID',
          `status` TINYINT UNSIGNED NOT NULL DEFAULT 0 COMMENT '0执行中 1成功 2失败 3跳过',
          `trigger_type` TINYINT UNSIGNED NOT NULL DEFAULT 1 COMMENT '1自动 2手动',
          `output` TEXT NULL COMMENT '执行输出',
          `error_message` TEXT NULL COMMENT '错误信息',
          `elapsed` DECIMAL(10,2) NOT NULL DEFAULT 0.00 COMMENT '耗时秒',
          `started_at` INT UNSIGNED NOT NULL DEFAULT 0 COMMENT '开始时间',
          `finished_at` INT UNSIGNED NOT NULL DEFAULT 0 COMMENT '结束时间',
          `create_time` INT UNSIGNED NOT NULL DEFAULT 0,
          `update_time` INT UNSIGNED NOT NULL DEFAULT 0,
          PRIMARY KEY (`id`),
          KEY `idx_action_status` (`action`, `status`),
          KEY `idx_create_time` (`create_time`),
          KEY `idx_task_key` (`task_key`)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='Docker爬虫任务执行日志'
        """
        with self.conn.cursor() as cur:
            cur.execute(sql)

    def create_running(self, task: CrawlerTask) -> int:
        now = int(time.time())
        with self.conn.cursor() as cur:
            cur.execute(
                f"""
                INSERT INTO {self.table}
                (`task_key`, `task_name`, `action`, `cron_expression`, `source`, `container_id`,
                 `status`, `trigger_type`, `output`, `error_message`, `elapsed`,
                 `started_at`, `finished_at`, `create_time`, `update_time`)
                VALUES (%s, %s, %s, %s, 'docker', %s, %s, 1, '', '', 0, %s, 0, %s, %s)
                """,
                (
                    task.task_key,
                    task.task_name,
                    task.action,
                    task.cron_expression,
                    socket.gethostname(),
                    STATUS_RUNNING,
                    now,
                    now,
                    now,
                ),
            )
            return int(cur.lastrowid)

    def update_progress(self, log_id: int, output: str, elapsed: float) -> None:
        with self.conn.cursor() as cur:
            cur.execute(
                f"UPDATE {self.table} SET output=%s, elapsed=%s, update_time=%s WHERE id=%s",
                (output[-OUTPUT_LIMIT:], round(elapsed, 2), int(time.time()), log_id),
            )

    def finish(self, log_id: int, status: int, output: str, elapsed: float, error_message: str = "") -> None:
        now = int(time.time())
        with self.conn.cursor() as cur:
            cur.execute(
                f"""
                UPDATE {self.table}
                SET status=%s, output=%s, error_message=%s, elapsed=%s, finished_at=%s, update_time=%s
                WHERE id=%s
                """,
                (status, output[-OUTPUT_LIMIT:], error_message[:2000], round(elapsed, 2), now, now, log_id),
            )


def load_tasks(path: Path) -> list[CrawlerTask]:
    raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
    rows = raw if isinstance(raw, list) else raw.get("tasks", [])
    return [CrawlerTask.from_dict(row) for row in rows]


def render_crontab(tasks: Iterable[CrawlerTask], python_bin: str = "/usr/local/bin/python", app_dir: str = "/app") -> str:
    lines = [
        "SHELL=/bin/bash",
        "BASH_ENV=/etc/sport-era-crawler/env.sh",
        "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
        "",
    ]
    for task in tasks:
        if not task.active:
            continue
        args = [
            python_bin,
            "scripts/docker_task_runner.py",
            "run",
            task.action,
            "--task-key",
            task.task_key,
            "--task-name",
            task.task_name,
            "--cron-expression",
            task.cron_expression,
        ]
        if task.timeout:
            args.extend(["--timeout", str(task.timeout)])
        cmd = " ".join(shlex.quote(str(arg)) for arg in args)
        lines.append(f"{task.cron_expression} root cd {shlex.quote(app_dir)} && {cmd}")
    lines.append("")
    return "\n".join(lines)


def run_main_action(action: str, timeout: int, output_callback: Callable[[str], None]) -> CommandResult:
    start = time.monotonic()
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["DQD_RUNNING_UNDER_TASK_RUNNER"] = "1"
    proc = subprocess.Popen(
        [sys.executable, "main.py", action],
        cwd=str(APP_DIR),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
        env=env,
    )
    output_parts: list[str] = []
    timed_out = False

    if os.name == "nt":
        try:
            stdout, _ = proc.communicate(timeout=timeout)
            stdout = stdout or ""
            output_parts.append(stdout)
            if stdout:
                output_callback(stdout)
        except subprocess.TimeoutExpired:
            timed_out = True
            proc.kill()
            stdout, _ = proc.communicate()
            stdout = stdout or ""
            output_parts.append(stdout)
            if stdout:
                output_callback(stdout)
    else:
        import selectors

        assert proc.stdout is not None
        selector = selectors.DefaultSelector()
        selector.register(proc.stdout, selectors.EVENT_READ)
        try:
            while proc.poll() is None:
                remaining = timeout - (time.monotonic() - start)
                if remaining <= 0:
                    timed_out = True
                    proc.kill()
                    break
                for key, _ in selector.select(timeout=min(1.0, remaining)):
                    line = key.fileobj.readline()
                    if line:
                        output_parts.append(line)
                        output_callback(line)
            rest = proc.stdout.read()
            if rest:
                output_parts.append(rest)
                output_callback(rest)
        finally:
            selector.close()

    elapsed = round(time.monotonic() - start, 2)
    return CommandResult(
        exit_code=proc.returncode if proc.returncode is not None else -1,
        output="".join(output_parts),
        elapsed=elapsed,
        timed_out=timed_out,
    )


def run_task(
    task: CrawlerTask,
    store=None,
    lock_dir: Optional[Path] = None,
    command_runner: Callable[[str, int, Callable[[str], None]], CommandResult] = run_main_action,
    flush_interval: float = 5.0,
) -> TaskRunResult:
    if store is None:
        store = TaskLogStore()
    if hasattr(store, "ensure_table"):
        store.ensure_table()

    log_id = int(store.create_running(task))
    start = time.monotonic()
    lock = TaskLock(task.action, Path(lock_dir or os.environ.get("DQD_LOCK_DIR") or APP_DIR / "data" / "locks"))
    output_parts: list[str] = []
    last_flush = 0.0

    def append_output(text: str) -> None:
        nonlocal last_flush
        output_parts.append(text)
        now = time.monotonic()
        elapsed = round(now - start, 2)
        if now - last_flush >= flush_interval:
            store.update_progress(log_id, "".join(output_parts), elapsed)
            last_flush = now

    if not lock.acquire():
        output = f"{task.action} 已有任务执行中，跳过本次"
        store.finish(log_id, STATUS_SKIPPED, output, 0, "")
        return TaskRunResult(log_id, STATUS_SKIPPED, output, 0, "")

    try:
        result = command_runner(task.action, task.resolved_timeout(), append_output)
        output = result.output or "".join(output_parts) or "执行完成"
        if result.timed_out:
            status = STATUS_FAILED
            error = f"{task.action} 执行超时，超过 {task.resolved_timeout()} 秒"
        elif result.exit_code == 0:
            status = STATUS_SUCCESS
            error = ""
        else:
            status = STATUS_FAILED
            error = f"{task.action} 执行失败，退出码 {result.exit_code}"
        store.finish(log_id, status, output, result.elapsed, error)
        return TaskRunResult(log_id, status, output, result.elapsed, error)
    except Exception as exc:
        elapsed = round(time.monotonic() - start, 2)
        output = "".join(output_parts) or str(exc)
        store.finish(log_id, STATUS_FAILED, output, elapsed, str(exc))
        return TaskRunResult(log_id, STATUS_FAILED, output, elapsed, str(exc))
    finally:
        lock.release()


def build_task_from_args(args) -> CrawlerTask:
    if args.tasks:
        for task in load_tasks(Path(args.tasks)):
            if task.action == args.action:
                if args.timeout:
                    task.timeout = args.timeout
                return task
    return CrawlerTask(
        task_key=args.task_key or args.action,
        task_name=args.task_name or args.action,
        action=args.action,
        cron_expression=args.cron_expression or "",
        active=True,
        timeout=args.timeout,
    )


def parse_args(argv: Optional[list[str]] = None):
    parser = argparse.ArgumentParser(description="Sport Era Docker crawler task runner")
    sub = parser.add_subparsers(dest="command", required=True)

    run = sub.add_parser("run")
    run.add_argument("action")
    run.add_argument("--task-key", default="")
    run.add_argument("--task-name", default="")
    run.add_argument("--cron-expression", default="")
    run.add_argument("--timeout", type=int, default=None)
    run.add_argument("--tasks", default="")

    render = sub.add_parser("render-cron")
    render.add_argument("--tasks", required=True)
    render.add_argument("--output", required=True)
    render.add_argument("--python-bin", default="/usr/local/bin/python")
    render.add_argument("--app-dir", default="/app")
    return parser.parse_args(argv)


def main(argv: Optional[list[str]] = None) -> int:
    args = parse_args(argv)
    if args.command == "render-cron":
        tasks = load_tasks(Path(args.tasks))
        content = render_crontab(tasks, python_bin=args.python_bin, app_dir=args.app_dir)
        output = Path(args.output)
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(content, encoding="utf-8")
        output.chmod(0o644)
        return 0

    if args.command == "run":
        store = TaskLogStore()
        try:
            result = run_task(build_task_from_args(args), store=store)
            return 0 if result.status in (STATUS_SUCCESS, STATUS_SKIPPED) else 1
        finally:
            store.close()

    return 1


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