from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable

import faiss
import numpy as np

from mikoshi.chunking import chunk_text
from mikoshi.config import Config, load_config
from mikoshi.hashing import sha256_bytes, sha256_text
from mikoshi.ignore import IgnoreMatcher
from mikoshi.indexing.file_scanner import scan_repo
from mikoshi.indexing.index_store import IndexStore
from mikoshi.retrieval.semantic import get_embeddings_provider, normalize_embeddings
from mikoshi.utils.timer import Timer
from mikoshi.utils.types import Chunk, IndexMeta, IndexResult


class IndexerError(RuntimeError):
    pass


@dataclass(frozen=True)
class IndexedFile:
    path: Path
    relpath: str
    file_hash: str
    text: str


def _read_file(path: Path) -> bytes:
    return path.read_bytes()


def _prepare_file(path: Path, repo_root: Path) -> IndexedFile:
    data = _read_file(path)
    file_hash = sha256_bytes(data)
    text = data.decode("utf-8", errors="ignore")
    relpath = path.relative_to(repo_root).as_posix()
    return IndexedFile(path=path, relpath=relpath, file_hash=file_hash, text=text)


def _chunks_from_text(
    relpath: str,
    file_hash: str,
    text: str,
    max_lines: int,
    overlap: int,
) -> list[Chunk]:
    spans = chunk_text(text, max_lines, overlap)
    chunks: list[Chunk] = []
    for span in spans:
        chunk_id = sha256_text(
            f"{relpath}:{span.start_line}:{span.end_line}:{span.text}"
        )
        chunks.append(
            Chunk(
                id=chunk_id,
                relpath=relpath,
                start_line=span.start_line,
                end_line=span.end_line,
                text=span.text,
                file_hash=file_hash,
                vector_idx=None,
            )
        )
    return chunks


def _with_vector_idx(chunks: Iterable[Chunk]) -> list[Chunk]:
    return [
        Chunk(
            id=chunk.id,
            relpath=chunk.relpath,
            start_line=chunk.start_line,
            end_line=chunk.end_line,
            text=chunk.text,
            file_hash=chunk.file_hash,
            vector_idx=idx,
        )
        for idx, chunk in enumerate(chunks)
    ]


def _reuse_chunks(
    prev_chunks: list[Chunk],
    prev_embeddings: np.ndarray | None,
) -> tuple[dict[str, list[Chunk]], dict[str, np.ndarray]]:
    chunks_by_file: dict[str, list[Chunk]] = {}
    embedding_by_id: dict[str, np.ndarray] = {}

    if not prev_chunks:
        return chunks_by_file, embedding_by_id

    if prev_embeddings is not None:
        for chunk in prev_chunks:
            if chunk.vector_idx is None:
                continue
            if 0 <= chunk.vector_idx < len(prev_embeddings):
                embedding_by_id[chunk.id] = prev_embeddings[chunk.vector_idx]

    for chunk in prev_chunks:
        chunks_by_file.setdefault(chunk.relpath, []).append(chunk)

    return chunks_by_file, embedding_by_id


def index_repo(repo_path: str, config: Config | None = None) -> IndexResult:
    config = config or load_config()
    repo_root = Path(repo_path).expanduser().resolve()

    matcher = IgnoreMatcher(repo_root)
    store = IndexStore(repo_root, config.index_root)

    prev_meta = store.load_meta()
    prev_chunks = store.load_chunks()
    prev_embeddings = store.load_embeddings()
    prev_files = prev_meta.files if prev_meta else {}

    if prev_meta:
        if (
            prev_meta.chunk_lines != config.chunk_lines
            or prev_meta.chunk_overlap != config.chunk_overlap
            or prev_meta.max_bytes != config.max_bytes
            or prev_meta.embedding_provider != config.embeddings.provider
            or prev_meta.model != config.embeddings.model
        ):
            prev_meta = None
            prev_chunks = []
            prev_embeddings = None
            prev_files = {}

    prev_chunks_by_file, prev_embedding_by_id = _reuse_chunks(
        prev_chunks,
        prev_embeddings,
    )

    with Timer() as timer:
        files = scan_repo(repo_root, matcher, config.max_bytes)
        indexed_files: list[IndexedFile] = []
        for path in files:
            try:
                indexed_files.append(_prepare_file(path, repo_root))
            except OSError:
                continue

        file_hashes = {item.relpath: item.file_hash for item in indexed_files}

        new_chunks: list[Chunk] = []
        embedding_slots: list[np.ndarray | None] = []
        pending_texts: list[str] = []
        pending_indices: list[int] = []

        for item in indexed_files:
            relpath = item.relpath
            unchanged = prev_files.get(relpath) == item.file_hash
            if unchanged and relpath in prev_chunks_by_file and prev_embedding_by_id:
                existing_chunks = prev_chunks_by_file[relpath]
                if all(chunk.id in prev_embedding_by_id for chunk in existing_chunks):
                    for chunk in existing_chunks:
                        new_chunks.append(chunk)
                        embedding_slots.append(prev_embedding_by_id[chunk.id])
                    continue

            chunks = _chunks_from_text(
                relpath,
                item.file_hash,
                item.text,
                config.chunk_lines,
                config.chunk_overlap,
            )
            for chunk in chunks:
                pending_indices.append(len(new_chunks))
                new_chunks.append(chunk)
                embedding_slots.append(None)
                pending_texts.append(chunk.text)

        provider = None
        if pending_texts:
            provider = get_embeddings_provider(config)
            pending_embeddings = provider.embed_texts(pending_texts)
        else:
            pending_embeddings = np.zeros((0, 0), dtype=np.float32)

        if provider is not None and provider.dimension:
            dimension = provider.dimension
        elif prev_embeddings is not None and prev_embeddings.size:
            dimension = int(prev_embeddings.shape[1])
        elif prev_meta and prev_meta.embedding_dim:
            dimension = prev_meta.embedding_dim
        else:
            dimension = int(pending_embeddings.shape[1]) if pending_embeddings.size else 1

        for slot_idx, embedding in zip(pending_indices, pending_embeddings):
            embedding_slots[slot_idx] = embedding

        if embedding_slots:
            all_embeddings = np.stack(embedding_slots).astype(np.float32)
        else:
            all_embeddings = np.zeros((0, dimension), dtype=np.float32)
        all_embeddings = normalize_embeddings(all_embeddings)

        final_chunks = _with_vector_idx(new_chunks)

        index = faiss.IndexFlatIP(dimension)
        if all_embeddings.size:
            index.add(all_embeddings)

        now = datetime.now(timezone.utc).isoformat()
        meta = IndexMeta(
            repo_id=store.repo_id,
            repo_path=str(repo_root),
            created_at=prev_meta.created_at if prev_meta else now,
            updated_at=now,
            embedding_provider=config.embeddings.provider,
            model=config.embeddings.model,
            embedding_dim=dimension,
            chunk_lines=config.chunk_lines,
            chunk_overlap=config.chunk_overlap,
            max_bytes=config.max_bytes,
            files=file_hashes,
            chunks=len(final_chunks),
        )

        store.save_chunks(final_chunks)
        store.save_embeddings(all_embeddings)
        store.save_faiss(index)
        store.save_meta(meta)

    return IndexResult(
        repo_id=store.repo_id,
        chunks_indexed=len(final_chunks),
        took_ms=timer.ms,
    )
