"""
NRC ADAMS Document Fetcher

Fetches documents from the NRC Agencywide Documents Access and Management System (ADAMS).
Supports NUREGs, Regulatory Guides, Standard Review Plans, and other NRC documents.

ADAMS Web API Documentation:
https://www.nrc.gov/site-help/developers.html
"""

import logging
import re
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path

import httpx
from tenacity import retry, stop_after_attempt, wait_exponential

logger = logging.getLogger(__name__)


class NRCDocumentType(Enum):
    """Types of NRC documents available in ADAMS."""

    NUREG = "NUREG"
    NUREG_CR = "NUREG/CR"
    REG_GUIDE = "Regulatory Guide"
    SRP = "Standard Review Plan"
    GENERIC_LETTER = "Generic Letter"
    INFORMATION_NOTICE = "Information Notice"
    INSPECTION_MANUAL = "Inspection Manual"
    OTHER = "Other"


@dataclass
class ADAMSDocument:
    """Represents a document from NRC ADAMS."""

    accession_number: str
    title: str
    document_date: datetime | None = None
    document_type: NRCDocumentType = NRCDocumentType.OTHER
    docket_number: str | None = None
    case_reference_number: str | None = None
    authors: list[str] = field(default_factory=list)
    page_count: int | None = None
    file_size: int | None = None
    mime_type: str = "application/pdf"
    download_url: str | None = None

    @property
    def doc_id(self) -> str:
        """Generate a standardized document ID."""
        return self.accession_number


