from __future__ import annotations

import json
import os
import re
import time
from pathlib import Path
from typing import Any
from uuid import uuid4

from pydantic import BaseModel

from .domain import TaskRecord


class TaskNotFoundError(FileNotFoundError):
    pass


class InvalidTaskId(ValueError):
    pass


#: Task ids are generated as uuid4 hex; anything else is rejected before it reaches the
#: filesystem, so an id from an HTTP path can never escape the runtime directory.
_VALID_TASK_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$")


def validate_task_id(task_id: str) -> str:
    if not isinstance(task_id, str) or not _VALID_TASK_ID.match(task_id):
        raise InvalidTaskId(f"invalid task id: {task_id!r}")
    return task_id


def _read_with_retry(path: Path, budget_s: float = 5.0) -> str:
    """Read a file that another thread may be atomically replacing right now.

    Windows denies access to both sides of a replace for the instant it takes, so a
    reader has to tolerate it exactly as the writer does.
    """
    deadline = time.monotonic() + budget_s
    delay = 0.005
    while True:
        try:
            return path.read_text(encoding="utf-8")
        except PermissionError:
            if time.monotonic() >= deadline:
                raise
            time.sleep(delay)
            delay = min(delay * 2, 0.1)


def _replace_with_retry(temporary: Path, path: Path, budget_s: float = 5.0) -> None:
    """Atomic replace, retried briefly.

    On Windows a replace fails while any reader still holds the destination open, which
    happens whenever a cancel saves state during a concurrent load. ponytail: a bounded
    retry, not a lock; revisit if writers ever outnumber this single-task design.
    """
    deadline = time.monotonic() + budget_s
    delay = 0.005
    while True:
        try:
            os.replace(temporary, path)
            return
        except PermissionError:
            if time.monotonic() >= deadline:
                temporary.unlink(missing_ok=True)
                raise
            time.sleep(delay)
            delay = min(delay * 2, 0.1)


class TaskStore:
    def __init__(self, root: str | Path) -> None:
        self.root = Path(root)

    def task_dir(self, task_id: str) -> Path:
        directory = (self.root / "tasks" / validate_task_id(task_id)).resolve()
        base = (self.root / "tasks").resolve()
        if base not in directory.parents:
            raise InvalidTaskId(f"task id escapes the runtime directory: {task_id!r}")
        return directory

    def create(self, record: TaskRecord) -> None:
        directory = self.task_dir(record.task_id)
        directory.mkdir(parents=True, exist_ok=False)
        (directory / "logs").mkdir(exist_ok=True)
        self.write_task_manifest(record)
        self._write_json(directory / "state.json", record.model_dump(mode="json"))
        self.append_event(record.task_id, {"event": "task_created", "state": record.state.value})

    def save(self, record: TaskRecord) -> None:
        if not self.task_dir(record.task_id).is_dir():
            raise TaskNotFoundError(record.task_id)
        self._write_json(self.task_dir(record.task_id) / "state.json", record.model_dump(mode="json"))

    def load(self, task_id: str) -> TaskRecord:
        path = self.task_dir(task_id) / "state.json"
        if not path.is_file():
            raise TaskNotFoundError(task_id)
        return TaskRecord.model_validate_json(_read_with_retry(path))

    def write_artifact(self, task_id: str, name: str, artifact: BaseModel) -> Path:
        path = self.task_dir(task_id) / name
        self._write_json(path, artifact.model_dump(mode="json"))
        return path

    def write_task_manifest(self, record: TaskRecord) -> None:
        """The immutable task definition. Mutable progress stays only in state.json."""
        self._write_json(
            self.task_dir(record.task_id) / "task.json",
            {
                "schema_version": record.schema_version,
                "task_id": record.task_id,
                "repo_path": str(record.repo_path),
                "goal": record.goal,
                "constraints": record.constraints,
                "baseline_commit": record.baseline_commit,
                "worktree_path": str(record.worktree_path) if record.worktree_path else None,
                "dirty_files_at_start": record.dirty_files,
                "created_at": record.created_at.isoformat(),
            },
        )

    def write_json(self, task_id: str, name: str, value: dict) -> Path:
        path = self.task_dir(task_id) / name
        self._write_json(path, value)
        return path

    def write_text(self, task_id: str, name: str, body: str) -> Path:
        path = self.task_dir(task_id) / name
        path.write_text(body, encoding="utf-8", newline=chr(10))
        return path

    def logs_dir(self, task_id: str) -> Path:
        directory = self.task_dir(task_id) / "logs"
        directory.mkdir(parents=True, exist_ok=True)
        return directory

    def append_event(self, task_id: str, event: dict[str, Any]) -> None:
        path = self.task_dir(task_id) / "events.jsonl"
        with path.open("a", encoding="utf-8", newline="\n") as stream:
            stream.write(json.dumps(event, sort_keys=True) + "\n")
            stream.flush()
            os.fsync(stream.fileno())

    @staticmethod
    def _write_json(path: Path, value: dict[str, Any]) -> None:
        # A unique temp name per write: cancel runs concurrently with the stage that is
        # saving state, and a shared temp path makes the two clobber each other.
        temporary = path.with_suffix(f"{path.suffix}.{os.getpid()}.{uuid4().hex[:8]}.tmp")
        with temporary.open("w", encoding="utf-8", newline="\n") as stream:
            json.dump(value, stream, sort_keys=True)
            stream.flush()
            os.fsync(stream.fileno())
        _replace_with_retry(temporary, path)