import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

STATE_FILE = ".qc-cli.json"


@dataclass(frozen=True)
class CliStateStore:
    config_dir: str = "."

    @property
    def path(self) -> Path:
        return Path(self.config_dir) / STATE_FILE

    def read(self) -> dict[str, Any]:
        if not self.path.exists():
            return {}
        with open(self.path) as f:
            value = json.load(f)
        return dict(value) if isinstance(value, dict) else {}

    def update(self, **updates: Any) -> None:
        state = self.read()
        state.update(updates)
        self._write(state)

    def get(self, key: str, default: Any = None) -> Any:
        return self.read().get(key, default)

    def get_last_training_job(self) -> str | None:
        value = self.get("last_training_job")
        return str(value) if value else None

    def get_last_model_artifact(self) -> str | None:
        value = self.get("last_model_artifact")
        return str(value) if value else None

    def get_last_quantized_model_id(self) -> str | None:
        value = self.get("last_quantized_model_id")
        return str(value) if value else None

    def get_last_compiled_model_id(self) -> str | None:
        value = self.get("last_compiled_model_id")
        return str(value) if value else None

    def get_last_downloaded_model(self) -> str | None:
        value = self.get("last_downloaded_model")
        return str(value) if value else None

    def set_last_training_job(self, job_name: str) -> None:
        self.update(last_training_job=job_name)

    def get_training_job(self, job_name: str) -> dict[str, Any]:
        jobs = self._training_jobs(self.read())
        value = jobs.get(job_name, {})
        return dict(value) if isinstance(value, dict) else {}

    def update_training_job(self, job_name: str, **updates: Any) -> None:
        state = self.read()
        jobs = self._training_jobs(state)
        jobs[job_name] = {**jobs.get(job_name, {}), **updates}
        state["training_jobs"] = jobs
        self._write(state)

    def set_latest_experiment_model_version(self, version: str) -> None:
        self.update(latest_experiment_model_version=version)

    def get_aihub_model_provenance(self, model_id: str) -> dict[str, Any]:
        provenance = self._aihub_model_provenance(self.read())
        value = provenance.get(model_id, {})
        return dict(value) if isinstance(value, dict) else {}

    def update_aihub_model_provenance(self, model_id: str, provenance: dict[str, Any]) -> None:
        state = self.read()
        model_provenance = self._aihub_model_provenance(state)
        model_provenance[model_id] = provenance
        state["aihub_model_provenance"] = model_provenance
        self._write(state)

    def _write(self, state: dict[str, Any]) -> None:
        with open(self.path, "w") as f:
            json.dump(state, f, indent=2)

    def _training_jobs(self, state: dict[str, Any]) -> dict[str, Any]:
        value = state.get("training_jobs", {})
        return dict(value) if isinstance(value, dict) else {}

    def _aihub_model_provenance(self, state: dict[str, Any]) -> dict[str, Any]:
        value = state.get("aihub_model_provenance", {})
        return dict(value) if isinstance(value, dict) else {}


def store(config_path: str) -> CliStateStore:
    config_dir = str(Path(config_path).parent)
    return CliStateStore(config_dir)