"""Fetch AI-related publications from NIST CSRC.

NIST has no formal API but the publications search returns structured HTML.
We scrape the search results and supplement with a curated catalog of key
AI publications (AI RMF, AI 100-series, etc.).
"""

from __future__ import annotations

import re
import time as time_mod

import httpx
from rich.console import Console

from ..config import Config
from .base import SourceDocument

console = Console()

NIST_SEARCH_URL = "https://csrc.nist.gov/publications/search"
NIST_PUB_BASE = "https://csrc.nist.gov/publications/detail/"

# Curated catalog of key NIST AI publications
NIST_AI_CATALOG = [
    # AI 100 series
    ("AI 100-1", "Artificial Intelligence Risk Management Framework (AI RMF 1.0)",
     "The AI RMF provides a framework for managing risks associated with AI systems throughout their lifecycle, "
     "addressing characteristics of trustworthy AI including validity, reliability, safety, security, resilience, "
     "accountability, transparency, explainability, interpretability, privacy, and fairness.",
     "https://csrc.nist.gov/pubs/ai/100/1/final", "2023-01-26", "Final"),
    ("AI 100-2 E2025", "Adversarial Machine Learning: A Taxonomy and Terminology of Attacks and Mitigations",
     "Establishes a taxonomy of adversarial machine learning attacks and mitigations, covering evasion, poisoning, "
     "and privacy attacks against predictive AI and generative AI systems including LLMs.",
     "https://csrc.nist.gov/pubs/ai/100/2/e2025/final", "2025-03-24", "Final"),
    ("AI 100-4", "Reducing Risks Posed by Synthetic Content",
     "Provides guidance on approaches to manage risks from synthetic content generated by AI, "
     "including authentication, provenance tracking, and content labeling.",
     "https://csrc.nist.gov/pubs/ai/100/4/final", "2024-07-26", "Final"),
    ("AI 600-1", "Artificial Intelligence Risk Management Framework: Generative AI Profile",
     "Companion resource for the AI RMF focusing on risks unique to or exacerbated by generative AI, "
     "covering hallucinations, data privacy, intellectual property, and CBRN information risks.",
     "https://csrc.nist.gov/pubs/ai/600/1/final", "2024-07-26", "Final"),
    # Special Publications
    ("SP 800-218A", "Secure Software Development Practices for Generative AI and Dual-Use Foundation Models",
     "An SSDF community profile providing secure software development practices specific to generative AI "
     "and dual-use foundation models, addressing data management, model training, and deployment security.",
     "https://csrc.nist.gov/pubs/sp/800/218/a/final", "2024-07-26", "Final"),
    ("IR 8596", "Cybersecurity Framework Profile for Artificial Intelligence",
     "A community profile mapping NIST Cybersecurity Framework functions to AI-specific risks and controls, "
     "helping organizations manage cybersecurity risks in AI systems.",
     "https://csrc.nist.gov/pubs/ir/8596/ipd", "2025-12-16", "Draft"),
    ("CSWP 31", "Proxy Validation and Verification for Critical AI Systems",
     "Describes a proxy design process for validation and verification of critical AI systems, "
     "addressing challenges of testing AI systems that are complex and opaque.",
     "https://csrc.nist.gov/pubs/cswp/31/final", "2024-09-26", "Final"),
    ("IR 8579", "Developing the NCCoE Chatbot: Technical and Security Learnings",
     "Documents technical and security lessons from implementing an LLM-based chatbot at the "
     "National Cybersecurity Center of Excellence, including prompt injection mitigations.",
     "https://csrc.nist.gov/pubs/ir/8579/ipd", "2025-07-31", "Draft"),
    # Concept papers and other
    ("NIST-AI-Agent-IAM", "Accelerating the Adoption of Software and AI Agent Identity and Authorization",
     "Concept paper on identity and authorization for AI agents, addressing how software agents "
     "authenticate, authorize, and maintain accountability in automated systems.",
     "https://csrc.nist.gov/publications/detail/other/2026/02/05/accelerating-the-adoption-of-software-and-artificial-intelligence-agent-identity-and-authorization/draft",
     "2026-02-05", "Draft"),
    # Trustworthy AI
    ("AI 100-6", "AI Measurement and Evaluation",
     "Guidance on measurement and evaluation approaches for AI systems to assess trustworthiness "
     "characteristics including accuracy, fairness, robustness, and security.",
     "https://csrc.nist.gov/pubs/ai/100/6/final", "2024-04-29", "Final"),
    # AI RMF playbook
    ("AI 100-1 Playbook", "AI RMF Playbook",
     "Practical companion to the AI RMF providing suggested actions for achieving AI risk management outcomes, "
     "organized by the framework's functions: Govern, Map, Measure, and Manage.",
     "https://airc.nist.gov/AI_RMF_Playbook", "2023-01-26", "Final"),
]


