from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import Any

from storage.artifact_store import ArtifactStore


@dataclass(frozen=True)
class TimelineRecord:
    index: int
    scene: dict[str, Any]
    # The prompt actually sent to the LLM for this step. For non-LLM planners
    # (or older records persisted before D9), this falls back to the task goal.
    prompt: str
    tool_call: dict[str, Any]
    result: dict[str, Any]
    timestamp: str
    screenshot_path: str | None = None


class Timeline:
    def __init__(self, artifact_store: ArtifactStore | None = None) -> None:
        self.artifact_store = artifact_store or ArtifactStore()

    def append(
        self,
        *,
        task_id: str,
        scene: Any,
        prompt: str,
        tool_call: dict[str, Any],
        result: dict[str, Any],
        screenshot: bytes | None = None,
    ) -> TimelineRecord:
        index = len(self.read(task_id)) + 1
        record = {
            "index": index,
            "scene": scene,
            "prompt": prompt,
            "tool_call": tool_call,
            "result": result,
            "timestamp": datetime.now().astimezone().isoformat(),
        }
        paths = self.artifact_store.write_step(
            task_id=task_id,
            index=index,
            screenshot=screenshot,
            record=record,
        )
        return TimelineRecord(
            index=index,
            scene=record["scene"],
            prompt=prompt,
            tool_call=tool_call,
            result=result,
            timestamp=record["timestamp"],
            screenshot_path=paths["screenshot_path"],
        )

    def read(self, task_id: str) -> list[dict[str, Any]]:
        return self.artifact_store.read_steps(task_id)

    def delete_task(self, task_id: str) -> None:
        """Delete all timeline records and screenshots for a task."""
        self.artifact_store.delete_task(task_id)