"""
OSTI.gov Document Fetcher

Fetches documents from the Office of Scientific and Technical Information (OSTI).
Supports LA-series reports, ORNL/TM reports, and other DOE technical documents.

OSTI API Documentation:
https://www.osti.gov/api/v1/docs
"""

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

import httpx
from tenacity import retry, stop_after_attempt, wait_exponential

logger = logging.getLogger(__name__)


class OSTIDocumentType(Enum):
    """Types of documents available in OSTI."""

    TECHNICAL_REPORT = "Technical Report"
    JOURNAL_ARTICLE = "Journal Article"
    CONFERENCE_PAPER = "Conference"
    THESIS = "Thesis/Dissertation"
    BOOK = "Book"
    SOFTWARE = "Software"
    PATENT = "Patent"
    OTHER = "Other"


@dataclass
class OSTIDocument:
    """Represents a document from OSTI."""

    osti_id: str
    title: str
    authors: list[str] = field(default_factory=list)
    publication_date: datetime | None = None
    document_type: OSTIDocumentType = OSTIDocumentType.TECHNICAL_REPORT
    report_number: str | None = None  # e.g., LA-10860-MS
    doi: str | None = None
    abstract: str | None = None
    subjects: list[str] = field(default_factory=list)
    sponsoring_org: str | None = None
    research_org: str | None = None
    country: str = "US"
    language: str = "English"
    availability: str | None = None
    download_url: str | None = None
    full_text_url: str | None = None

    @property
    def doc_id(self) -> str:
        """Generate standardized document ID."""
        return self.report_number or self.osti_id


