"""Document ingestion pipeline - crawl, chunk, embed, store, integrate.

Orchestrates the full ingestion flow:
1. Crawl documentation source
2. Store raw documents
3. Chunk documents into retrievable segments
4. Generate embeddings
5. Store chunks with embeddings
6. Extract entities and link to knowledge graph (Graph-RAG integration)

Supports both single-document and bulk source ingestion.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Protocol, cast

import structlog

from sibyl.crawler.chunker import ChunkStrategy, DocumentChunker
from sibyl.crawler.embedder import EmbeddingService
from sibyl.crawler.graph_integration import (
    GraphIntegrationService,
    create_graph_integration_service,
)
from sibyl.crawler.local import LocalFileCrawler
from sibyl.crawler.service import CrawlerService
from sibyl.persistence.content_common import (
    ContentConflictError,
    CrawledDocumentRecord,
    CrawlSourceRecord,
    DocumentChunkRecord,
)
from sibyl.persistence.content_runtime import (
    get_content_read_session,
    get_crawl_source_by_id,
    get_document_by_url_for_org,
    list_source_documents,
    save_crawl_source_record,
    save_crawled_document_record,
    save_document_chunks,
)
from sibyl_core.models.entities import Entity, EntityType
from sibyl_core.models.sources import SourceType

if TYPE_CHECKING:
    from uuid import UUID

    from sibyl_core.services.native_graph import NativeEntityManager


class _RollbackSession(Protocol):
    async def rollback(self) -> object: ...


log = structlog.get_logger()

# Type alias for progress callback: receives (stats, chunks_in_last_doc)
ProgressCallback = Callable[["IngestionStats", int], Awaitable[None]]


@dataclass
class IngestionStats:
    """Statistics from an ingestion run."""

    source_id: UUID
    source_name: str
    documents_crawled: int = 0
    documents_stored: int = 0
    chunks_created: int = 0
    embeddings_generated: int = 0
    entities_extracted: int = 0
    entities_linked: int = 0
    errors: int = 0
    duration_seconds: float = 0.0

    def __str__(self) -> str:
        base = (
            f"Ingested {self.source_name}: "
            f"{self.documents_stored} docs, "
            f"{self.chunks_created} chunks, "
            f"{self.embeddings_generated} embeddings"
        )
        if self.entities_extracted > 0:
            base += f", {self.entities_extracted} entities ({self.entities_linked} linked)"
        base += f" in {self.duration_seconds:.1f}s"
        return base


class IngestionPipeline:
    """Full document ingestion pipeline.

    Coordinates crawling, chunking, embedding, storage, and graph integration.
    Handles batching and error recovery.
    """

    def __init__(
        self,
        organization_id: str,
        *,
        chunk_strategy: ChunkStrategy = ChunkStrategy.SEMANTIC,
        generate_embeddings: bool = True,
        embedding_batch_size: int = 50,
        integrate_with_graph: bool = True,
    ) -> None:
        """Initialize the ingestion pipeline.

        Args:
            organization_id: Organization ID for graph operations
            chunk_strategy: Strategy for chunking documents
            generate_embeddings: Whether to generate embeddings
            embedding_batch_size: Batch size for embedding generation
            integrate_with_graph: Whether to extract entities and link to graph
        """
        self.organization_id = organization_id
        self.chunk_strategy = chunk_strategy
        self.generate_embeddings = generate_embeddings
        self.embedding_batch_size = embedding_batch_size
        self.integrate_with_graph = integrate_with_graph

        self._crawler: CrawlerService | None = None
        self._embedder: EmbeddingService | None = None
        self._graph_integration: GraphIntegrationService | None = None
        self._entity_manager: NativeEntityManager | None = None
        self._chunker = DocumentChunker()

    async def start(self) -> None:
        """Start pipeline services."""
        self._crawler = CrawlerService()
        await self._crawler.start()

        if self.generate_embeddings:
            self._embedder = EmbeddingService(batch_size=self.embedding_batch_size)

        if self.integrate_with_graph:
            try:
                self._graph_integration = await create_graph_integration_service(
                    self.organization_id,
                    extract_entities=True,
                    create_new_entities=False,  # Only link to existing entities for now
                )
                self._entity_manager = self._graph_integration.get_entity_manager()
                log.info(
                    "Graph integration enabled",
                    extract_entities=True,
                    entity_manager=bool(self._entity_manager),
                )
            except Exception as e:
                log.warning("Graph integration unavailable", error=str(e), exc_info=True)
                self._graph_integration = None
                self._entity_manager = None

        log.info("Ingestion pipeline started")

    async def stop(self) -> None:
        """Stop pipeline services."""
        if self._crawler:
            await self._crawler.stop()
            self._crawler = None

        log.info("Ingestion pipeline stopped")

    async def __aenter__(self) -> IngestionPipeline:
        """Async context manager entry."""
        await self.start()
        return self

    async def __aexit__(self, *args: object) -> None:
        """Async context manager exit."""
        await self.stop()

    async def ingest_source(
        self,
        source: CrawlSourceRecord,
        *,
        max_pages: int = 100,
        max_depth: int = 3,
        on_progress: ProgressCallback | None = None,
    ) -> IngestionStats:
        """Ingest a full documentation source.

        Crawls all pages, chunks content, generates embeddings,
        and stores everything in the database.

        Args:
            source: CrawlSource to ingest
            max_pages: Maximum pages to crawl
            max_depth: Maximum link depth
            on_progress: Optional callback called after each document with (stats, chunks_created)

        Returns:
            IngestionStats with results
        """
        if not self._crawler:
            raise RuntimeError("Pipeline not started. Use 'async with' or call start()")

        start_time = datetime.now(UTC)
        stats = IngestionStats(
            source_id=source.id,
            source_name=source.name,
        )

        log.info(
            "Starting source ingestion",
            source=source.name,
            url=source.url,
            max_pages=max_pages,
        )

        try:
            # Select crawler and method based on source type
            if source.source_type == SourceType.LOCAL:
                # Local files don't need llms.txt discovery
                crawler: CrawlerService | LocalFileCrawler = LocalFileCrawler()
                doc_stream = crawler.crawl_source(
                    source,
                    max_pages=max_pages,
                    max_depth=max_depth,
                )
            else:
                # Web sources use discovery (probes for llms.txt first)
                if not self._crawler:
                    raise RuntimeError("Web crawler not started")
                doc_stream = self._crawler.crawl_with_discovery(
                    source,
                    max_pages=max_pages,
                    max_depth=max_depth,
                )

            # Crawl and process documents
            async for doc in doc_stream:
                stats.documents_crawled += 1
                chunks_before = stats.chunks_created

                try:
                    # Store document and create chunks
                    await self._process_document(doc, stats, source.source_type)
                    stats.documents_stored += 1

                    # Report progress after each successful document
                    if on_progress:
                        chunks_added = stats.chunks_created - chunks_before
                        try:
                            await on_progress(stats, chunks_added)
                        except Exception as cb_err:
                            log.warning("Progress callback failed", error=str(cb_err))

                except Exception as e:
                    stats.errors += 1
                    log.error(  # noqa: TRY400
                        "Failed to process document",
                        url=doc.url,
                        error=str(e),
                    )

        except Exception as e:
            log.error("Source ingestion failed", source=source.name, error=str(e))  # noqa: TRY400
            stats.errors += 1

        # Auto-tag source based on crawled documents
        if stats.documents_stored > 0:
            await self._update_source_tags(source)

        stats.duration_seconds = (datetime.now(UTC) - start_time).total_seconds()

        log.info(
            "Source ingestion complete",
            source=source.name,
            stats=str(stats),
        )

        return stats

    async def _create_guidance_entity(
        self,
        document: CrawledDocumentRecord,
    ) -> str | None:
        """Create a guidance entity in the knowledge graph for a local file.

        Args:
            document: The crawled document from a local directory

        Returns:
            Entity ID if created, None otherwise
        """
        if not self._entity_manager:
            log.debug("Skipping guidance entity - entity manager not available")
            return None

        try:
            # Extract filename from file:// URL for the entity name
            file_path = document.url.replace("file://", "")
            file_name = file_path.split("/")[-1].replace(".md", "").replace("-", " ").title()

            # Create a summary from the first ~500 chars of content
            description = document.content[:500] if document.content else document.title
            if len(description) > 500:
                description = description[:497] + "..."

            entity = Entity(
                id=f"guidance:{document.id!s}",
                name=file_name,
                entity_type=EntityType.GUIDE,
                description=description,
                content=document.content,
                source_file=file_path,
                metadata={
                    "source_id": str(document.source_id),
                    "url": document.url,
                    "word_count": document.word_count,
                    "headings": document.headings[:10] if document.headings else [],
                },
            )

            entity_id = await self._entity_manager.create_direct(entity)
            log.debug("Created guidance entity", entity_id=entity_id, name=file_name)
            return entity_id

        except Exception as e:
            log.warning(
                "Failed to create guidance entity",
                url=document.url,
                error=str(e),
            )
            return None

    async def _update_source_tags(self, source: CrawlSourceRecord) -> None:
        """Update source with auto-detected tags and favicon.

        Fetches all documents for the source, extracts tags using heuristics,
        fetches favicon, and updates the CrawlSource record.
        """
        from sibyl.crawler.tagger import aggregate_source_tags

        try:
            async with get_content_read_session() as session:
                documents = await list_source_documents(session, source_id=source.id)

                if not documents:
                    return

                # Extract and aggregate tags
                tags, categories = aggregate_source_tags(list(documents))

                # Fetch favicon (only for web sources)
                favicon_url = None
                if self._crawler and source.url.startswith("http"):
                    try:
                        favicon_url = await self._crawler.fetch_favicon(source.url)
                    except Exception as e:
                        log.debug("Failed to fetch favicon", source=source.name, error=str(e))

                # Update source in database
                db_source = await get_crawl_source_by_id(session, source_id=source.id)
                if db_source:
                    db_source.tags = tags
                    db_source.categories = categories
                    if favicon_url:
                        db_source.favicon_url = favicon_url
                    await save_crawl_source_record(session, source=db_source)
                    log.info(
                        "Auto-detected source metadata",
                        source=source.name,
                        tags=tags[:10],
                        categories=categories,
                        favicon=bool(favicon_url),
                    )

        except Exception as e:
            log.warning("Failed to update source metadata", source=source.name, error=str(e))

    async def _process_document(
        self,
        document: CrawledDocumentRecord,
        stats: IngestionStats,
        source_type: SourceType | None = None,
    ) -> None:
        """Process a single document - store, chunk, embed, integrate with graph.

        Args:
            document: Document to process
            stats: Stats to update
            source_type: Type of source (LOCAL creates guidance entities)
        """
        db_chunks: list[DocumentChunkRecord] = []
        saved_chunks: list[DocumentChunkRecord] = []
        stored_document = document

        async with get_content_read_session() as session:
            source = await get_crawl_source_by_id(session, source_id=document.source_id)
            if source is None:
                msg = f"Source not found for document {document.source_id}"
                raise ValueError(msg)

            existing = await get_document_by_url_for_org(
                session,
                url=document.url,
                organization_id=source.organization_id,
            )
            if existing is not None:
                log.debug("Document already exists, skipping", url=document.url)
                return

            try:
                stored_document = await save_crawled_document_record(session, document=document)
            except ContentConflictError as exc:
                if session is not None and hasattr(session, "rollback"):
                    await cast("_RollbackSession", session).rollback()
                existing = await get_document_by_url_for_org(
                    session,
                    url=document.url,
                    organization_id=source.organization_id,
                )
                if existing is None:
                    raise
                log.debug(
                    "Document inserted by concurrent crawl, skipping",
                    url=document.url,
                    error=str(exc),
                )
                return

            chunks = self._chunker.chunk_document(
                stored_document,
                strategy=self.chunk_strategy,
            )
            stats.chunks_created += len(chunks)

            if not chunks:
                log.debug("No chunks created for document", url=document.url)
                return

            # Generate embeddings if enabled
            embeddings = None
            if self.generate_embeddings and self._embedder:
                try:
                    embeddings = await self._embedder.embed_chunks(chunks)
                    stats.embeddings_generated += len(embeddings)
                except Exception as e:
                    log.warning(
                        "Failed to generate embeddings",
                        url=document.url,
                        error=str(e),
                    )

            # Store chunks
            for i, chunk in enumerate(chunks):
                db_chunk = DocumentChunkRecord(
                    document_id=stored_document.id,
                    chunk_index=chunk.chunk_index,
                    chunk_type=chunk.chunk_type,
                    content=chunk.content,
                    context=chunk.context,
                    token_count=chunk.token_count,
                    start_char=chunk.start_char,
                    end_char=chunk.end_char,
                    heading_path=chunk.heading_path,
                    language=chunk.language,
                    embedding=embeddings[i] if embeddings and i < len(embeddings) else None,
                    is_complete=True,
                    has_entities=False,
                    entity_ids=[],
                )
                db_chunks.append(db_chunk)

            saved_chunks = await save_document_chunks(session, chunks=db_chunks)

        # Graph integration: extract entities and link to knowledge graph
        if self._graph_integration and saved_chunks:
            try:
                integration_stats = await self._graph_integration.process_chunks(
                    saved_chunks,
                    source_name=stored_document.url,
                )
                stats.entities_extracted += integration_stats.entities_extracted
                stats.entities_linked += integration_stats.entities_linked

                # Create DOCUMENTED_IN relationships for linked entities
                entity_uuids = []
                for chunk in saved_chunks:
                    if chunk.entity_ids:
                        entity_uuids.extend(chunk.entity_ids)

                if entity_uuids:
                    await self._graph_integration.create_doc_relationships(
                        stored_document.id,
                        list(set(entity_uuids)),  # Dedupe
                        document_title=stored_document.title,
                        document_url=stored_document.url,
                    )

            except Exception as e:
                log.warning(
                    "Graph integration failed for document",
                    url=stored_document.url,
                    error=str(e),
                )

        # Create guidance entity for local sources
        if source_type == SourceType.LOCAL:
            await self._create_guidance_entity(stored_document)

        log.debug(
            "Processed document",
            url=stored_document.url,
            chunks=len(chunks),
            embeddings=len(embeddings) if embeddings else 0,
            entities=stats.entities_extracted,
        )

    async def ingest_url(
        self,
        url: str,
        source: CrawlSourceRecord,
    ) -> CrawledDocumentRecord | None:
        """Ingest a single URL.

        Args:
            url: URL to crawl
            source: Parent source

        Returns:
            Created document or None if failed
        """
        if not self._crawler:
            raise RuntimeError("Pipeline not started")

        result = await self._crawler.crawl_page(url)
        if not result.success:
            log.warning("Failed to crawl URL", url=url, error=result.error_message)
            return None

        # Create document
        doc = self._crawler.result_to_document(result, source)

        # Store and process
        stats = IngestionStats(source_id=source.id, source_name=source.name)
        await self._process_document(doc, stats, source.source_type)

        return doc


async def ingest_documentation(
    name: str,
    url: str,
    *,
    organization_id: str,
    max_pages: int = 100,
    max_depth: int = 3,
    include_patterns: list[str] | None = None,
) -> IngestionStats:
    """Convenience function to ingest a documentation site.

    Creates a source if needed and runs full ingestion.

    Args:
        name: Human-readable name for the source
        url: Base URL to crawl
        organization_id: Organization UUID for multi-tenant isolation.
        max_pages: Maximum pages to crawl
        max_depth: Maximum link depth
        include_patterns: URL patterns to include

    Returns:
        IngestionStats with results
    """
    from sibyl.crawler.service import create_source, get_source_by_url

    # Get or create source
    source = await get_source_by_url(url)
    if source is None:
        source = await create_source(
            name=name,
            url=url,
            organization_id=organization_id,
            include_patterns=include_patterns,
        )

    # Run ingestion
    async with IngestionPipeline(organization_id) as pipeline:
        return await pipeline.ingest_source(
            source,
            max_pages=max_pages,
            max_depth=max_depth,
        )


async def reingest_source(source_id: UUID, organization_id: str) -> IngestionStats:
    """Re-ingest an existing source.

    Useful for refreshing stale documentation.

    Args:
        source_id: UUID of source to re-ingest
        organization_id: Organization ID for graph operations

    Returns:
        IngestionStats with results
    """
    async with get_content_read_session() as session:
        source = await get_crawl_source_by_id(session, source_id=source_id)
        if source is None:
            raise ValueError(f"Source not found: {source_id}")

    async with IngestionPipeline(organization_id) as pipeline:
        return await pipeline.ingest_source(source)
