"""
Embedding pipeline orchestrator.

Coordinates the full ingestion flow:
parse → chunk → embed → store → graph construction.
"""

import hashlib
import logging
import time
from typing import Optional

from django.conf import settings
from django.core.cache import cache
from django.core.files.base import ContentFile
from django.core.files.storage import default_storage

from library.metrics import (
    PIPELINE_DURATION,
    PIPELINE_ITEMS_IN_PROGRESS,
    PIPELINE_ITEMS_TOTAL,
)

logger = logging.getLogger(__name__)

# S3 key patterns
ORIGINAL_S3_KEY = "items/{item_uid}/original.{ext}"
CHUNK_S3_KEY = "chunks/{item_uid}/chunk_{index}.txt"
IMAGE_S3_KEY = "images/{item_uid}/{index}.{ext}"

# Batch sizes
EMBEDDING_BATCH_SIZE = getattr(settings, "EMBEDDING_BATCH_SIZE", 8)


class EmbeddingPipeline:
    """
    Orchestrates the complete embedding pipeline for a single Item.

    Stages:
    1. Parse document (text + images)
    2. Chunk text (content-type-aware)
    3. Store chunks in S3 + Neo4j
    4. Embed text chunks
    5. Store images in S3 + Neo4j
    6. Embed images (multimodal, if available)
    7. Extract concepts (if system chat model available)
    """

    def __init__(self, user=None):
        """
        :param user: Optional Django user for usage tracking.
        """
        self.user = user

    def process_item(
        self,
        item_uid: str,
        progress_callback: Optional[callable] = None,
    ) -> dict:
        """
        Run the full embedding pipeline for an Item.

        :param item_uid: UID of the Item node to process.
        :param progress_callback: Optional callback(percent, message).
        :returns: Dict with processing results.
        :raises ValueError: If item not found or no embedding model configured.
        """
        from library.models import Item

        start_time = time.time()
        PIPELINE_ITEMS_IN_PROGRESS.inc()

        try:
            item = Item.nodes.get(uid=item_uid)
        except Exception:
            PIPELINE_ITEMS_TOTAL.labels(status="failed").inc()
            PIPELINE_ITEMS_IN_PROGRESS.dec()
            raise ValueError(f"Item not found: {item_uid}")

        logger.info(
            "Pipeline starting item_uid=%s title='%s' file_type=%s",
            item_uid,
            item.title,
            item.file_type,
        )

        # Mark as processing
        item.embedding_status = "processing"
        item.error_message = ""
        item.save()

        try:
            result = self._run_pipeline(item, progress_callback)

            # Mark as completed
            item.embedding_status = "completed"
            item.chunk_count = result.get("chunks_created", 0)
            item.image_count = result.get("images_stored", 0)
            item.embedding_model_name = result.get("model_name", "")
            item.error_message = ""
            item.save()

            elapsed = time.time() - start_time
            PIPELINE_ITEMS_TOTAL.labels(status="completed").inc()
            PIPELINE_DURATION.observe(elapsed)

            logger.info(
                "Pipeline completed item_uid=%s chunks=%d images=%d concepts=%d elapsed=%.2fs",
                item_uid,
                result.get("chunks_created", 0),
                result.get("images_stored", 0),
                result.get("concepts_extracted", 0),
                elapsed,
            )

            if progress_callback:
                progress_callback(100, "Completed")

            return result

        except Exception as exc:
            item.embedding_status = "failed"
            item.error_message = str(exc)[:500]
            item.save()

            PIPELINE_ITEMS_TOTAL.labels(status="failed").inc()
            logger.error(
                "Pipeline failed item_uid=%s: %s",
                item_uid,
                exc,
                exc_info=True,
            )
            raise

        finally:
            PIPELINE_ITEMS_IN_PROGRESS.dec()

    def _run_pipeline(self, item, progress_callback) -> dict:
        """
        Execute pipeline stages sequentially.

        :param item: Item node instance.
        :param progress_callback: Optional progress callback.
        :returns: Results dict.
        """
        from llm_manager.models import LLMModel

        from .chunker import ContentTypeChunker
        from .concepts import ConceptExtractor
        from .embedding_client import EmbeddingClient
        from .parsers import DocumentParser

        result = {
            "chunks_created": 0,
            "chunks_embedded": 0,
            "images_stored": 0,
            "images_analyzed": 0,
            "images_embedded": 0,
            "concepts_extracted": 0,
            "model_name": "",
        }

        # --- Resolve library context ---
        library = self._get_item_library(item)
        chunking_config = library.chunking_config if library else {}
        embedding_instruction = library.embedding_instruction if library else ""
        library_type = library.library_type if library else ""

        # --- Get system embedding model ---
        embedding_model = LLMModel.get_system_embedding_model()
        if not embedding_model:
            raise ValueError(
                "No system embedding model configured. "
                "Set one via Django admin > LLM Models > Set as System Embedding Model."
            )

        result["model_name"] = embedding_model.name
        embed_client = EmbeddingClient(embedding_model, user=self.user)

        # --- Check dimension compatibility ---
        if embedding_model.vector_dimensions:
            self._check_dimension_compatibility(embedding_model.vector_dimensions)

        if progress_callback:
            progress_callback(5, "Parsing document")

        # --- Stage 1: Parse ---
        parser = DocumentParser()
        file_data = self._read_item_from_s3(item)
        if not file_data:
            logger.warning("No file data for item_uid=%s, skipping", item.uid)
            return result

        parse_result = parser.parse_bytes(
            file_data,
            item.file_type,
            filename=item.title,
        )

        if progress_callback:
            progress_callback(20, "Chunking text")

        # --- Stage 2: Chunk ---
        chunker = ContentTypeChunker()
        chunk_result = chunker.chunk(parse_result, chunking_config, library_type)

        if progress_callback:
            progress_callback(30, "Storing chunks")

        # --- Stage 3: Store chunks in S3 + Neo4j ---
        chunk_nodes = self._store_chunks(item, chunk_result)
        result["chunks_created"] = len(chunk_nodes)

        if progress_callback:
            progress_callback(40, "Embedding text chunks")

        # --- Stage 4: Embed text chunks ---
        if chunk_result.chunks:
            self._embed_chunks(
                item,
                chunk_nodes,
                chunk_result.chunks,
                embed_client,
                embedding_instruction,
                progress_callback,
            )
            result["chunks_embedded"] = len(chunk_nodes)

        if progress_callback:
            progress_callback(70, "Storing images")

        # --- Stage 5: Store images ---
        image_nodes = self._store_images(item, parse_result.images)
        result["images_stored"] = len(image_nodes)

        # Associate images with nearby chunks
        self._associate_images_with_chunks(
            chunk_nodes, image_nodes, chunk_result, parse_result
        )

        if progress_callback:
            progress_callback(75, "Analyzing images with vision model")

        # --- Stage 5.5: Vision analysis (Phase 2B) ---
        vision_model = LLMModel.get_system_vision_model()
        if vision_model and image_nodes:
            from .vision import VisionAnalyzer

            vision_prompt = self._get_vision_prompt(library)
            analyzer = VisionAnalyzer(vision_model, user=self.user)
            images_analyzed = analyzer.analyze_images(
                image_nodes,
                vision_prompt=vision_prompt,
                item=item,
            )
            result["images_analyzed"] = images_analyzed
        elif image_nodes:
            # No vision model — mark images as skipped
            for img_node in image_nodes:
                img_node.analysis_status = "skipped"
                img_node.save()

        if progress_callback:
            progress_callback(85, "Embedding images")

        # --- Stage 6: Embed images (multimodal) ---
        if image_nodes and embedding_model.supports_multimodal:
            embedded_count = self._embed_images(image_nodes, embed_client)
            result["images_embedded"] = embedded_count

        if progress_callback:
            progress_callback(90, "Extracting concepts")

        # --- Stage 7: Concept extraction ---
        chat_model = LLMModel.get_system_chat_model()
        if chat_model and chunk_result.chunks:
            document_text = "\n\n".join(
                block.text for block in parse_result.text_blocks
            )
            extractor = ConceptExtractor(chat_model, user=self.user)
            concepts_count = extractor.extract_for_item(
                item, chunk_nodes, chunk_result.chunks, document_text
            )
            result["concepts_extracted"] = concepts_count

        # Update content hash to prevent redundant re-processing
        if file_data:
            item.content_hash = hashlib.sha256(file_data).hexdigest()
            item.save()

        return result

    def _get_item_library(self, item):
        """
        Walk the graph to find the Library containing this Item.

        :param item: Item node.
        :returns: Library node, or None.
        """
        from library.models import Collection

        try:
            # Item <- Collection <- Library
            from neomodel import db

            results, _ = db.cypher_query(
                "MATCH (l:Library)-[:CONTAINS]->(c:Collection)-[:CONTAINS]->(i:Item {uid: $uid}) "
                "RETURN l.uid, l.library_type, l.chunking_config, l.embedding_instruction",
                {"uid": item.uid},
            )
            if results:
                from library.models import Library

                return Library.nodes.get(uid=results[0][0])
        except Exception as exc:
            logger.warning("Could not resolve library for item_uid=%s: %s", item.uid, exc)

        return None

    def _read_item_from_s3(self, item) -> Optional[bytes]:
        """
        Read the original file from S3 storage.

        :param item: Item node with s3_key.
        :returns: File bytes, or None.
        """
        if not item.s3_key:
            logger.warning("Item has no s3_key item_uid=%s", item.uid)
            return None

        try:
            with default_storage.open(item.s3_key, "rb") as f:
                data = f.read()
            logger.debug(
                "Read item from S3 key=%s size=%d", item.s3_key, len(data)
            )
            return data
        except Exception as exc:
            logger.error(
                "Failed to read from S3 key=%s: %s", item.s3_key, exc
            )
            raise

    def _store_chunks(self, item, chunk_result) -> list:
        """
        Store chunk text in S3 and create Chunk nodes in Neo4j.

        :param item: Item node.
        :param chunk_result: ChunkResult from chunker.
        :returns: List of Chunk node instances.
        """
        from library.models import Chunk

        # Delete existing chunks for this item
        for old_chunk in item.chunks.all():
            # Clean up S3
            try:
                default_storage.delete(old_chunk.chunk_s3_key)
            except Exception:
                pass
            old_chunk.delete()

        chunk_nodes = []
        for idx, chunk_text in enumerate(chunk_result.chunks):
            s3_key = CHUNK_S3_KEY.format(item_uid=item.uid, index=idx)

            # Store chunk text in S3
            try:
                default_storage.save(s3_key, ContentFile(chunk_text.encode("utf-8")))
            except Exception as exc:
                logger.error("Failed to store chunk %d to S3: %s", idx, exc)
                raise

            # Create Chunk node
            chunk_node = Chunk(
                chunk_index=idx,
                chunk_s3_key=s3_key,
                chunk_size=len(chunk_text),
                text_preview=chunk_text[:500],
            )
            chunk_node.save()
            item.chunks.connect(chunk_node)
            chunk_nodes.append(chunk_node)

        logger.info(
            "Stored %d chunks for item_uid=%s", len(chunk_nodes), item.uid
        )
        return chunk_nodes

    def _embed_chunks(
        self,
        item,
        chunk_nodes: list,
        chunk_texts: list[str],
        embed_client,
        embedding_instruction: str,
        progress_callback: Optional[callable],
    ):
        """
        Generate embeddings for chunks and update Chunk nodes.

        :param item: Item node.
        :param chunk_nodes: List of Chunk nodes.
        :param chunk_texts: List of chunk text strings.
        :param embed_client: EmbeddingClient instance.
        :param embedding_instruction: Instruction prefix for embedding.
        :param progress_callback: Optional progress callback.
        """
        # Prepend embedding instruction if configured
        if embedding_instruction:
            texts_to_embed = [
                f"{embedding_instruction}\n\n{text}" for text in chunk_texts
            ]
        else:
            texts_to_embed = chunk_texts

        batch_size = EMBEDDING_BATCH_SIZE
        total_batches = (len(texts_to_embed) + batch_size - 1) // batch_size

        for batch_idx in range(0, len(texts_to_embed), batch_size):
            batch_texts = texts_to_embed[batch_idx : batch_idx + batch_size]
            batch_nodes = chunk_nodes[batch_idx : batch_idx + batch_size]
            batch_num = batch_idx // batch_size + 1

            logger.debug(
                "Embedding batch %d/%d size=%d item_uid=%s",
                batch_num,
                total_batches,
                len(batch_texts),
                item.uid,
            )

            embeddings = embed_client.embed_texts(batch_texts)

            for node, embedding in zip(batch_nodes, embeddings):
                node.embedding = embedding
                node.save()

            if progress_callback:
                pct = 40 + (30 * (batch_idx + len(batch_texts)) / len(texts_to_embed))
                progress_callback(
                    int(pct),
                    f"Embedded {batch_idx + len(batch_texts)}/{len(texts_to_embed)} chunks",
                )

        logger.info(
            "Embedded %d chunks for item_uid=%s", len(chunk_nodes), item.uid
        )

    def _store_images(self, item, extracted_images) -> list:
        """
        Store extracted images in S3 and create Image nodes in Neo4j.

        :param item: Item node.
        :param extracted_images: List of ExtractedImage from parser.
        :returns: List of Image node instances.
        """
        from library.models import Image

        # Delete existing images for this item
        for old_image in item.images.all():
            try:
                default_storage.delete(old_image.s3_key)
            except Exception:
                pass
            old_image.delete()

        image_nodes = []
        for img in extracted_images:
            s3_key = IMAGE_S3_KEY.format(
                item_uid=item.uid,
                index=img.source_index,
                ext=img.ext,
            )

            try:
                default_storage.save(s3_key, ContentFile(img.data))
            except Exception as exc:
                logger.warning(
                    "Failed to store image %d to S3: %s", img.source_index, exc
                )
                continue

            image_node = Image(
                s3_key=s3_key,
                image_type="diagram",  # Default; could be refined by content analysis
                metadata={
                    "width": img.width,
                    "height": img.height,
                    "source_page": img.source_page,
                    "content_type": f"image/{img.ext}",
                },
            )
            image_node.save()
            item.images.connect(image_node)
            image_nodes.append(image_node)

        if image_nodes:
            logger.info(
                "Stored %d images for item_uid=%s", len(image_nodes), item.uid
            )

        return image_nodes

    def _associate_images_with_chunks(
        self, chunk_nodes, image_nodes, chunk_result, parse_result
    ):
        """
        Create HAS_NEARBY_IMAGE relationships between chunks and images.

        Associates images with chunks from the same page/section.

        :param chunk_nodes: List of Chunk nodes.
        :param image_nodes: List of Image nodes.
        :param chunk_result: ChunkResult with page mapping.
        :param parse_result: ParseResult with image source pages.
        """
        if not chunk_nodes or not image_nodes:
            return

        # Build page -> images mapping
        page_images: dict[int, list] = {}
        for img_node, ext_img in zip(image_nodes, parse_result.images):
            page_images.setdefault(ext_img.source_page, []).append(img_node)

        # Connect chunks to images on the same page
        connected = 0
        for chunk_idx, chunk_node in enumerate(chunk_nodes):
            page = chunk_result.chunk_page_map.get(chunk_idx, -1)
            nearby = page_images.get(page, [])
            for img_node in nearby:
                chunk_node.nearby_images.connect(
                    img_node, {"proximity": "same_page"}
                )
                connected += 1

        if connected:
            logger.debug(
                "Created %d chunk-image associations", connected
            )

    def _embed_images(self, image_nodes: list, embed_client) -> int:
        """
        Generate multimodal embeddings for Image nodes.

        :param image_nodes: List of Image nodes.
        :param embed_client: EmbeddingClient with multimodal support.
        :returns: Number of images successfully embedded.
        """
        from library.models import ImageEmbedding

        embedded_count = 0
        for img_node in image_nodes:
            try:
                img_data = default_storage.open(img_node.s3_key, "rb").read()
                ext = img_node.s3_key.rsplit(".", 1)[-1] if "." in img_node.s3_key else "png"

                embedding = embed_client.embed_image(img_data, ext)
                if embedding:
                    emb_node = ImageEmbedding(embedding=embedding)
                    emb_node.save()
                    img_node.embeddings.connect(emb_node)
                    embedded_count += 1
            except Exception as exc:
                logger.warning(
                    "Image embedding failed s3_key=%s: %s",
                    img_node.s3_key,
                    exc,
                )

        if embedded_count:
            logger.info("Embedded %d images", embedded_count)

        return embedded_count

    def _get_vision_prompt(self, library) -> str:
        """
        Get the content-type-aware vision prompt for a library.

        :param library: Library node, or None.
        :returns: Vision prompt string, or empty string.
        """
        if not library:
            return ""

        try:
            from library.content_types import get_library_type_config

            config = get_library_type_config(library.library_type)
            return config.get("vision_prompt", "")
        except Exception:
            return ""

    def _check_dimension_compatibility(self, model_dimensions: int):
        """
        Check if the model's vector dimensions match the Neo4j index.

        :param model_dimensions: Expected embedding dimensions.
        """
        # Log a warning — actual enforcement is in setup_neo4j_indexes
        logger.debug(
            "System embedding model dimensions=%d", model_dimensions
        )

    def reprocess_item(self, item_uid: str, progress_callback=None) -> dict:
        """
        Re-embed an item: delete existing chunks/images, then re-process.

        :param item_uid: UID of the Item to re-embed.
        :param progress_callback: Optional progress callback.
        :returns: Processing results dict.
        """
        from library.models import Item

        try:
            item = Item.nodes.get(uid=item_uid)
        except Exception:
            raise ValueError(f"Item not found: {item_uid}")

        # Clear content hash to force re-processing
        item.content_hash = ""
        item.save()

        logger.info("Re-processing item_uid=%s title='%s'", item_uid, item.title)
        return self.process_item(item_uid, progress_callback)
