import json
import logging
from abc import ABC, abstractmethod
from typing import Any, ClassVar, Optional, cast

from django.conf import settings
from django.core.files.base import ContentFile
from plasmapdf.models.PdfDataLayer import build_translation_layer

from opencontractserver.annotations.models import RELATIONSHIP_LABEL
from opencontractserver.documents.models import Document
from opencontractserver.pipeline.base.exceptions import DocumentParsingError
from opencontractserver.pipeline.base.file_types import FileTypeEnum
from opencontractserver.types.dicts import OpenContractDocExport
from opencontractserver.utils.compact_pawls import compact_pawls_pages
from opencontractserver.utils.importing import (
    import_annotations,
    import_relationships,
    load_or_create_labels,
)
from opencontractserver.utils.structural_sets import create_structural_annotation_set
from opencontractserver.utils.subtree_groups import build_subtree_groups_for_document

from .base_component import PipelineComponentBase

logger = logging.getLogger(__name__)


class BaseParser(PipelineComponentBase, ABC):
    """
    Abstract base parser class. Parsers should inherit from this class.
    Handles automatic loading of settings from Django settings.PIPELINE_SETTINGS.
    """

    supported_file_types: ClassVar[list[FileTypeEnum]] = []

    # Whether this parser can split paginated input into page-range chunks
    # (see BaseChunkedParser). Introspected by the ingestion orchestrator to
    # decide whether the chunked parse path applies. Non-paginated parsers
    # (TXT, DOCX) and remote parsers that self-batch leave this False.
    supports_chunking: ClassVar[bool] = False

    def __init__(self, **kwargs):
        """
        Initializes the Parser.
        Kwargs are passed to the superclass constructor (PipelineComponentBase).
        """
        super().__init__(**kwargs)

    @abstractmethod
    def _parse_document_impl(
        self, user_id: int, doc_id: int, **all_kwargs
    ) -> Optional[OpenContractDocExport]:
        """
        Abstract internal method to parse a document.
        Concrete subclasses must implement this method.

        Args:
            user_id (int): ID of the user.
            doc_id (int): ID of the document to parse.
            **all_kwargs: All keyword arguments, including those from
                          PIPELINE_SETTINGS and direct call-time arguments.

        Returns:
            Optional[OpenContractDocExport]: The parsed document data, or None if parsing failed.
        """
        pass

    def parse_document(
        self, user_id: int, doc_id: int, **direct_kwargs
    ) -> Optional[OpenContractDocExport]:
        """
        Parses a document, automatically injecting settings from PIPELINE_SETTINGS.

        Args:
            user_id (int): ID of the user.
            doc_id (int): ID of the document to parse.
            **direct_kwargs: Arbitrary keyword arguments that may be provided
                             for specific parser functionalities at call time.
                             These will override settings from PIPELINE_SETTINGS.

        Returns:
            Optional[OpenContractDocExport]: The parsed document data, or None if parsing failed.
        """
        # Merge component settings with direct kwargs, direct_kwargs take precedence
        merged_kwargs = {**self.get_component_settings(), **direct_kwargs}
        logger.info(
            f"Calling _parse_document_impl for doc_id {doc_id} with merged kwargs: {merged_kwargs}"
        )
        return self._parse_document_impl(user_id, doc_id, **merged_kwargs)

    def save_parsed_data(
        self,
        user_id: int,
        doc_id: int,
        open_contracts_data: OpenContractDocExport,
        corpus_id: Optional[int] = None,
        annotation_type: Optional[str] = None,
    ) -> None:
        """
        Saves the parsed data (both annotations and relationships) to the Document model.

        Args:
            user_id (int): ID of the user.
            doc_id (int): ID of the document.
            open_contracts_data (OpenContractDocExport): The parsed document data, including:
                - labelled_text (List[OpenContractsAnnotationPythonType]):
                  Annotation list to be imported.
                - relationships (Optional[List[OpenContractsRelationshipPythonType]]):
                  Relationship list to be imported.
                - page_count, pawls_file_content, content, etc.
            corpus_id (Optional[int]): ID of the corpus, if the document should be associated with one.
            annotation_type (Optional[str]): The fallback annotation_type (e.g., SPAN_LABEL or TOKEN_LABEL).
                If the annotation data doesn't specify an annotation_type, this one is used.
        """
        logger = logging.getLogger(__name__)
        logger.info(f"Saving parsed data for doc {doc_id}")

        document = Document.objects.get(pk=doc_id)

        # Get user for corpus operations
        from django.contrib.auth import get_user_model

        User = get_user_model()
        user = User.objects.get(pk=user_id)

        # Associate with corpus if provided
        if corpus_id:
            # Use Django's lazy-loading with string reference to avoid circular import
            from django.apps import apps

            Corpus = apps.get_model("corpuses", "Corpus")
            DocumentPath = apps.get_model("documents", "DocumentPath")
            corpus_obj = Corpus.objects.get(id=corpus_id)

            # Check if document is ALREADY in this corpus via DocumentPath
            # (e.g., created by import_content in process_documents_zip)
            existing_path = DocumentPath.objects.filter(
                document=document,
                corpus=corpus_obj,
                is_current=True,
                is_deleted=False,
            ).first()

            if existing_path:
                # Document already in corpus - don't create another copy
                logger.info(
                    f"Document {doc_id} already in corpus {corpus_obj.title} "
                    f"via DocumentPath {existing_path.id}, skipping add_document"
                )
            else:
                # Document not yet in corpus - use add_document for corpus isolation
                document, status, _ = corpus_obj.add_document(
                    document=document, user=user
                )
                logger.info(
                    f"Associated document with corpus: {corpus_obj.title} (status: {status})"
                )
        else:
            corpus_obj = None

        # Save content to txt_extract_file
        txt_content = open_contracts_data.get("content", "")
        txt_file = ContentFile(txt_content.encode("utf-8"))
        document.txt_extract_file.save(f"doc_{doc_id}.txt", txt_file)

        # Handle PAWLS content if any
        pawls_file_content = open_contracts_data.get("pawls_file_content")
        if pawls_file_content:
            compact_data = compact_pawls_pages(pawls_file_content)
            pawls_string = json.dumps(compact_data)
            pawls_file = ContentFile(pawls_string.encode("utf-8"))
            document.pawls_parse_file.save(f"doc_{doc_id}.pawls", pawls_file)

            # Create text layer from PAWLS tokens (use original v1 data)
            span_translation_layer = build_translation_layer(pawls_file_content)
            # Optionally overwrite txt_extract_file with text from PAWLS
            txt_file = ContentFile(span_translation_layer.doc_text.encode("utf-8"))
            document.txt_extract_file.save(f"doc_{doc_id}.txt", txt_file)
            document.page_count = len(pawls_file_content)
        else:
            # Handle cases without PAWLS content
            document.page_count = open_contracts_data.get("page_count", 1)

        document.save()

        # Determine fallback label type (for annotations) if annotation types aren't specified in data
        logger.info(
            f"Loading or creating labels for document {doc_id} with file type {document.file_type}"
        )

        if annotation_type is not None:
            target_label_type = annotation_type
        else:
            target_label_type = settings.ANNOTATION_LABELS.get(
                document.file_type, "SPAN_LABEL"
            )

        logger.info(f"Target label type for textual annotations: {target_label_type}")

        # 1) Build a data dict for text annotation labels
        text_labels_data_dict = {
            label_data["annotationLabel"]: {
                "label_type": target_label_type,
                "color": "grey",
                "description": "Parser Structural Label",
                "icon": "expand",
                "text": label_data["annotationLabel"],
                "creator_id": user_id,
                "read_only": True,
            }
            for label_data in open_contracts_data.get("labelled_text", [])
        }

        logger.info(f"Text label data dict: {text_labels_data_dict}")

        # 2) Create or load text labels
        existing_text_labels = load_or_create_labels(
            user_id=user_id,
            labelset_obj=None,  # No labelset in this context
            label_data_dict=text_labels_data_dict,
            existing_labels={},
        )

        logger.info(f"Existing text label lookup: {existing_text_labels}")

        # 3) Import annotations & store mapping of old annotation IDs to new DB IDs
        # Pass PAWLs data for pre-extracting image content (faster embeddings).
        # Cast PawlsPagePythonType list to dict-list since import_annotations'
        # signature predates the typed-dict surface.
        annotation_id_map = import_annotations(
            user_id=user_id,
            doc_obj=document,
            corpus_obj=corpus_obj,
            annotations_data=open_contracts_data.get("labelled_text", []),
            label_lookup=existing_text_labels,
            label_type=target_label_type,
            pawls_data=cast("list[dict[Any, Any]]", pawls_file_content),
        )

        # 4) If there are relationships, load/create relationship labels and then import
        relationship_data = open_contracts_data.get("relationships", [])
        if relationship_data:
            # Build label data dict for relationship labels
            relationship_label_texts = {
                rel["relationshipLabel"] for rel in relationship_data
            }
            relationship_label_data_dict = {
                label_text: {
                    "label_type": RELATIONSHIP_LABEL,  # Distinct from text annotation type
                    "color": "gray",
                    "description": "Parser Relationship Label",
                    "icon": "share-alt",
                    "text": label_text,
                    "creator_id": user_id,
                    "read_only": True,
                }
                for label_text in relationship_label_texts
            }

            existing_relationship_labels = load_or_create_labels(
                user_id=user_id,
                labelset_obj=None,  # No labelset in this context
                label_data_dict=relationship_label_data_dict,
                existing_labels={},
            )

            # Now import relationships
            import_relationships(
                user_id=user_id,
                doc_obj=document,
                corpus_obj=corpus_obj,
                relationships_data=relationship_data,
                label_lookup=existing_relationship_labels,
                annotation_id_map=annotation_id_map,
            )

        # 4.5) Materialise subtree-group relationships from the structural
        # annotation parent-child tree. Runs BEFORE the structural-set
        # migration so the new rows (document=document, structural=True)
        # are picked up and re-homed to the structural_set automatically
        # by _create_structural_annotation_set's existing filter.
        build_subtree_groups_for_document(document=document, user_id=user_id)

        # 5) Create StructuralAnnotationSet and migrate structural annotations to it
        self._create_structural_annotation_set(document, user)

        logger.info(
            f"Document {doc_id} parsed (with annotations & relationships) and saved successfully."
        )

    def _create_structural_annotation_set(self, document: Document, user) -> None:
        """
        Create a StructuralAnnotationSet for the document and migrate structural
        annotations and relationships to it.

        This ensures structural annotations are properly isolated per document and
        can be embedded using corpus-specific embedders. The migration logic lives
        in :func:`opencontractserver.utils.structural_sets.create_structural_annotation_set`
        so the worker-upload ingestion path produces an identical structural layer.

        Args:
            document: The Document object
            user: The user creating the set
        """
        create_structural_annotation_set(
            document,
            user,
            parser_name=self.title or self.__class__.__name__,
            parser_version="1.0",
        )

    def _run_enrichment_stage(
        self, user_id: int, doc_id: int, parsed_data: OpenContractDocExport
    ) -> OpenContractDocExport:
        """
        Run the configured enricher chain over freshly-parsed document data.

        Enrichers are resolved per the document's MIME type from
        ``PipelineSettings.preferred_enrichers`` and run as a chain (each
        receives the previous enricher's output). Enrichment is additive and
        non-fatal: any failure here — including failing to resolve the
        enricher list itself — is logged and the un-enriched ``parsed_data``
        is returned so the parsed document is still saved.

        Args:
            user_id: ID of the user the document is being ingested for.
            doc_id: ID of the document being ingested.
            parsed_data: The OpenContractDocExport produced by parse_document().

        Returns:
            OpenContractDocExport: The parsed data, enriched if any enrichers
            are configured for the document's file type.
        """
        try:
            from opencontractserver.documents.models import PipelineSettings
            from opencontractserver.pipeline.utils import run_enrichers

            file_type = Document.objects.values_list("file_type", flat=True).get(
                pk=doc_id
            )
            enricher_paths = PipelineSettings.get_instance().get_preferred_enrichers(
                file_type
            )
            if not enricher_paths:
                return parsed_data
            logger.info(
                f"Running {len(enricher_paths)} enricher(s) for document "
                f"{doc_id} (file_type={file_type})"
            )
            return run_enrichers(enricher_paths, user_id, doc_id, parsed_data)
        except Exception as e:
            # Defence-in-depth: even resolving the enricher list must not fail
            # ingestion. run_enrichers already isolates per-enricher failures.
            logger.warning(
                f"Enrichment stage skipped for document {doc_id}: {e}",
                exc_info=True,
            )
            return parsed_data

    def process_document(
        self, user_id: int, doc_id: int, **kwargs
    ) -> OpenContractDocExport:
        """
        Process a document by parsing it and then saving the parsed data.
        This method calls parse_document(...) and then save_parsed_data(...).

        Args:
            user_id (int): ID of the user.
            doc_id (int): ID of the document to process.
            **kwargs: Arbitrary keyword arguments that may be provided
                      for specific parser functionalities. Notably:
                      - corpus_id (Optional[int]): ID of corpus for corpus-specific embeddings

        Returns:
            OpenContractDocExport: The parsed document data.

        Raises:
            DocumentParsingError: If parsing fails. The exception's is_transient
                attribute indicates whether the error might succeed on retry.
        """
        # Extract corpus_id for save_parsed_data (not needed by parse_document)
        corpus_id = kwargs.pop("corpus_id", None)

        logger.info(
            f"Processing document {doc_id} with possible parser kwargs: {kwargs}"
            + (f" (corpus_id={corpus_id})" if corpus_id else "")
        )

        parsed_data = self.parse_document(user_id, doc_id, **kwargs)
        if parsed_data is None:
            raise DocumentParsingError(
                f"Parser {self.__class__.__name__} returned None for document {doc_id}",
                is_transient=True,  # Assume transient by default; subclasses override
            )

        # Enrichment stage: chainable OpenContractDocExport transforms run
        # between parsing and persistence. Additive and non-fatal.
        parsed_data = self._run_enrichment_stage(user_id, doc_id, parsed_data)

        self.save_parsed_data(user_id, doc_id, parsed_data, corpus_id=corpus_id)
        logger.info(f"Document {doc_id} processed successfully.")

        return parsed_data
