from __future__ import annotations

import asyncio
import json
import os
import time
from pathlib import Path
from typing import Any, Callable

from ..utils import make_ulid as _make_ulid, iso_now as _iso_now

_SNAPSHOT_EVERY_N = 50
_SNAPSHOT_EVERY_MS = 30_000
_SNAPSHOT_KEEP = 10
_LOCK_TIMEOUT = 5.0


class FileEventStore:
    def __init__(self, base_dir: str | Path = "./runs"):
        self._base = Path(base_dir)
        self._last_snapshot: dict[str, tuple[int, int]] = {}  # run_id -> (event_count, ts_ms)
        self._append_locks: dict[str, asyncio.Lock] = {}
        self._event_counts: dict[str, int] = {}

    def _run_dir(self, run_id: str) -> Path:
        d = self._base / run_id
        d.mkdir(parents=True, exist_ok=True)
        return d

    def events_path(self, run_id: str) -> Path:
        return self._run_dir(run_id) / "events.ndjson"

    def _append_to_file(self, path: Path, line: str) -> None:
        """Sync file write with cross-process lock. Runs inside asyncio.to_thread."""
        import sys
        lock_path = path.with_suffix(".lock")
        if sys.platform == "win32":
            import time as _time
            deadline = _time.time() + _LOCK_TIMEOUT
            while _time.time() < deadline:
                try:
                    lf = open(lock_path, "x")
                    break
                except FileExistsError:
                    _time.sleep(0.05)
            else:
                raise TimeoutError(f"Could not acquire lock: {lock_path}")
            try:
                with open(path, "a", encoding="utf-8") as f:
                    f.write(line + "\n")
                    f.flush()
                    os.fsync(f.fileno())
            finally:
                lf.close()
                try:
                    lock_path.unlink()
                except FileNotFoundError:
                    pass
        else:
            import fcntl
            with open(path, "a", encoding="utf-8") as f:
                fcntl.flock(f, fcntl.LOCK_EX)
                try:
                    f.write(line + "\n")
                    f.flush()
                    os.fsync(f.fileno())
                finally:
                    fcntl.flock(f, fcntl.LOCK_UN)

    def _snapshots_dir(self, run_id: str) -> Path:
        d = self._run_dir(run_id) / "snapshots"
        d.mkdir(exist_ok=True)
        return d

    async def append(self, run_id: str, event_type: str, payload: dict) -> str:
        event_id = _make_ulid()
        event = {
            "event_id": event_id,
            "run_id": run_id,
            "ts": _iso_now(),
            "type": event_type,
            "payload": payload,
        }
        line = json.dumps(event, ensure_ascii=False)
        path = self.events_path(run_id)
        async with self._append_locks.setdefault(run_id, asyncio.Lock()):
            await asyncio.to_thread(self._append_to_file, path, line)
        self._event_counts[run_id] = self._event_counts.get(run_id, 0) + 1
        return event_id

    def _read_events(self, run_id: str, after_event_id: str | None = None) -> list[dict]:
        path = self.events_path(run_id)
        if not path.exists():
            return []
        events = []
        found = after_event_id is None
        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    ev = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if not found:
                    if ev.get("event_id") == after_event_id:
                        found = True
                    continue
                events.append(ev)
        return events

    async def replay_all(self, run_id: str, reducer: Callable[[Any, dict], Any]) -> Any:
        state = None
        for ev in self._read_events(run_id):
            state = reducer(state, ev)
        return state

    async def save_snapshot(self, run_id: str, state: Any | None = None, last_event_id: str = "") -> None:
        if not last_event_id:
            events = self._read_events(run_id)
            last_event_id = events[-1]["event_id"] if events else "EMPTY"

        snap = {"last_event_id": last_event_id, "state": state}
        snap_path = self._snapshots_dir(run_id) / f"{last_event_id}.json"

        def _write_and_cleanup():
            tmp = snap_path.with_suffix(".tmp")
            tmp.write_text(json.dumps(snap, ensure_ascii=False), encoding="utf-8")
            tmp.replace(snap_path)
            # 只保留最新 _SNAPSHOT_KEEP 个快照
            snap_dir = self._snapshots_dir(run_id)
            old_snaps = sorted(snap_dir.glob("*.json"), key=lambda p: p.stem)
            for old in old_snaps[:-_SNAPSHOT_KEEP]:
                old.unlink(missing_ok=True)

        await asyncio.to_thread(_write_and_cleanup)

        self._last_snapshot[run_id] = (
            self._event_counts.get(run_id, 0),
            int(time.time() * 1000),
        )

    async def replay_from_snapshot(self, run_id: str, reducer: Callable[[Any, dict], Any]) -> Any:
        snap_dir = self._snapshots_dir(run_id)

        def _load_latest_snapshot():
            snaps = sorted(snap_dir.glob("*.json"), key=lambda p: p.stem)
            if not snaps:
                return None, None
            latest = snaps[-1]
            try:
                data = json.loads(latest.read_text(encoding="utf-8"))
                return data.get("state"), data.get("last_event_id", "")
            except Exception:
                return None, None

        state, last_event_id = await asyncio.to_thread(_load_latest_snapshot)

        if state is None and last_event_id is None:
            return await self.replay_all(run_id, reducer)

        for ev in self._read_events(run_id, after_event_id=last_event_id):
            state = reducer(state, ev)
        return state

    def should_snapshot(self, run_id: str) -> bool:
        count = self._event_counts.get(run_id)
        if count is None:
            # 冷启动兜底：一次性从文件计数并缓存
            count = len(self._read_events(run_id))
            self._event_counts[run_id] = count
        last_count, last_ts = self._last_snapshot.get(run_id, (0, 0))
        elapsed_ms = int(time.time() * 1000) - last_ts
        return (count - last_count >= _SNAPSHOT_EVERY_N) or (
            last_ts > 0 and elapsed_ms >= _SNAPSHOT_EVERY_MS
        )
