"""Cooperative, session-scoped task guidance and device lease parking."""

import asyncio
from collections.abc import Awaitable, Callable
import json
import logging
import os
from pathlib import Path
import threading
import time
import uuid

from artemis.config import TRACES_PATH
from artemis.runtime.device_lock import DeviceExecutionLock

logger = logging.getLogger(__name__)

# Windows readers may briefly deny rename/delete even when directory permissions
# allow writes. Bound the total wait to 1.55s; other IO failures remain immediate.
_WINDOWS_ACCESS_RETRY_DELAYS = (0.05, 0.1, 0.2, 0.4, 0.8)


def _retry_windows_access[T](operation: Callable[[], T]) -> T:
    delays = iter(_WINDOWS_ACCESS_RETRY_DELAYS)
    while True:
        try:
            return operation()
        except OSError as exc:
            delay = next(delays, None)
            if delay is None or getattr(exc, "winerror", None) not in {5, 32, 33}:
                raise
            time.sleep(delay)


def _unlink_control_file(path: Path) -> None:
    _retry_windows_access(lambda: path.unlink(missing_ok=True))


def control_dir(session_id: str) -> Path:
    if not session_id or not all(c.isalnum() or c in "-_" for c in session_id):
        raise ValueError("无效的任务编号")
    return TRACES_PATH / session_id


