"""
Citation Tracking System

Maintains provenance and citation information for all chunks,
enabling accurate source attribution and cross-referencing.
"""

import logging
import uuid
from dataclasses import dataclass
from datetime import datetime
from enum import Enum

logger = logging.getLogger(__name__)


class CitationType(Enum):
    """Type of citation reference."""

    STANDARD = "standard"  # ANSI/ANS standards
    REGULATION = "regulation"  # NRC regulations
    REPORT = "report"  # Technical reports
    JOURNAL = "journal"  # Journal articles
    HANDBOOK = "handbook"  # Reference handbooks
    OTHER = "other"


@dataclass
class ChunkCitation:
    """
    Citation information for a document chunk.

    Provides complete provenance for any piece of content,
    enabling accurate source attribution in RAG responses.
    """

    # Document identity
    doc_id: str  # Human-readable ID (e.g., "ANSI/ANS-8.1-2014")
    document_uuid: str = ""  # Internal UUID

    # Location in document
    section: str = ""  # e.g., "Table 1", "Section 2.3"
    section_path: str = ""  # Full path: "Chapter 2 > Section 2.1 > Table 1"
    page: int | None = None
    paragraph: int | None = None

    # Content type
    citation_type: CitationType = CitationType.OTHER
    has_table: bool = False
    has_equation: bool = False
    has_figure: bool = False

    # Quality indicators
    confidence: float = 1.0  # Extraction confidence (0.0 to 1.0)
    ocr_quality: float | None = None  # OCR confidence if applicable

    # Timestamps
    extraction_date: datetime | None = None
    last_verified: datetime | None = None

    def __post_init__(self):
        if not self.extraction_date:
            self.extraction_date = datetime.now()

    @property
    def short_citation(self) -> str:
        """Generate short citation string."""
        parts = [self.doc_id]
        if self.section:
            parts.append(self.section)
        return ", ".join(parts)

    @property
    def full_citation(self) -> str:
        """Generate full citation string."""
        parts = [self.doc_id]
        if self.section_path:
            parts.append(self.section_path)
        elif self.section:
            parts.append(self.section)
        if self.page:
            parts.append(f"p. {self.page}")
        return ", ".join(parts)

    def to_dict(self) -> dict:
        """Convert to dictionary for storage."""
        return {
            "doc_id": self.doc_id,
            "document_uuid": self.document_uuid,
            "section": self.section,
            "section_path": self.section_path,
            "page": self.page,
            "paragraph": self.paragraph,
            "citation_type": self.citation_type.value,
            "has_table": self.has_table,
            "has_equation": self.has_equation,
            "has_figure": self.has_figure,
            "confidence": self.confidence,
            "ocr_quality": self.ocr_quality,
            "extraction_date": self.extraction_date.isoformat() if self.extraction_date else None,
            "last_verified": self.last_verified.isoformat() if self.last_verified else None,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "ChunkCitation":
        """Create from dictionary."""
        citation_type = CitationType(data.get("citation_type", "other"))

        extraction_date = None
        if data.get("extraction_date"):
            extraction_date = datetime.fromisoformat(data["extraction_date"])

        last_verified = None
        if data.get("last_verified"):
            last_verified = datetime.fromisoformat(data["last_verified"])

        return cls(
            doc_id=data.get("doc_id", ""),
            document_uuid=data.get("document_uuid", ""),
            section=data.get("section", ""),
            section_path=data.get("section_path", ""),
            page=data.get("page"),
            paragraph=data.get("paragraph"),
            citation_type=citation_type,
            has_table=data.get("has_table", False),
            has_equation=data.get("has_equation", False),
            has_figure=data.get("has_figure", False),
            confidence=data.get("confidence", 1.0),
            ocr_quality=data.get("ocr_quality"),
            extraction_date=extraction_date,
            last_verified=last_verified,
        )


@dataclass
class CrossReference:
    """
    Cross-reference between documents.

    Tracks when one document cites or references another.
    """

    id: str = ""
    source_doc_id: str = ""  # Document making the reference
    source_document_uuid: str = ""
    target_doc_id: str = ""  # Document being referenced (may not be in corpus)
    citation_text: str = ""  # Actual citation text
    citation_type: CitationType = CitationType.OTHER
    section_ref: str = ""  # Section in source where reference appears

    def __post_init__(self):
        if not self.id:
            self.id = str(uuid.uuid4())

    def to_dict(self) -> dict:
        """Convert to dictionary for storage."""
        return {
            "id": self.id,
            "source_doc_id": self.source_doc_id,
            "source_document_uuid": self.source_document_uuid,
            "target_doc_id": self.target_doc_id,
            "citation_text": self.citation_text,
            "citation_type": self.citation_type.value,
            "section_ref": self.section_ref,
        }


class CitationTracker:
    """
    Tracks citations and cross-references for the SAGE knowledge base.

    Maintains a registry of all citations and enables:
    - Citation lookup for any chunk
    - Cross-reference tracking between documents
    - Citation validation and verification

    Example:
        tracker = CitationTracker()

        # Create citation for a chunk
        citation = tracker.create_citation(
            doc_id="ANSI/ANS-8.1-2014",
            section="Table 1",
            page=12
        )

        # Track cross-reference
        tracker.add_cross_reference(
            source_doc_id="NUREG-1520",
            target_doc_id="ANSI/ANS-8.1-2014",
            citation_text="See ANSI/ANS-8.1, Table 1"
        )
    """

    def __init__(self) -> None:
        """Initialize citation tracker."""
        self._citations: dict[str, ChunkCitation] = {}
        self._cross_references: list[CrossReference] = []
        self._doc_citations: dict[str, list[str]] = {}  # doc_id -> citation_ids

    def create_citation(
        self,
        doc_id: str,
        section: str = "",
        page: int | None = None,
        **kwargs,
    ) -> ChunkCitation:
        """
        Create a new citation.

        Args:
            doc_id: Document identifier
            section: Section reference
            page: Page number
            **kwargs: Additional citation attributes

        Returns:
            ChunkCitation object
        """
        citation = ChunkCitation(
            doc_id=doc_id,
            section=section,
            page=page,
            **kwargs,
        )

        citation_id = str(uuid.uuid4())
        self._citations[citation_id] = citation

        # Index by document
        if doc_id not in self._doc_citations:
            self._doc_citations[doc_id] = []
        self._doc_citations[doc_id].append(citation_id)

        return citation

    def get_citation(self, citation_id: str) -> ChunkCitation | None:
        """Get citation by ID."""
        return self._citations.get(citation_id)

    def get_document_citations(self, doc_id: str) -> list[ChunkCitation]:
        """Get all citations for a document."""
        citation_ids = self._doc_citations.get(doc_id, [])
        return [self._citations[cid] for cid in citation_ids if cid in self._citations]

    def add_cross_reference(
        self,
        source_doc_id: str,
        target_doc_id: str,
        citation_text: str = "",
        citation_type: CitationType = CitationType.OTHER,
        section_ref: str = "",
        source_document_uuid: str = "",
    ) -> CrossReference:
        """
        Add a cross-reference between documents.

        Args:
            source_doc_id: Document making the reference
            target_doc_id: Document being referenced
            citation_text: Actual citation text
            citation_type: Type of citation
            section_ref: Section where reference appears
            source_document_uuid: Internal UUID of source

        Returns:
            CrossReference object
        """
        cross_ref = CrossReference(
            source_doc_id=source_doc_id,
            source_document_uuid=source_document_uuid,
            target_doc_id=target_doc_id,
            citation_text=citation_text,
            citation_type=citation_type,
            section_ref=section_ref,
        )

        self._cross_references.append(cross_ref)
        return cross_ref

    def get_references_from(self, doc_id: str) -> list[CrossReference]:
        """Get all references made by a document."""
        return [
            ref for ref in self._cross_references
            if ref.source_doc_id == doc_id
        ]

    def get_references_to(self, doc_id: str) -> list[CrossReference]:
        """Get all references to a document."""
        return [
            ref for ref in self._cross_references
            if ref.target_doc_id == doc_id
        ]

    def format_citations(
        self,
        citations: list[ChunkCitation],
        style: str = "short",
    ) -> str:
        """
        Format citations for display.

        Args:
            citations: List of citations to format
            style: "short" or "full"

        Returns:
            Formatted citation string
        """
        if not citations:
            return ""

        if style == "full":
            formatted = [c.full_citation for c in citations]
        else:
            formatted = [c.short_citation for c in citations]

        # Remove duplicates while preserving order
        seen = set()
        unique = []
        for c in formatted:
            if c not in seen:
                seen.add(c)
                unique.append(c)

        return "; ".join(unique)

    def validate_citation(self, citation: ChunkCitation) -> list[str]:
        """
        Validate a citation for completeness.

        Args:
            citation: Citation to validate

        Returns:
            List of validation issues (empty if valid)
        """
        issues = []

        if not citation.doc_id:
            issues.append("Missing document ID")

        if citation.confidence < 0.7:
            issues.append(f"Low extraction confidence ({citation.confidence:.1%})")

        if citation.ocr_quality is not None and citation.ocr_quality < 0.7:
            issues.append(f"Low OCR quality ({citation.ocr_quality:.1%})")

        if not citation.section and not citation.section_path:
            issues.append("No section reference provided")

        return issues

    def export_citations(self) -> dict:
        """Export all citations and cross-references."""
        return {
            "citations": {
                cid: c.to_dict()
                for cid, c in self._citations.items()
            },
            "cross_references": [
                ref.to_dict() for ref in self._cross_references
            ],
        }

    def import_citations(self, data: dict) -> None:
        """Import citations and cross-references."""
        for cid, cdata in data.get("citations", {}).items():
            self._citations[cid] = ChunkCitation.from_dict(cdata)

            # Rebuild document index
            doc_id = cdata.get("doc_id", "")
            if doc_id not in self._doc_citations:
                self._doc_citations[doc_id] = []
            self._doc_citations[doc_id].append(cid)

        for ref_data in data.get("cross_references", []):
            self._cross_references.append(CrossReference(
                id=ref_data.get("id", str(uuid.uuid4())),
                source_doc_id=ref_data.get("source_doc_id", ""),
                source_document_uuid=ref_data.get("source_document_uuid", ""),
                target_doc_id=ref_data.get("target_doc_id", ""),
                citation_text=ref_data.get("citation_text", ""),
                citation_type=CitationType(ref_data.get("citation_type", "other")),
                section_ref=ref_data.get("section_ref", ""),
            ))

    def get_stats(self) -> dict:
        """Get citation statistics."""
        unique_docs = set()
        for citation in self._citations.values():
            unique_docs.add(citation.doc_id)

        return {
            "total_citations": len(self._citations),
            "unique_documents": len(unique_docs),
            "cross_references": len(self._cross_references),
            "documents_with_citations": len(self._doc_citations),
        }
