from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from time_machine.domain.events import TemporalEvent


class JsonlTraceSink:
    def __init__(self, path: Path) -> None:
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)
        if not self.path.exists():
            self.path.touch()

    def record_event(self, event: TemporalEvent) -> None:
        self._write({"kind": "event", "event": event.model_dump(mode="json")})

    def record_model_io(
        self,
        component: str,
        prompt: str,
        raw_output: str,
        parsed_output: Any,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self._write(
            {
                "kind": "model_io",
                "component": component,
                "prompt": prompt,
                "raw_output": raw_output,
                "parsed_output": parsed_output,
                "metadata": metadata or {},
            }
        )

    def _write(self, record: dict[str, Any]) -> None:
        record["created_at"] = datetime.now(timezone.utc).isoformat()
        with self.path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, ensure_ascii=True) + "\n")