from __future__ import annotations

import json
import shutil
from dataclasses import asdict, is_dataclass
from datetime import date, datetime
from pathlib import Path
from typing import Any


class ArtifactStore:
    def __init__(self, root: str | Path = "tasks/history") -> None:
        self.root = Path(root)

    def task_dir(self, task_id: str) -> Path:
        return self.root / task_id

    def write_step(
        self,
        *,
        task_id: str,
        index: int,
        before_screenshot: bytes | None,
        after_screenshot: bytes | None,
        record: dict[str, Any],
    ) -> dict[str, str | None]:
        task_dir = self.task_dir(task_id)
        task_dir.mkdir(parents=True, exist_ok=True)
        stem = f"{index:03d}"
        before_screenshot_path: Path | None = None
        after_screenshot_path: Path | None = None
        if before_screenshot is not None:
            before_screenshot_path = task_dir / f"{stem}-before.png"
            before_screenshot_path.write_bytes(before_screenshot)
        if after_screenshot is not None:
            # Preserve the original filename as the compatibility alias for the
            # screenshot captured after a step.
            after_screenshot_path = task_dir / f"{stem}.png"
            after_screenshot_path.write_bytes(after_screenshot)

        json_path = task_dir / f"{stem}.json"
        payload = {
            **record,
            "before_screenshot_path": (
                str(before_screenshot_path) if before_screenshot_path else None
            ),
            "after_screenshot_path": (
                str(after_screenshot_path) if after_screenshot_path else None
            ),
            "screenshot_path": (
                str(after_screenshot_path) if after_screenshot_path else None
            ),
        }
        json_path.write_text(
            json.dumps(_jsonable(payload), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        return {
            "json_path": str(json_path),
            "before_screenshot_path": (
                str(before_screenshot_path) if before_screenshot_path else None
            ),
            "after_screenshot_path": (
                str(after_screenshot_path) if after_screenshot_path else None
            ),
            "screenshot_path": (
                str(after_screenshot_path) if after_screenshot_path else None
            ),
        }

    def read_steps(self, task_id: str) -> list[dict[str, Any]]:
        task_dir = self.task_dir(task_id)
        if not task_dir.exists():
            return []
        steps = []
        for path in sorted(task_dir.glob("*.json")):
            steps.append(json.loads(path.read_text(encoding="utf-8")))
        return steps

    def delete_task(self, task_id: str) -> None:
        """Remove all on-disk artifacts (JSON + screenshots) for a task."""
        shutil.rmtree(self.task_dir(task_id), ignore_errors=True)


def _jsonable(value: Any) -> Any:
    if hasattr(value, "to_dict"):
        return value.to_dict()
    if is_dataclass(value):
        return asdict(value)
    if isinstance(value, dict):
        return {key: _jsonable(inner) for key, inner in value.items()}
    if isinstance(value, list):
        return [_jsonable(inner) for inner in value]
    if isinstance(value, tuple):
        return [_jsonable(inner) for inner in value]
    if isinstance(value, (datetime, date)):
        return value.isoformat()
    return value