from __future__ import annotations

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


def default_embeddings_store() -> Path:
    """Default location for the embeddings configuration file."""
    cache_dir = Path.home() / ".cache" / "hfschema"
    return cache_dir / "embeddings.json"


@dataclass
class EmbeddingsConfig:
    spaces: Dict[str, Dict[str, Any]]
    links: Dict[str, Dict[str, str]]


def _empty_config() -> EmbeddingsConfig:
    return EmbeddingsConfig(spaces={}, links={})


def load_embeddings_config(path: Path | None = None) -> EmbeddingsConfig:
    """Load embeddings configuration from disk, returning an empty config if missing or invalid."""
    if path is None:
        path = default_embeddings_store()
    if not path.exists():
        return _empty_config()
    try:
        raw = json.loads(path.read_text())
    except json.JSONDecodeError:
        return _empty_config()

    spaces = raw.get("spaces") or {}
    links = raw.get("links") or {}
    # Ensure basic shape.
    if not isinstance(spaces, dict):
        spaces = {}
    if not isinstance(links, dict):
        links = {}
    return EmbeddingsConfig(spaces=spaces, links=links)


def save_embeddings_config(config: EmbeddingsConfig, path: Path | None = None) -> Path:
    """Persist embeddings configuration to disk."""
    if path is None:
        path = default_embeddings_store()
    path.parent.mkdir(parents=True, exist_ok=True)
    data = {"spaces": config.spaces, "links": config.links}
    path.write_text(json.dumps(data, indent=2, sort_keys=True))
    return path


def add_space(
    name: str,
    model: str,
    provider: str | None = None,
    dim: int | None = None,
    *,
    store_path: Path | None = None,
) -> EmbeddingsConfig:
    """Add or update an embedding space definition."""
    config = load_embeddings_config(store_path)
    spec: Dict[str, Any] = {"model": model}
    if provider is not None:
        spec["provider"] = provider
    if dim is not None:
        spec["dim"] = dim
    config.spaces[name] = spec
    save_embeddings_config(config, store_path)
    return config


def link_field(
    dataset_key: str,
    field: str,
    space: str,
    *,
    store_path: Path | None = None,
) -> EmbeddingsConfig:
    """Associate a schema field with an embedding space."""
    config = load_embeddings_config(store_path)
    # Do not require the space to exist, but keep the behavior simple and predictable.
    links_for_dataset = config.links.get(dataset_key)
    if links_for_dataset is None:
        links_for_dataset = {}
        config.links[dataset_key] = links_for_dataset
    links_for_dataset[field] = space
    save_embeddings_config(config, store_path)
    return config

