from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Sequence

from common.runtime_paths import PROJECT_ROOT, app_data_root
from common.task_runtime import (
    MAX_CLOUD_CONCURRENT_TASKS,
    account_fingerprint,
    atomic_write_json,
    cleanup_stale_task_locks,
    make_task_id,
    read_status_files,
    safe_name,
    task_account_dir,
    task_data_root,
    task_lock_keys,
    task_lock_root,
)


class TaskStartError(RuntimeError):
    """任务启动失败，界面可以直接把错误提示给用户。"""


LOCK_START_GRACE_SECONDS = 60


@dataclass
class ManagedTask:
    task_id: str
    cloud: str
    account: str
    title: str
    task_dir: Path
    process: subprocess.Popen
    lock_paths: List[Path] = field(default_factory=list)


class TaskProcessManager:
    """壳体内的统一任务管理器，负责并发限制、独立目录和子进程启动。"""

    def __init__(self) -> None:
        self.tasks: Dict[str, ManagedTask] = {}
        task_data_root()
        task_lock_root()

    def _worker_command(self, spec_path: Path) -> List[str]:
        if getattr(sys, "frozen", False):
            return [sys.executable, "--task-worker", str(spec_path)]
        return [sys.executable, "-B", str(PROJECT_ROOT / "main.py"), "--task-worker", str(spec_path)]

    def _running_tasks(self) -> List[ManagedTask]:
        self.refresh()
        return [task for task in self.tasks.values() if task.process.poll() is None]

    def has_running_tasks(self) -> bool:
        return bool(self._running_tasks())

    def running_count_for_cloud(self, cloud: str) -> int:
        cloud_key = safe_name(cloud)
        return sum(1 for task in self._running_tasks() if safe_name(task.cloud) == cloud_key)

    def _pid_running(self, pid: int) -> bool:
        """判断锁文件里记录的子进程是否仍在运行，用于清理旧锁。"""
        if pid <= 0:
            return False
        if os.name == "nt":
            try:
                import ctypes

                kernel32 = ctypes.windll.kernel32
                handle = kernel32.OpenProcess(0x1000, False, int(pid))
                if not handle:
                    return False
                exit_code = ctypes.c_ulong()
                ok = kernel32.GetExitCodeProcess(handle, ctypes.byref(exit_code))
                kernel32.CloseHandle(handle)
                return bool(ok and exit_code.value == 259)
            except Exception:
                return False
        try:
            os.kill(int(pid), 0)
            return True
        except ProcessLookupError:
            return False
        except PermissionError:
            return True
        except Exception:
            return False

    def _lock_is_stale(self, path: Path) -> bool:
        """判断任务锁是否已失效，避免软件异常退出后永久挡住后续任务。"""
        try:
            data = json.loads(path.read_text(encoding="utf-8-sig"))
        except Exception:
            try:
                return (time.time() - path.stat().st_mtime) > LOCK_START_GRACE_SECONDS
            except OSError:
                return True
        pid = int(data.get("pid") or 0)
        if pid:
            return not self._pid_running(pid)
        try:
            created_at = float(data.get("created_at") or 0)
        except (TypeError, ValueError):
            created_at = 0
        return created_at > 0 and (time.time() - created_at) > LOCK_START_GRACE_SECONDS

    def _acquire_locks(self, keys: Sequence[str], task_id: str, task_dir: Path) -> List[Path]:
        cleanup_stale_task_locks()
        acquired: List[Path] = []
        for key in keys:
            path = task_lock_root() / f"{key}.lock"
            while True:
                try:
                    fd = os.open(str(path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                    break
                except FileExistsError as exc:
                    if self._lock_is_stale(path):
                        try:
                            path.unlink(missing_ok=True)
                            continue
                        except OSError:
                            pass
                    for item in acquired:
                        item.unlink(missing_ok=True)
                    detail = f"同账号同地区已有资源任务运行中:\n{key}\n\n锁文件: {path}"
                    raise TaskStartError(detail) from exc
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                handle.write(
                    json.dumps(
                        {
                            "task_id": task_id,
                            "task_dir": str(task_dir),
                            "status_path": str(task_dir / "status.json"),
                            "pid": 0,
                            "created_at": time.time(),
                        },
                        ensure_ascii=False,
                    )
                )
            acquired.append(path)
        return acquired

    def _mark_locks_started(self, lock_paths: Sequence[Path], pid: int) -> None:
        for path in lock_paths:
            try:
                data = json.loads(path.read_text(encoding="utf-8-sig"))
            except Exception:
                data = {}
            data["pid"] = int(pid)
            data["started_at"] = time.time()
            path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    def _release_locks(self, task: ManagedTask) -> None:
        for path in task.lock_paths:
            try:
                path.unlink(missing_ok=True)
            except OSError:
                pass
        task.lock_paths = []

    def start_task(
        self,
        *,
        cloud: str,
        ak: str,
        sk: str,
        action: str,
        title: str,
        entrypoint: str,
        regions: Sequence[str],
        risk_group: str,
        payload: Dict[str, Any] | None = None,
    ) -> ManagedTask:
        if self.running_count_for_cloud(cloud) >= MAX_CLOUD_CONCURRENT_TASKS:
            raise TaskStartError(f"{cloud} 当前运行任务已达到 {MAX_CLOUD_CONCURRENT_TASKS} 个，请等待前面的任务完成。")
        account = account_fingerprint(ak, sk)
        task_id = make_task_id(action)
        task_dir = task_account_dir(cloud, account) / task_id
        task_dir.mkdir(parents=True, exist_ok=True)
        lock_keys = task_lock_keys(cloud, account, regions, risk_group)
        lock_paths = self._acquire_locks(lock_keys, task_id, task_dir)
        spec = {
            "task_id": task_id,
            "cloud": cloud,
            "account": account,
            "action": action,
            "title": title,
            "entrypoint": entrypoint,
            "regions": list(regions),
            "risk_group": risk_group,
            "task_dir": str(task_dir),
            "payload": payload or {},
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        spec_path = task_dir / "task.json"
        atomic_write_json(spec_path, spec)
        atomic_write_json(
            task_dir / "status.json",
            {
                **{key: spec[key] for key in ("task_id", "cloud", "account", "action", "title", "regions", "risk_group", "task_dir")},
                "status": "starting",
                "message": "等待子进程启动",
                "updated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
            },
        )
        env = os.environ.copy()
        env["PYTHONUTF8"] = "1"
        env["PYTHONIOENCODING"] = "utf-8"
        env["WAWA_TASK_DIR"] = str(task_dir)
        env["WAWA_TASK_AK"] = ak
        env["WAWA_TASK_SK"] = sk
        kwargs: Dict[str, Any] = {}
        if os.name == "nt":
            kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
        stdout_handle = None
        stderr_handle = None
        try:
            stdout_handle = (task_dir / "worker_stdout.log").open("ab")
            stderr_handle = (task_dir / "worker_stderr.log").open("ab")
            process = subprocess.Popen(
                self._worker_command(spec_path),
                cwd=str(PROJECT_ROOT),
                env=env,
                stdout=stdout_handle,
                stderr=stderr_handle,
                **kwargs,
            )
            self._mark_locks_started(lock_paths, process.pid)
        except Exception:
            for path in lock_paths:
                path.unlink(missing_ok=True)
            raise
        finally:
            if stdout_handle is not None:
                stdout_handle.close()
            if stderr_handle is not None:
                stderr_handle.close()
        task = ManagedTask(task_id, cloud, account, title, task_dir, process, lock_paths)
        self.tasks[task_id] = task
        return task

    def stop_task(self, task_id: str) -> None:
        task = self.tasks.get(task_id)
        if not task:
            return
        (task.task_dir / "stop.flag").write_text("stop\n", encoding="utf-8")
        if task.process.poll() is None:
            try:
                task.process.terminate()
            except Exception:
                pass

    def refresh(self) -> None:
        for task_id, task in list(self.tasks.items()):
            return_code = task.process.poll()
            if return_code is None:
                continue
            self._release_locks(task)
            self.tasks.pop(task_id, None)

    def list_task_statuses(self, *, limit: int = 80) -> List[Dict[str, Any]]:
        self.refresh()
        status_paths = sorted(app_data_root().glob("tasks/*/accounts/*/tasks/*/status.json"), key=lambda path: path.stat().st_mtime, reverse=True)
        rows = read_status_files(status_paths[:limit])
        rows.sort(key=lambda row: str(row.get("updated_at") or ""), reverse=True)
        return rows