class OSTIClient:
    """
    Client for OSTI.gov API.

    Provides access to DOE scientific and technical information including:
    - National laboratory reports (LA, ORNL, SAND, etc.)
    - DOE-sponsored research
    - Technical standards

    Example:
        client = OSTIClient()
        docs = await client.search(
            keywords=["nuclear criticality", "safety"],
            research_org="Los Alamos National Laboratory"
        )
        for doc in docs:
            await client.download_document(doc, output_dir)
    """

    BASE_URL = "https://www.osti.gov/api/v1"

    # Known NCS-relevant report numbers
    NCS_RELEVANT_REPORTS = [
        "LA-10860-MS",  # Critical Dimensions of Systems Containing U-235, Pu-239, and U-233
        "LA-12808",  # Nuclear Criticality Safety Guide
        "ORNL/TM-12292",  # SCALE 4.3 - Modular Code System
        "SAND2017-6962",  # MCNP 6.2 User Manual
    ]

    def __init__(
        self,
        api_key: str | None = None,
        timeout: float = 30.0,
    ):
        """
        Initialize OSTI client.

        Args:
            api_key: Optional OSTI API key (increases rate limits)
            timeout: Request timeout in seconds
        """
        self.api_key = api_key
        self.timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> "OSTIClient":
        """Async context manager entry."""
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        self._client = httpx.AsyncClient(
            timeout=self.timeout,
            headers=headers,
        )
        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."""
        if self._client is None:
            headers = {}
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"
            self._client = httpx.AsyncClient(timeout=self.timeout, headers=headers)
        return self._client

    def _parse_document(self, data: dict[str, Any]) -> OSTIDocument:
        """Parse OSTI API response into OSTIDocument."""
        # Parse publication date
        pub_date = None
        if date_str := data.get("publication_date"):
            try:
                pub_date = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
            except ValueError:
                pass

        # Parse document type
        doc_type = OSTIDocumentType.OTHER
        if type_str := data.get("product_type"):
            for dt in OSTIDocumentType:
                if dt.value.lower() in type_str.lower():
                    doc_type = dt
                    break

        # Get download URL
        download_url = None
        full_text_url = None
        if links := data.get("links"):
            for link in links:
                if link.get("rel") == "fulltext":
                    full_text_url = link.get("href")
                elif link.get("rel") == "citation":
                    download_url = link.get("href")

        return OSTIDocument(
            osti_id=str(data.get("osti_id", "")),
            title=data.get("title", ""),
            authors=data.get("authors", []),
            publication_date=pub_date,
            document_type=doc_type,
            report_number=data.get("report_number"),
            doi=data.get("doi"),
            abstract=data.get("description"),
            subjects=data.get("subjects", []),
            sponsoring_org=data.get("sponsor_org"),
            research_org=data.get("research_org"),
            country=data.get("country", "US"),
            language=data.get("language", "English"),
            availability=data.get("availability"),
            download_url=download_url,
            full_text_url=full_text_url,
        )

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
    )
    async def get_document(self, osti_id: str) -> OSTIDocument | None:
        """
        Get document by OSTI ID.

        Args:
            osti_id: OSTI document ID

        Returns:
            OSTIDocument or None if not found
        """
        url = f"{self.BASE_URL}/records/{osti_id}"

        try:
            response = await self.client.get(url)
            response.raise_for_status()
            data = response.json()
            return self._parse_document(data)
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                logger.warning(f"Document {osti_id} not found")
                return None
            logger.error(f"HTTP error fetching {osti_id}: {e}")
            raise
        except httpx.RequestError as e:
            logger.error(f"Request error fetching {osti_id}: {e}")
            raise

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
    )
    async def search(
        self,
        keywords: list[str] | None = None,
        title: str | None = None,
        author: str | None = None,
        report_number: str | None = None,
        research_org: str | None = None,
        sponsoring_org: str | None = None,
        subject: str | None = None,
        date_start: datetime | None = None,
        date_end: datetime | None = None,
        document_type: OSTIDocumentType | None = None,
        page: int = 0,
        page_size: int = 25,
    ) -> list[OSTIDocument]:
        """
        Search OSTI for documents.

        Args:
            keywords: Full-text search keywords
            title: Title search
            author: Author name
            report_number: Report number (e.g., LA-10860-MS)
            research_org: Research organization
            sponsoring_org: Sponsoring organization
            subject: Subject/topic
            date_start: Start of publication date range
            date_end: End of publication date range
            document_type: Document type filter
            page: Page number (0-indexed)
            page_size: Results per page (max 100)

        Returns:
            List of matching OSTIDocument objects
        """
        params: dict[str, Any] = {
            "page": page,
            "rows": min(page_size, 100),
        }

        # Build search query
        if keywords:
            params["q"] = " ".join(keywords)

        if title:
            params["title"] = title

        if author:
            params["author"] = author

        if report_number:
            params["report_number"] = report_number

        if research_org:
            params["research_org"] = research_org

        if sponsoring_org:
            params["sponsor_org"] = sponsoring_org

        if subject:
            params["subject"] = subject

        if date_start:
            params["publication_date_start"] = date_start.strftime("%Y-%m-%d")

        if date_end:
            params["publication_date_end"] = date_end.strftime("%Y-%m-%d")

        if document_type:
            params["product_type"] = document_type.value

        logger.info(f"Searching OSTI with params: {params}")

        url = f"{self.BASE_URL}/records"

        try:
            response = await self.client.get(url, params=params)
            response.raise_for_status()
            data = response.json()

            # OSTI API returns a list directly, not {"records": [...]}
            records = data if isinstance(data, list) else data.get("records", [])

            documents = []
            for record in records:
                documents.append(self._parse_document(record))

            logger.info(f"Found {len(documents)} documents")
            return documents

        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error searching OSTI: {e}")
            raise
        except httpx.RequestError as e:
            logger.error(f"Request error searching OSTI: {e}")
            raise

    async def search_by_report_number(
        self, report_number: str
    ) -> OSTIDocument | None:
        """
        Search for a document by its report number.

        Args:
            report_number: Report number (e.g., LA-10860-MS)

        Returns:
            OSTIDocument or None if not found
        """
        docs = await self.search(report_number=report_number, page_size=1)
        return docs[0] if docs else None

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

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

        Returns:
            Path to downloaded file, or None if download failed
        """
        download_url = document.full_text_url or document.download_url

        if not download_url:
            logger.error(f"No download URL for document {document.osti_id}")
            return None

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

        if filename is None:
            # Use report number or OSTI ID for filename
            base_name = document.report_number or document.osti_id
            # Sanitize filename
            base_name = base_name.replace("/", "_").replace("\\", "_")
            filename = f"{base_name}.pdf"

        output_path = output_dir / filename

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

            async with self.client.stream("GET", 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.doc_id}")
            return output_path

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

    async def fetch_ncs_documents(
        self,
        output_dir: Path,
    ) -> list[Path]:
        """
        Fetch known NCS-relevant documents.

        Args:
            output_dir: Directory to save documents

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

        for report_num in self.NCS_RELEVANT_REPORTS:
            doc = await self.search_by_report_number(report_num)
            if doc:
                path = await self.download_document(doc, output_dir)
                if path:
                    downloaded.append(path)
            else:
                logger.warning(f"Could not find report {report_num}")

        return downloaded
