from __future__ import annotations

import json
from typing import Any, Mapping, Sequence

from .interfaces import (
    ChunkingStrategy,
    DocumentDownloader,
    EmbeddingStrategy,
    MetadataBuilder,
    TextExtractor,
    VectorStoreWriter,
)
from .models import (
    Chunk,
    DocumentIngestContext,
    DocumentIngestRequest,
    DocumentIngestResult,
)


class DocumentIngestPipeline:
    """
    High-level orchestrator that wires the modular strategies together.
    """

    def __init__(
        self,
        *,
        downloader: DocumentDownloader,
        extractor: TextExtractor,
        chunker: ChunkingStrategy,
        metadata_builder: MetadataBuilder,
        vector_store_writer: VectorStoreWriter,
        embedding_strategies: Mapping[str, EmbeddingStrategy],
        default_embedding_backend: str,
    ) -> None:
        self._downloader = downloader
        self._extractor = extractor
        self._chunker = chunker
        self._metadata_builder = metadata_builder
        self._vector_store_writer = vector_store_writer
        self._embedding_strategies = dict(embedding_strategies)
        self._default_embedding_backend = default_embedding_backend

    async def run(
        self,
        request: DocumentIngestRequest,
        context: DocumentIngestContext,
    ) -> DocumentIngestResult:
        document = await self._downloader.fetch(request)
        text = await self._extractor.extract(document)
        chunks = self._chunker.chunk(text, request, context)
        if not chunks:
            raise ValueError("No chunks produced from document; nothing to persist.")

        enriched_chunks = self._build_metadata(chunks, request, context)
        embedding_backend = request.embedding_backend or self._default_embedding_backend
        embedding_strategy = self._embedding_strategies.get(embedding_backend)
        if embedding_strategy is None:
            raise ValueError(
                f"Unsupported embedding backend '{embedding_backend}'. "
                f"Available: {list(self._embedding_strategies)}"
            )

        embeddings = await embedding_strategy.embed(
            [chunk.text for chunk in enriched_chunks]
        )
        if len(embeddings) != len(enriched_chunks):
            raise ValueError("Embedding strategy returned unexpected vector count.")

        chunk_ids = await self._vector_store_writer.write(
            request.collection_name,
            enriched_chunks,
            embeddings,
        )
        return DocumentIngestResult(
            inserted_chunks=len(chunk_ids),
            collection_name=request.collection_name,
            chunk_ids=chunk_ids,
            embedding_backend=embedding_strategy.name,
            metadata={
                "source_url": request.source_url,
                "mime_type": request.mime_type,
            },
        )

    def _build_metadata(
        self,
        chunks: Sequence[Chunk],
        request: DocumentIngestRequest,
        context: DocumentIngestContext,
    ) -> Sequence[Chunk]:
        normalized: list[Chunk] = []
        for chunk in chunks:
            metadata_dict: dict[str, Any] = dict(context.default_metadata)
            if request.metadata:
                metadata_dict.update(request.metadata)
            limited_metadata = _enforce_metadata_limit(metadata_dict)
            merged = dict(
                self._metadata_builder.build(
                    Chunk(
                        chunk_id=chunk.chunk_id,
                        text=chunk.text,
                        token_count=chunk.token_count,
                        byte_count=chunk.byte_count,
                        metadata=limited_metadata,
                    ),
                    request,
                    context,
                )
            )
            normalized.append(
                Chunk(
                    chunk_id=chunk.chunk_id,
                    text=chunk.text,
                    token_count=chunk.token_count,
                    byte_count=chunk.byte_count,
                    metadata=merged,
                )
            )
        return normalized


def _enforce_metadata_limit(
    metadata: Mapping[str, Any], limit_bytes: int = 4096
) -> dict[str, Any]:
    """
    Trim metadata dict so it stays within Chroma's metadata constraints.
    """
    working = dict(metadata)
    serialized = json.dumps(working, default=str).encode("utf-8")
    if len(serialized) <= limit_bytes:
        return working

    trimmed: dict[str, Any] = {}
    for key, value in working.items():
        candidate = dict(trimmed)
        candidate[key] = value
        size = len(json.dumps(candidate, default=str).encode("utf-8"))
        if size <= limit_bytes:
            trimmed[key] = value
        else:
            trimmed["metadata_truncated"] = True
            break
    return trimmed