class NRCAdamsFetcher:
    """
    Fetches documents from NRC ADAMS.

    Supports searching by:
    - Document type (NUREG, Regulatory Guide, etc.)
    - Docket number
    - Date range
    - Keyword search

    Example:
        fetcher = NRCAdamsFetcher()
        docs = await fetcher.search_documents(
            document_type=NRCDocumentType.NUREG,
            keywords=["criticality", "safety"]
        )
        for doc in docs:
            await fetcher.download_document(doc, output_dir)
    """

    # NRC document server base URL (more reliable than ADAMS web interface)
    BASE_URL = "https://www.nrc.gov/docs/"

    # Known NUREG documents relevant to NCS
    # Accession numbers verified against NRC document server (www.nrc.gov/docs/)
    # Note: NRC server availability is intermittent - some documents may return 503 errors
    NCS_RELEVANT_NUREGS = [
        "ML15176A258",  # NUREG-1520 Rev 2 - Standard Review Plan for Fuel Cycle Facilities (2.2 MB)
        "ML050250061",  # NUREG/CR-6698 - Guide for Validation of Nuclear Criticality Safety (2.3 MB)
    ]

    @staticmethod
    def get_document_url(accession_number: str) -> str:
        """
        Generate the direct download URL for an NRC document.

        The NRC document server uses a folder structure based on the first 6
        characters of the accession number.

        Args:
            accession_number: ADAMS accession number (e.g., "ML15176A258")

        Returns:
            Direct URL to the PDF document
        """
        # URL format: https://www.nrc.gov/docs/ML1517/ML15176A258.pdf
        folder = accession_number[:6]
        return f"https://www.nrc.gov/docs/{folder}/{accession_number}.pdf"

    def __init__(
        self,
        timeout: float = 30.0,
        max_retries: int = 3,
    ):
        """
        Initialize the NRC ADAMS fetcher.

        Args:
            timeout: Request timeout in seconds
            max_retries: Maximum number of retry attempts
        """
        self.timeout = timeout
        self.max_retries = max_retries
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> "NRCAdamsFetcher":
        """Async context manager entry."""
        self._client = httpx.AsyncClient(timeout=self.timeout)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        if self._client:
            await self._client.aclose()

    @property
    def client(self) -> httpx.AsyncClient:
        """Get HTTP client, creating if necessary."""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=self.timeout)
        return self._client

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
    )
    async def get_document_metadata(
        self, accession_number: str
    ) -> ADAMSDocument | None:
        """
        Fetch metadata for a document by its accession number.

        Uses a HEAD request to verify the document exists and get basic info.

        Args:
            accession_number: ADAMS accession number (e.g., "ML15176A258")

        Returns:
            ADAMSDocument with metadata, or None if not found
        """
        download_url = self.get_document_url(accession_number)

        try:
            # Use HEAD request to verify document exists and get size
            response = await self.client.head(download_url, follow_redirects=True)
            response.raise_for_status()

            # Extract file size from Content-Length header
            file_size = None
            if content_length := response.headers.get("content-length"):
                file_size = int(content_length)

            # Determine mime type
            content_type = response.headers.get("content-type", "application/pdf")
            mime_type = content_type.split(";")[0].strip()

            return ADAMSDocument(
                accession_number=accession_number,
                title=f"Document {accession_number}",
                file_size=file_size,
                mime_type=mime_type,
                download_url=download_url,
            )
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                logger.warning(f"Document {accession_number} not found at {download_url}")
            elif e.response.status_code == 503:
                logger.warning(
                    f"NRC server unavailable for {accession_number} (503). "
                    "Try again later or use ADAMS web interface."
                )
            else:
                logger.error(f"HTTP error fetching {accession_number}: {e}")
            return None
        except httpx.RequestError as e:
            logger.error(f"Request error fetching {accession_number}: {e}")
            raise

    async def search_documents(
        self,
        keywords: list[str] | None = None,
        document_type: NRCDocumentType | None = None,
        docket_number: str | None = None,
        date_start: datetime | None = None,
        date_end: datetime | None = None,
        limit: int = 100,
    ) -> list[ADAMSDocument]:
        """
        Search ADAMS for documents matching criteria.

        Args:
            keywords: Search keywords
            document_type: Type of document to filter by
            docket_number: Docket number to filter by
            date_start: Start of date range
            date_end: End of date range
            limit: Maximum number of results

        Returns:
            List of matching ADAMSDocument objects
        """
        # Build search query
        # ADAMS uses a specific query syntax
        query_parts = []

        if keywords:
            query_parts.append(" AND ".join(keywords))

        if document_type:
            query_parts.append(f'DocumentType:"{document_type.value}"')

        if docket_number:
            query_parts.append(f'DocketNumber:"{docket_number}"')

        if date_start:
            query_parts.append(f"DocumentDate>={date_start.strftime('%m/%d/%Y')}")

        if date_end:
            query_parts.append(f"DocumentDate<={date_end.strftime('%m/%d/%Y')}")

        query = " AND ".join(query_parts) if query_parts else "*"

        logger.info(f"Searching ADAMS with query: {query}")

        # In production, this would make the actual ADAMS search API call
        # For now, return known NCS-relevant documents
        documents = []
        for acc_num in self.NCS_RELEVANT_NUREGS[:limit]:
            doc = await self.get_document_metadata(acc_num)
            if doc:
                documents.append(doc)

        return documents

    async def download_document(
        self,
        document: ADAMSDocument,
        output_dir: Path,
        filename: str | None = None,
    ) -> Path | None:
        """
        Download a document from ADAMS.

        Args:
            document: ADAMSDocument to download
            output_dir: Directory to save the file
            filename: Optional custom filename

        Returns:
            Path to downloaded file, or None if download failed
        """
        if not document.download_url:
            logger.error(f"No download URL for document {document.accession_number}")
            return None

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        if filename is None:
            # Default filename based on accession number
            extension = ".pdf" if "pdf" in document.mime_type else ".txt"
            filename = f"{document.accession_number}{extension}"

        output_path = output_dir / filename

        try:
            logger.info(f"Downloading {document.accession_number} to {output_path}")

            async with self.client.stream("GET", document.download_url) as response:
                response.raise_for_status()

                with open(output_path, "wb") as f:
                    async for chunk in response.aiter_bytes(chunk_size=8192):
                        f.write(chunk)

            logger.info(f"Successfully downloaded {document.accession_number}")
            return output_path

        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error downloading {document.accession_number}: {e}")
            return None
        except httpx.RequestError as e:
            logger.error(f"Request error downloading {document.accession_number}: {e}")
            raise
        except OSError as e:
            logger.error(f"IO error saving {document.accession_number}: {e}")
            return None

    async def fetch_ncs_documents(
        self,
        output_dir: Path,
        include_types: list[NRCDocumentType] | None = None,
    ) -> list[Path]:
        """
        Fetch all NCS-relevant documents.

        This is a convenience method that downloads known NCS-relevant
        documents from ADAMS.

        Args:
            output_dir: Directory to save documents
            include_types: Document types to include (default: all)

        Returns:
            List of paths to downloaded documents
        """
        downloaded = []

        # Fetch known NCS-relevant NUREGs
        for acc_num in self.NCS_RELEVANT_NUREGS:
            doc = await self.get_document_metadata(acc_num)
            if doc:
                if include_types is None or doc.document_type in include_types:
                    path = await self.download_document(doc, output_dir)
                    if path:
                        downloaded.append(path)

        return downloaded


def parse_nureg_id(text: str) -> tuple[str, str | None] | None:
    """
    Parse NUREG identifier from text.

    Args:
        text: Text potentially containing NUREG ID

    Returns:
        Tuple of (nureg_number, revision) or None
    """
    # Pattern for NUREG-XXXX or NUREG/CR-XXXX
    pattern = r"NUREG(?:/CR)?-(\d+)(?:\s*Rev(?:ision)?\.?\s*(\d+))?"

    match = re.search(pattern, text, re.IGNORECASE)
    if match:
        nureg_num = match.group(1)
        revision = match.group(2)
        return (nureg_num, revision)

    return None
