from __future__ import annotations

import json
import logging
import os
import tempfile
from pathlib import Path
from typing import Dict

from solutions.graph.models import (
    ConnectionConfiguration,
    GraphAlreadyExistsError,
    GraphDoesNotExistError,
    GraphId,
    GraphMetadata,
    VersionedGraphMetadata,
)
from stores.graph_metadata_store import (
    AbstractGraphMetadataStore,
    GraphMetadataStoreError,
)
from stores.utils import hash_dict

logger = logging.getLogger(__name__)


class JsonFileGraphMetadataStore(AbstractGraphMetadataStore):
    """
    Store graph metadata as individual JSON files in a directory.
    Each graph configuration is stored in a separate file named {graph_id}.json
    """

    def __init__(self, storage_dir: str) -> None:
        logger.debug(f"Initializing JSON graph store with storage at %s.", storage_dir)

        self.__storage_dir = Path(storage_dir)
        self.__ensure_storage_dir_exists()

    def __ensure_storage_dir_exists(self) -> None:
        try:
            self.__storage_dir.mkdir(parents=True, exist_ok=True)
        except Exception as ex:
            raise GraphMetadataStoreError(f"Failed to create storage directory {self.__storage_dir}.") from ex

    def __get_file_path(self, graph_id: GraphId) -> Path:
        return self.__storage_dir / f"{graph_id}.json"

    def __load_metadata_from_file(self, file_path: Path) -> VersionedGraphMetadata:
        try:
            logger.debug(f"Loading JSON graph metadata from file at %s.", file_path)
            with open(file_path, "r", encoding="utf-8") as f:
                data = json.load(f)
                graph_metadata: GraphMetadata = self.__from_dict(data)
                versioned_metadata: VersionedGraphMetadata = graph_metadata  # type: ignore
                versioned_metadata["version_token"] = hash_dict(graph_metadata)
                return versioned_metadata
        except json.JSONDecodeError as ex:
            raise GraphMetadataStoreError(f"Failed to parse JSON from file {file_path}.") from ex
        except Exception as ex:
            raise GraphMetadataStoreError(f"Failed to load metadata from file {file_path}.") from ex

    def __save_metadata_to_file(self, file_path: Path, metadata: GraphMetadata) -> None:
        """
        Write data to a temporary file and then atomically move it to the target location.
        This ensures that the file is not left in a corrupted state if the write operation fails.
        """
        logger.debug(f"Saving JSON graph metadata to file at %s.", file_path)
        tmp_file_path: str | None = None
        try:
            data = self.__to_dict(metadata)
            with tempfile.NamedTemporaryFile("w", delete=False, dir=file_path.parent, encoding="utf-8") as tmp_file:
                json.dump(data, tmp_file, indent=2, ensure_ascii=False)
                tmp_file_path = tmp_file.name

                os.replace(tmp_file_path, file_path)
        except Exception as ex:
            if tmp_file_path and os.path.exists(tmp_file_path):
                os.remove(tmp_file_path)
            raise GraphMetadataStoreError(f"Failed to save metadata to file {file_path}.") from ex

    def __to_dict(self, metadata: GraphMetadata) -> dict:
        return {
            "id": metadata["id"],
            "name": metadata["name"],
            "nodes": metadata["nodes"],
            "edges": metadata["edges"],
            "nodes_view": metadata["nodes_view"],
            "edges_view": metadata["edges_view"],
            "sampling": metadata["sampling"],
            "cypher_queries": metadata["cypher_queries"],
            "connection": metadata.get("connection", {"type": "kuzu", "params": None}),
        }

    def __from_dict(self, data: dict) -> GraphMetadata:
        connection_data = data.get("connection", {"type": "kuzu", "params": None})
        connection = ConnectionConfiguration(
            type=connection_data.get("type", "kuzu")
        )

        return {
            "id": data["id"],
            "name": data["name"],
            "nodes": data["nodes"],
            "edges": data["edges"],
            "nodes_view": data["nodes_view"],
            "edges_view": data["edges_view"],
            "sampling": data.get("sampling", {"sampling": False, "max_rows": 1000}),
            "cypher_queries": data.get("cypher_queries", []),
            "connection": connection,
        }

    def init_dataset_schema(self) -> None:
        # Initialize the storage directory. No schema needed for JSON files.
        self.__ensure_storage_dir_exists()

    def exists(self, id: GraphId) -> bool:
        file_path = self.__get_file_path(id)
        return file_path.exists()

    def get_all(self) -> Dict[GraphId, VersionedGraphMetadata]:
        try:
            graphs_metadata = {}

            for file_path in self.__storage_dir.glob("*.json"):
                try:
                    metadata = self.__load_metadata_from_file(file_path)
                    graphs_metadata[metadata["id"]] = metadata
                except Exception as ex:
                    logger.warning(f"Failed to load metadata from {file_path}: {ex}")
                    continue

            return graphs_metadata
        except Exception as ex:
            raise GraphMetadataStoreError(f"Failed to load graph metadata from directory {self.__storage_dir}.") from ex

    def get(self, id: GraphId) -> VersionedGraphMetadata:
        file_path = self.__get_file_path(id)

        if not file_path.exists():
            raise GraphDoesNotExistError(id)

        return self.__load_metadata_from_file(file_path)

    def delete(self, id: GraphId) -> None:
        file_path = self.__get_file_path(id)

        if not file_path.exists():
            raise GraphDoesNotExistError(id)

        try:
            file_path.unlink()
        except Exception as ex:
            raise GraphMetadataStoreError(f"Failed to delete metadata file for graph {id} at {file_path}.") from ex

    def update(self, metadata: GraphMetadata, ref_version_token: str) -> VersionedGraphMetadata:
        graph_id = metadata["id"]

        file_path = self.__get_file_path(graph_id)

        if not file_path.exists():
            raise GraphDoesNotExistError(graph_id)

        self.__save_metadata_to_file(file_path, metadata)

        return metadata | {"version_token": hash_dict(metadata)}  # type: ignore

    def create(self, metadata: GraphMetadata) -> VersionedGraphMetadata:
        graph_id = metadata["id"]

        file_path = self.__get_file_path(graph_id)

        if file_path.exists():
            raise GraphAlreadyExistsError(graph_id)

        self.__save_metadata_to_file(file_path, metadata)

        return metadata | {"version_token": hash_dict(metadata)}  # type: ignore