def _write_atomic(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    try:
        temporary.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
        _retry_windows_access(lambda: os.replace(temporary, path))
    except OSError:
        logger.exception("Cannot write task control file %s", path)
        raise
    finally:
        try:
            _unlink_control_file(temporary)
        except OSError:
            # Cleanup must never replace the original write error or turn a
            # successful atomic write into a failed task/control submission.
            logger.warning(
                "Could not remove task control temporary file %s", temporary, exc_info=True
            )


def read_state(session_id: str) -> dict:
    try:
        return json.loads(
            (control_dir(session_id) / "task_control_state.json").read_text(encoding="utf-8")
        )
    except (OSError, ValueError):
        return {}


def request_pause(session_id: str, *, reason: str | None = None) -> None:
    _unlink_control_file(control_dir(session_id) / "resume.request")
    _write_atomic(
        control_dir(session_id) / "pause.request", {"created_at": time.time(), "reason": reason}
    )


def request_resume(session_id: str) -> None:
    _write_atomic(control_dir(session_id) / "resume.request", {"created_at": time.time()})


def add_instruction(
    session_id: str,
    instruction: str,
    *,
    release_loop: bool = False,
    source: str = "user",
    instruction_id: str | None = None,
) -> dict:
    if not instruction.strip() or len(instruction) > 10000:
        raise ValueError("指导内容必须在 1–10000 字之间")
    directory = control_dir(session_id) / "instructions"
    directory.mkdir(parents=True, exist_ok=True)
    # Unique names make simultaneous web and MCP submissions additive, never last-write-wins.
    if instruction_id is not None:
        if not instruction_id or not all(c.isalnum() or c in "-_" for c in instruction_id):
            raise ValueError("无效的指导编号")
        for existing in (
            directory / "consumed" / f"{instruction_id}.json",
            directory / f"{instruction_id}.reading",
            directory / f"{instruction_id}.json",
        ):
            if existing.exists():
                record = json.loads(existing.read_text(encoding="utf-8"))
                if record.get("instruction") != instruction.strip():
                    raise ValueError("指导编号已用于不同内容")
                return record
    path = directory / f"{instruction_id or f'{time.time_ns():020d}-{uuid.uuid4().hex}'}.json"
    record = {
        "id": path.stem,
        "instruction": instruction.strip(),
        "release_loop": bool(release_loop),
        "source": source,
        "timestamp": time.time(),
    }
    _write_atomic(path, record)
    return record


def list_instructions(session_id: str) -> list[dict]:
    """Show submissions separately from guidance the runner actually consumed."""
    directory = control_dir(session_id) / "instructions"
    records = []
    for folder, status in ((directory, "pending"), (directory / "consumed", "consumed")):
        for path in folder.glob("*.json") if folder.exists() else []:
            try:
                record = json.loads(path.read_text(encoding="utf-8"))
                if isinstance(record, dict) and record.get("instruction"):
                    records.append(
                        {**record, "id": record.get("id") or path.stem, "status": status}
                    )
            except (OSError, ValueError):
                continue
    return sorted(records, key=lambda record: (record.get("timestamp", 0), record["id"]))


def has_pending_instruction(base_dir: str | Path) -> bool:
    directory = Path(base_dir) / "instructions"
    return any(directory.glob("*.json")) or (Path(base_dir) / "injected_instruction.json").exists()


def consume_instructions(base_dir: str | Path) -> dict | None:
    directory = Path(base_dir) / "instructions"
    found = []
    for path in sorted(directory.glob("*.json")) if directory.exists() else []:
        try:
            if (directory / "consumed" / path.name).exists():
                path.unlink(missing_ok=True)
                continue
            claimed = path.with_suffix(".reading")
            path.rename(claimed)
            record = json.loads(claimed.read_text(encoding="utf-8"))
            record["consumed_at"] = time.time()
            _write_atomic(directory / "consumed" / path.name, record)
            claimed.unlink()
            found.append(record)
        except (OSError, ValueError):
            continue
    legacy = Path(base_dir) / "injected_instruction.json"
    if legacy.exists():
        try:
            claimed = legacy.with_suffix(".reading")
            legacy.rename(claimed)
            found.append(json.loads(claimed.read_text(encoding="utf-8")))
            claimed.unlink()
        except (OSError, ValueError):
            pass
    if not found:
        return None
    return {
        "instruction": "\n\n".join(str(item.get("instruction") or "") for item in found).strip(),
        "release_loop": any(bool(item.get("release_loop")) for item in found),
    }


class TaskControl:
    """Yield a device lease at the next safe turn and reacquire at the FIFO tail."""

    def __init__(self, session_id: str, device_lock: DeviceExecutionLock):
        self.session_id = session_id
        self.device_lock = device_lock
        self.directory = control_dir(session_id)
        self.budget_epoch = 0
        self.pause_reason: str | None = None

    def _state(self, status: str) -> None:
        _write_atomic(
            self.directory / "task_control_state.json",
            {
                "status": status,
                "updated_at": time.time(),
                "pid": os.getpid(),
                "budget_epoch": self.budget_epoch,
                "pause_reason": self.pause_reason,
            },
        )

    async def checkpoint(self, *, on_pause=None, on_resume=None) -> bool:
        pause = self.directory / "pause.request"
        if not pause.exists():
            return False
        try:
            self.pause_reason = json.loads(pause.read_text(encoding="utf-8")).get("reason")
        except (ValueError, OSError):
            self.pause_reason = None
        _unlink_control_file(pause)
        if on_pause is not None:
            await on_pause()
        self.device_lock.release()
        from artemis.runtime.awake_service import screen_awake_service

        was_awake_owner = self.device_lock.device_id in screen_awake_service.device_ids
        await asyncio.to_thread(screen_awake_service.release_device, self.device_lock.device_id)
        self._state("task_paused")
        resume = self.directory / "resume.request"
        while not resume.exists():
            if self.pause_reason == "question" and self.question_answered():
                break
            await asyncio.sleep(0.2)
        _unlink_control_file(resume)
        self._state("resuming")
        previous = self.device_lock
        self.device_lock = DeviceExecutionLock(
            previous.device_id,
            description=previous.description,
            max_concurrency=previous.max_concurrency,
            concurrency_mode=previous.concurrency_mode,
            session_id=self.session_id,
            ingress=previous.ingress,
            lock_scope=previous.lock_scope,
            resource_ids=getattr(previous, "resource_ids", ()),
        )
        cancelled = threading.Event()
        acquiring = asyncio.create_task(
            asyncio.to_thread(self.device_lock.acquire, cancel_event=cancelled)
        )
        try:
            await asyncio.shield(acquiring)
        except asyncio.CancelledError:
            cancelled.set()
            try:
                await asyncio.shield(acquiring)
            except Exception:
                logger.warning("Resume acquisition ended during cancellation", exc_info=True)
            self.device_lock.release()
            raise
        if was_awake_owner:
            await asyncio.to_thread(screen_awake_service.ensure_device, self.device_lock.device_id)
        if on_resume is not None:
            await on_resume()
        if self.pause_reason == "execution_budget":
            self.budget_epoch += 1
        self.pause_reason = None
        self._state("running")
        return True

    def question_answered(self) -> bool:
        return has_pending_instruction(self.directory)

    async def run_interruptibly[T](
        self, operation: Awaitable[T], *, on_pause=None, on_resume=None
    ) -> tuple[bool, T | None]:
        """Interrupt a stale model/helper call for pause or fresh user guidance.

        Only use for operations whose cancellation cannot leave a device action
        half executed. The caller must return to fresh perception after resume.
        """
        running = asyncio.ensure_future(operation)
        pause = self.directory / "pause.request"

        async def stop_operation():
            running.cancel()
            result = (await asyncio.gather(running, return_exceptions=True))[0]
            if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError):
                raise result

        try:
            while True:
                if running.done():
                    # Preserve a completed tool result (including a permission failure
                    # that requested pause). The next checkpoint handles the request.
                    return False, running.result()
                if pause.exists():
                    await stop_operation()
                    await self.checkpoint(on_pause=on_pause, on_resume=on_resume)
                    return True, None
                if has_pending_instruction(self.directory):
                    await stop_operation()
                    return True, None
                await asyncio.wait({running}, timeout=0.2)
        finally:
            if not running.done():
                await stop_operation()

    def release(self) -> None:
        self.device_lock.release()
        state_path = self.directory / "task_control_state.json"
        try:
            _unlink_control_file(state_path)
        except OSError:
            # The lease is already released. An ephemeral status file left
            # behind must not override the task's result or block the queue.
            logger.warning("Could not remove task control state file %s", state_path, exc_info=True)