from __future__ import annotations

import json
import shutil
from pathlib import Path
from typing import Iterable

import numpy as np
import faiss

from mikoshi.hashing import sha256_text
from mikoshi.utils.types import Chunk, IndexMeta


class IndexStore:
    def __init__(self, repo_root: Path, index_root: Path) -> None:
        self.repo_root = repo_root
        self.repo_id = repo_id_for_path(repo_root)
        self.store_dir = index_root / self.repo_id
        self.meta_path = self.store_dir / "meta.json"
        self.chunks_path = self.store_dir / "chunks.jsonl"
        self.embeddings_path = self.store_dir / "embeddings.npy"
        self.faiss_path = self.store_dir / "index.faiss"

    def exists(self) -> bool:
        return self.meta_path.exists() and self.chunks_path.exists()

    def ensure_dir(self) -> None:
        self.store_dir.mkdir(parents=True, exist_ok=True)

    def clear(self) -> None:
        if self.store_dir.exists():
            shutil.rmtree(self.store_dir)

    def load_meta(self) -> IndexMeta | None:
        if not self.meta_path.exists():
            return None
        data = json.loads(self.meta_path.read_text(encoding="utf-8"))
        return IndexMeta.from_dict(data)

    def save_meta(self, meta: IndexMeta) -> None:
        self.ensure_dir()
        self.meta_path.write_text(
            json.dumps(meta.to_dict(), indent=2, sort_keys=True),
            encoding="utf-8",
        )

    def load_chunks(self) -> list[Chunk]:
        if not self.chunks_path.exists():
            return []
        chunks: list[Chunk] = []
        with self.chunks_path.open("r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                chunks.append(Chunk.from_dict(json.loads(line)))
        return chunks

    def save_chunks(self, chunks: Iterable[Chunk]) -> None:
        self.ensure_dir()
        with self.chunks_path.open("w", encoding="utf-8") as handle:
            for chunk in chunks:
                handle.write(json.dumps(chunk.to_dict(), ensure_ascii=False))
                handle.write("\n")

    def load_embeddings(self) -> np.ndarray | None:
        if not self.embeddings_path.exists():
            return None
        return np.load(self.embeddings_path)

    def save_embeddings(self, embeddings: np.ndarray) -> None:
        self.ensure_dir()
        np.save(self.embeddings_path, embeddings)

    def load_faiss(self) -> faiss.Index | None:
        if not self.faiss_path.exists():
            return None
        return faiss.read_index(str(self.faiss_path))

    def save_faiss(self, index: faiss.Index) -> None:
        self.ensure_dir()
        faiss.write_index(index, str(self.faiss_path))


def repo_id_for_path(repo_root: Path) -> str:
    return sha256_text(str(repo_root.resolve()))[:16]