def _nist_id_to_name(nist_id: str) -> str:
    """Convert NIST ID to slug. E.g. 'AI 100-1' -> 'nist-ai-100-1'."""
    slug = nist_id.lower().replace(" ", "-").replace("/", "-").replace(".", "-")
    return f"nist-{slug}"


class NISTFetcher:
    """Fetch AI-related publications from NIST CSRC.

    Combines a curated catalog with search scraping for discovery.
    """

    def __init__(self, config: Config | None = None):
        self.config = config or Config.load()
        self.client = httpx.Client(timeout=30, follow_redirects=True)

    def search(
        self, keywords: list[str], since: str | None = None
    ) -> list[SourceDocument]:
        """Return AI-relevant NIST publications."""
        seen: dict[str, SourceDocument] = {}

        # Strategy 1: Curated catalog
        console.print("  Loading NIST AI publication catalog...")
        for nist_id, title, abstract, url, date, status in NIST_AI_CATALOG:
            if since and date and date < since:
                continue
            name = _nist_id_to_name(nist_id)
            seen[name] = SourceDocument(
                name=name,
                title=f"NIST {nist_id}: {title}",
                abstract=abstract,
                source="nist",
                source_id=nist_id,
                source_url=url,
                time=date,
                doc_status=status.lower(),
            )

        # Strategy 2: Search CSRC for additional AI publications
        console.print("  Searching NIST CSRC for AI publications...")
        search_terms = ["artificial intelligence", "machine learning", "large language model"]
        for term in search_terms:
            new_docs = self._search_csrc(term, since)
            for doc in new_docs:
                if doc.name not in seen:
                    seen[doc.name] = doc
            time_mod.sleep(0.5)

        console.print(f"  Found [bold green]{len(seen)}[/] NIST publications")
        return list(seen.values())

    def _search_csrc(self, keyword: str, since: str | None) -> list[SourceDocument]:
        """Search NIST CSRC publications page."""
        docs = []
        try:
            resp = self.client.get(
                NIST_SEARCH_URL,
                params={
                    "keywords": keyword,
                    "status": "Final,Draft",
                    "sortBy": "relevance",
                },
            )
            if resp.status_code != 200:
                return docs

            # Parse search results HTML
            # Results are in structured divs with title, series, date, abstract
            # Pattern: <a href="/publications/detail/...">TITLE</a>
            entries = re.findall(
                r'<a[^>]*href="(/publications/detail/[^"]+)"[^>]*>([^<]+)</a>',
                resp.text,
            )

            for href, title in entries:
                title = title.strip()
                if not title or len(title) < 10:
                    continue

                # Extract a usable ID from the URL path
                parts = href.rstrip("/").split("/")
                # e.g. /publications/detail/sp/800/218/a/final -> sp-800-218-a
                slug_parts = [p for p in parts[3:] if p not in ("final", "draft", "ipd", "fpd")]
                nist_id = "-".join(slug_parts).upper()
                name = f"nist-{'-'.join(slug_parts).lower()}"

                if not name or name == "nist-":
                    continue

                docs.append(SourceDocument(
                    name=name,
                    title=title,
                    abstract=title,  # Will be enriched later if needed
                    source="nist",
                    source_id=nist_id,
                    source_url=f"https://csrc.nist.gov{href}",
                    time="",
                    doc_status="published",
                ))

        except httpx.HTTPError as e:
            console.print(f"[yellow]NIST search error: {e}[/]")
        return docs

    def download_text(self, doc: SourceDocument) -> str | None:
        """NIST publications are free PDFs — try to fetch and extract."""
        url = doc.source_url
        if not url:
            return None
        try:
            # First get the publication page to find the PDF link
            resp = self.client.get(url)
            if resp.status_code != 200:
                return None

            # Look for PDF download link
            pdf_match = re.search(r'href="([^"]+\.pdf)"', resp.text)
            if not pdf_match:
                # Extract abstract from page instead
                abstract_match = re.search(
                    r'(?:Abstract|Summary|Description)[:\s]*</[^>]+>\s*<[^>]+>(.+?)</(?:p|div)',
                    resp.text, re.DOTALL | re.IGNORECASE,
                )
                if abstract_match:
                    text = re.sub(r'<[^>]+>', '', abstract_match.group(1)).strip()
                    return text[:5000] if text else None
                return None

            pdf_url = pdf_match.group(1)
            if not pdf_url.startswith("http"):
                pdf_url = f"https://csrc.nist.gov{pdf_url}"

            resp = self.client.get(pdf_url)
            resp.raise_for_status()
            try:
                from io import BytesIO
                from pdfminer.high_level import extract_text
                text = extract_text(BytesIO(resp.content))
                return text[:100000] if text else None
            except ImportError:
                return f"[PDF document: {doc.title}. Install pdfminer.six to extract text.]"
        except httpx.HTTPError as e:
            console.print(f"[dim]Could not download {doc.name}: {e}[/]")
        return None

    def close(self) -> None:
        self.client.close()
