from __future__ import annotations

from dataclasses import dataclass, field
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
    ocr_results: list[dict[str, Any]] = field(default_factory=list)
    before_screenshot_path: str | None = None
    after_screenshot_path: str | None = None
    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,
        before_screenshot: bytes | None = None,
        after_screenshot: bytes | None = None,
        ocr_results: list[dict[str, Any]] | None = None,
    ) -> TimelineRecord:
        index = len(self.read(task_id)) + 1
        resolved_after_screenshot = (
            after_screenshot if after_screenshot is not None else screenshot
        )
        resolved_ocr_results = list(ocr_results or [])
        record = {
            "index": index,
            "scene": scene,
            "prompt": prompt,
            "tool_call": tool_call,
            "result": result,
            "ocr_results": resolved_ocr_results,
            "timestamp": datetime.now().astimezone().isoformat(),
        }
        paths = self.artifact_store.write_step(
            task_id=task_id,
            index=index,
            before_screenshot=before_screenshot,
            after_screenshot=resolved_after_screenshot,
            record=record,
        )
        return TimelineRecord(
            index=index,
            scene=record["scene"],
            prompt=prompt,
            tool_call=tool_call,
            result=result,
            timestamp=record["timestamp"],
            ocr_results=resolved_ocr_results,
            before_screenshot_path=paths["before_screenshot_path"],
            after_screenshot_path=paths["after_screenshot_path"],
            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)