from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from ai_orchestrator.schemas.task_state import ArtifactRecord


class ArtifactStore:
    def __init__(self, base_dir: Path, project_root: Path) -> None:
        self.base_dir = base_dir
        self.project_root = project_root
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def task_dir(self, task_id: str) -> Path:
        path = self.base_dir / task_id
        path.mkdir(parents=True, exist_ok=True)
        return path

    def write_text(self, task_id: str, name: str, content: str) -> ArtifactRecord:
        path = self.task_dir(task_id) / name
        path.write_text(content, encoding="utf-8")
        return ArtifactRecord(type=self._infer_type(name), path=self._relative(path))

    def write_json(self, task_id: str, name: str, payload: Any) -> ArtifactRecord:
        path = self.task_dir(task_id) / name
        path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        return ArtifactRecord(type=self._infer_type(name), path=self._relative(path))

    def _relative(self, path: Path) -> str:
        resolved = path.resolve()
        try:
            return resolved.relative_to(self.project_root.resolve()).as_posix()
        except ValueError:
            return str(resolved)

    @staticmethod
    def _infer_type(name: str) -> str:
        lower_name = name.lower()
        if lower_name.endswith(".diff"):
            return "patch"
        if "test" in lower_name and lower_name.endswith(".txt"):
            return "test_log"
        if lower_name.endswith(".json"):
            return "json"
        if lower_name.endswith(".txt"):
            return "text"
        return "artifact"
