"""Document indexing to search backend."""

import json
import time
from typing import Callable

from rich.console import Console
from rich.progress import Progress

from mykb.core.chunker import chunk_document, content_hash
from mykb.core.collection import get_collection, list_collections
from mykb.core.config import get_state_file, get_embedding_config, get_backend_config
from mykb.backends import get_backend, BaseBackend
from mykb.core.models import MykbDocument
from mykb.sources import get_source

console = Console()

# Default embedding dimension (google/embeddinggemma-300m)
EMBEDDING_DIMENSION = 768


def get_default_backend() -> BaseBackend:
    """Get the configured backend instance."""
    backend_type, backend_config = get_backend_config()
    return get_backend(backend_type, **backend_config)


def load_state() -> dict:
    state_file = get_state_file()
    if state_file.exists():
        return json.loads(state_file.read_text())
    return {}


def save_state(state: dict):
    state_file = get_state_file()
    state_file.parent.mkdir(parents=True, exist_ok=True)
    state_file.write_text(json.dumps(state))


EmbedFn = Callable[[list[str]], list[list[float]]]


def index_collection(
    collection_name: str,
    backend: BaseBackend,
    full: bool = False,
    embed_fn: EmbedFn | None = None,
):
    """Index a single collection into the search backend.

    Args:
        collection_name: Name of the collection to index.
        backend: Backend instance to store documents.
        full: Force full re-index.
        embed_fn: Optional callback for generating embeddings.
    """
    col = get_collection(collection_name)
    embed_config = get_embedding_config()
    state = load_state()
    file_states = state.setdefault('files', {})

    source = get_source(col.source, collection_name, col.model_dump())

    # Collect items first to get a count for the progress bar
    items = list(source.scan())
    if not items:
        console.print(f"[yellow]No items found for collection '{collection_name}'[/yellow]")
        return 0

    doc_count = 0

    with Progress(console=console) as progress:
        task = progress.add_task(f'Indexing {collection_name}', total=len(items))

        for item in items:
            state_key = source.item_state_key(item)
            c_hash = content_hash(item.content)

            # Incremental check
            if not full and file_states.get(state_key) == c_hash:
                progress.advance(task)
                continue
            file_states[state_key] = c_hash

            chunks = chunk_document(
                item.content,
                max_tokens=embed_config['chunk_size'],
                overlap_ratio=embed_config['chunk_overlap'],
            )

            # Build documents
            docs = []
            for chunk in chunks:
                doc_id = f'{collection_name}_{c_hash}_{chunk.index}'
                doc = MykbDocument(
                    id=doc_id,
                    collection=collection_name,
                    source_path=item.source_id,
                    source_uri=item.source_uri,
                    title=item.title,
                    content=chunk.text,
                    chunk_index=chunk.index,
                    category=item.category,
                    tags=item.tags,
                    created_at=item.created_at,
                    modified_at=item.modified_at,
                )
                docs.append(doc)

            # Generate embeddings if embed_fn is provided
            if embed_fn:
                texts = [d.content for d in docs]
                vectors = embed_fn(texts)
                backend_docs = []
                for d, vec in zip(docs, vectors):
                    dd = d.model_dump()
                    dd['_vectors'] = {'gemma': vec}
                    backend_docs.append(dd)
            else:
                backend_docs = [d.model_dump() for d in docs]

            if backend_docs:
                backend.add_documents(collection_name, backend_docs)
                doc_count += len(backend_docs)

            progress.advance(task)

    save_state(state)
    return doc_count


def run_index(
    collection_name: str | None = None,
    full: bool = False,
    embed_fn: EmbedFn | None = None,
):
    """Main index entry point.

    Args:
        collection_name: Index specific collection, or all if None.
        full: Force full re-index.
        embed_fn: Optional callback for generating embeddings.
                  Signature: (texts: list[str]) -> list[list[float]]
                  If None, no embeddings are generated.
    """
    t_start = time.perf_counter()
    backend = get_default_backend()

    # Initialize backend (create index if needed)
    backend.initialize('mykb', EMBEDDING_DIMENSION)

    if collection_name:
        collections = {collection_name: get_collection(collection_name)}
    else:
        collections = list_collections()

    if not collections:
        console.print("[yellow]No collections configured. Use 'mykb collection add' first.[/yellow]")
        return

    total = 0
    for name in collections:
        console.print(f'\n[bold]Indexing collection: {name}[/bold]')
        count = index_collection(name, backend, full=full, embed_fn=embed_fn)
        total += count
        console.print(f'  [green]Indexed {count} chunks[/green]')

    elapsed = time.perf_counter() - t_start
    console.print(f'\n[bold green]Total: {total} chunks indexed ({elapsed:.2f}s)[/bold green]')
