"""LangChain-powered AI insights for the Repository (PRD REPO-SP9).

Chains (all live in this module; each is callable independently):

* ``summarise_document`` — map-reduce summary over a document's chunks.
* ``classify_document`` — structured-output collection + keyword suggestion.
* ``extract_metadata`` — structured-output metadata extraction (authors, dates…).
* ``embed_document_chunks`` — splits, embeds and persists chunks to pgvector.
* ``analyse_research_trends`` — clustered RAG analysis across the corpus.

The implementations degrade gracefully when ``OPENAI_API_KEY`` is missing
(stub output, no network calls). This keeps local/CI workflows self-contained
while still giving production the real LangChain pipeline.
"""

from __future__ import annotations

import logging
import statistics
from datetime import timedelta
from typing import Any

from django.conf import settings
from django.db.models import Count, Q
from django.utils import timezone

logger = logging.getLogger(__name__)


# ── Availability & clients ─────────────────────────────────────────────────


def _has_api_key() -> bool:
    return bool(getattr(settings, "OPENAI_API_KEY", "").strip())


def _chat_fast():
    from langchain_openai import ChatOpenAI

    return ChatOpenAI(
        model=settings.LANGCHAIN_CHAT_MODEL_FAST,
        temperature=0,
        api_key=settings.OPENAI_API_KEY,
    )


def _chat_deep():
    from langchain_openai import ChatOpenAI

    return ChatOpenAI(
        model=settings.LANGCHAIN_CHAT_MODEL_DEEP,
        temperature=0.2,
        api_key=settings.OPENAI_API_KEY,
    )


def _embeddings():
    from langchain_openai import OpenAIEmbeddings

    return OpenAIEmbeddings(
        model=settings.LANGCHAIN_EMBEDDING_MODEL,
        api_key=settings.OPENAI_API_KEY,
    )


def _pg_connection_string() -> str:
    db = settings.DATABASES["default"]
    return (
        f"postgresql+psycopg://{db['USER']}:{db['PASSWORD']}@{db['HOST']}:{db['PORT']}/{db['NAME']}"
    )


def get_pgvector_store():
    """Return the shared PGVector store used by embed + trend chains."""
    from langchain_postgres import PGVector

    return PGVector(
        collection_name=settings.REPOSITORY_PGVECTOR_COLLECTION,
        connection=_pg_connection_string(),
        embeddings=_embeddings(),
        use_jsonb=True,
    )


def _splitter():
    from langchain_text_splitters import RecursiveCharacterTextSplitter

    return RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)


# ── Chains ─────────────────────────────────────────────────────────────────


def _audit_ai_generation(action: str, document, insight) -> None:
    """FRREP-AR013 — AI generation events get a distinct audit-log action code.

    These chains run from signal receivers (see analytics/receivers.py) with
    no request/actor in scope, so ``actor=None`` — the audit row still gives
    managers a "what AI activity happened, on what, and when" trail, which
    previously didn't exist at all for AI insight generation (only dashboard/
    export ``ACCESS`` entries were logged, via middleware).
    """
    from apps.repository.documents.audit_helper import audit_repository

    try:
        audit_repository(
            actor=None,
            action=action,
            target_model="AIDocumentInsight",
            object_id=insight.pk,
            object_repr=document.title[:200] if document else "",
            changes={
                "document_id": document.pk if document else None,
                "insight_kind": insight.kind,
                "model_used": insight.model_used,
            },
        )
    except Exception:  # pragma: no cover — auditing must never break generation
        logger.exception("ai_insights: failed to write audit log for action=%s", action)


def summarise_document(document_id: int):
    """Generate a 3–5 sentence abstract for documents without one."""
    from apps.repository.analytics.models import AIDocumentInsight
    from apps.repository.documents.models import Document

    try:
        document = Document.objects.get(pk=document_id)
    except Document.DoesNotExist:
        return None

    text = (document.extracted_text or document.abstract or document.title).strip()
    if not text:
        return None

    if not _has_api_key():
        summary = text[:800] + ("…" if len(text) > 800 else "")
        model_name = "stub"
    else:
        from langchain_core.messages import HumanMessage, SystemMessage

        chunks = _splitter().split_text(text)
        partials: list[str] = []
        for chunk in chunks[:16]:
            msg = _chat_deep().invoke(
                [
                    SystemMessage(
                        "You are a research librarian writing neutral, factual mini-abstracts."
                    ),
                    HumanMessage(f"Summarise this in 2 sentences:\n\n{chunk}"),
                ]
            )
            partials.append(msg.content.strip())
        combined = " ".join(partials)
        final = _chat_deep().invoke(
            [
                {
                    "role": "system",
                    "content": (
                        "You are compiling a final abstract (3–5 sentences, 80–150 words) "
                        "from the partial summaries below. No marketing language, no editorialising."
                    ),
                },
                {"role": "user", "content": combined},
            ]
        )
        summary = final.content.strip()
        model_name = settings.LANGCHAIN_CHAT_MODEL_DEEP

    insight = AIDocumentInsight.objects.create(
        document=document,
        kind=AIDocumentInsight.Kind.SUMMARY,
        payload={"summary": summary},
        model_used=model_name,
        confidence=None,
    )
    Document.objects.filter(pk=document.pk).update(
        ai_summary=summary,
        ai_summary_generated_at=timezone.now(),
    )
    _audit_ai_generation("AI_SUMMARY_GENERATED", document, insight)
    return insight


def classify_document(document_id: int):
    """Suggest collection + subject keywords (REPO-SP9)."""
    from apps.repository.analytics.models import AIDocumentInsight
    from apps.repository.documents.models import Collection, Document

    try:
        document = Document.objects.get(pk=document_id)
    except Document.DoesNotExist:
        return None

    text = (document.abstract or document.extracted_text[:4000] or document.title).strip()
    available_slugs = list(Collection.objects.values_list("slug", flat=True))

    if not _has_api_key() or not available_slugs:
        payload = {
            "suggested_collection_slug": document.collection.slug if document.collection_id else (available_slugs[0] if available_slugs else ""),
            "keywords": [w.lower() for w in (document.title or "").split()[:6]],
            "confidence": 0.0,
        }
        model_name = "stub"
    else:
        from pydantic import BaseModel, Field

        class ClassificationResult(BaseModel):
            suggested_collection_slug: str = Field(..., description="One of the provided collection slugs")
            keywords: list[str] = Field(..., min_length=3, max_length=10)
            confidence: float = Field(..., ge=0.0, le=1.0)

        structured = _chat_fast().with_structured_output(ClassificationResult)
        result = structured.invoke(
            [
                {
                    "role": "system",
                    "content": (
                        "You classify research documents into a fixed taxonomy. "
                        f"Valid collection slugs: {', '.join(available_slugs)}. "
                        "Always pick the single most appropriate slug. "
                        "Return 3–10 subject keywords (lower-case, noun phrases, no stop-words)."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Title: {document.title}\n\nText:\n{text[:3500]}",
                },
            ]
        )
        payload = result.model_dump()
        model_name = settings.LANGCHAIN_CHAT_MODEL_FAST

    insight = AIDocumentInsight.objects.create(
        document=document,
        kind=AIDocumentInsight.Kind.CLASSIFICATION,
        payload=payload,
        confidence=payload.get("confidence"),
        model_used=model_name,
    )
    _audit_ai_generation("AI_CLASSIFICATION_GENERATED", document, insight)
    return insight


def extract_metadata(document_id: int):
    """Pull authors, publication date, funding sources, themes from the file text."""
    from apps.repository.analytics.models import AIDocumentInsight
    from apps.repository.documents.models import Document

    try:
        document = Document.objects.get(pk=document_id)
    except Document.DoesNotExist:
        return None

    text = (document.extracted_text or document.abstract or document.title).strip()
    if not text:
        return None

    if not _has_api_key():
        payload = {
            "authors": [],
            "publication_date": None,
            "funding_sources": [],
            "key_themes": [],
            "language": document.language,
        }
        model_name = "stub"
    else:
        from pydantic import BaseModel, Field

        class ExtractedMetadata(BaseModel):
            authors: list[str] = Field(default_factory=list)
            publication_date: str | None = None
            funding_sources: list[str] = Field(default_factory=list)
            key_themes: list[str] = Field(default_factory=list)
            language: str | None = None

        structured = _chat_fast().with_structured_output(ExtractedMetadata)
        result = structured.invoke(
            [
                {
                    "role": "system",
                    "content": (
                        "Extract structured metadata from the first pages of a research document. "
                        "Return empty lists / null when something is not present."
                    ),
                },
                {"role": "user", "content": text[:6000]},
            ]
        )
        payload = result.model_dump()
        model_name = settings.LANGCHAIN_CHAT_MODEL_FAST

    insight = AIDocumentInsight.objects.create(
        document=document,
        kind=AIDocumentInsight.Kind.METADATA,
        payload=payload,
        model_used=model_name,
    )
    Document.objects.filter(pk=document.pk).update(ai_metadata=payload)
    _audit_ai_generation("AI_METADATA_EXTRACTED", document, insight)
    return insight


def embed_document_chunks(document_id: int) -> int:
    """Chunk + embed a document and store vectors in pgvector."""
    from apps.repository.analytics.models import DocumentEmbedding
    from apps.repository.documents.models import Document

    try:
        document = Document.objects.get(pk=document_id)
    except Document.DoesNotExist:
        return 0

    text = (document.extracted_text or document.abstract or document.title).strip()
    if not text:
        return 0

    chunks = _splitter().split_text(text)
    if not chunks:
        return 0

    if not _has_api_key():
        # Store zero-embeddings so tests / dev runs still exercise the data path.
        DocumentEmbedding.objects.filter(document=document).delete()
        stored = 0
        for idx, chunk in enumerate(chunks):
            DocumentEmbedding.objects.create(
                document=document,
                chunk_index=idx,
                chunk_text=chunk,
                **(
                    {"embedding": [0.0] * settings.LANGCHAIN_EMBEDDING_DIMENSIONS}
                    if _has_pgvector()
                    else {"embedding_raw": []}
                ),
                metadata={
                    "public_document_id": document.public_document_id,
                    "collection": document.collection.slug if document.collection_id else "",
                    "document_type": document.document_type,
                },
            )
            stored += 1
        return stored

    from langchain_core.documents import Document as LCDocument

    store = get_pgvector_store()
    lc_docs = [
        LCDocument(
            page_content=chunk,
            metadata={
                "document_id": document.pk,
                "public_document_id": document.public_document_id,
                "collection": document.collection.slug if document.collection_id else "",
                "document_type": document.document_type,
                "chunk_index": idx,
            },
        )
        for idx, chunk in enumerate(chunks)
    ]
    store.add_documents(lc_docs)

    # Also persist in the Django-owned table for analytics joins.
    emb_model = _embeddings()
    vectors = emb_model.embed_documents(chunks)
    DocumentEmbedding.objects.filter(document=document).delete()
    for idx, (chunk, vec) in enumerate(zip(chunks, vectors)):
        DocumentEmbedding.objects.create(
            document=document,
            chunk_index=idx,
            chunk_text=chunk,
            **({"embedding": vec} if _has_pgvector() else {"embedding_raw": list(vec)}),
            metadata={
                "public_document_id": document.public_document_id,
                "collection": document.collection.slug if document.collection_id else "",
                "document_type": document.document_type,
            },
        )
    return len(chunks)


def analyse_research_trends(period_days: int = 180, n_clusters: int = 8) -> list[Any]:
    """Produce all three ``ResearchTrend`` kinds for the AI insights page (REPO-SP9 AR010):

    * ``CLUSTER``     — KMeans over recent embeddings, LLM-labelled (unchanged
      behaviour, factored out into :func:`_build_cluster_trends`).
    * ``GROWTH_AREA`` — a subject keyword whose share of published output grew
      significantly between the prior window and this one
      (:func:`_detect_growth_areas`).
    * ``GAP``         — a curated subject keyword whose document count sits
      well below the corpus median (:func:`_detect_gaps`).

    Growth/gap detection is deterministic (no embeddings or LLM involved) and
    deliberately keyword-based rather than cluster-based: KMeans cluster ids
    aren't stable across independent fits, so there's no reliable way to say
    "cluster 3 today is the same topic as cluster 3 last period". The
    Repository's AGROVOC subject-keyword taxonomy is already curated on every
    document and gives a stable comparison unit instead. These two therefore
    still run even when there isn't enough embedded content yet for KMeans.
    """
    period_start = timezone.now().date() - timedelta(days=period_days)
    end_date = timezone.now().date()

    trends: list[Any] = _build_cluster_trends(period_start, end_date, n_clusters=n_clusters)
    trends.extend(_detect_growth_areas(period_start, end_date))
    trends.extend(_detect_gaps(period_start, end_date))
    return trends


def _build_cluster_trends(period_start, end_date, *, n_clusters: int = 8) -> list[Any]:
    """Cluster recent embeddings and let the LLM label each cluster as a research trend.

    Unchanged behaviour from the original ``analyse_research_trends`` — split
    out so growth/gap detection (below) can run independently of it.
    """
    from apps.repository.analytics.models import DocumentEmbedding, ResearchTrend
    from apps.repository.documents.models import Document

    doc_ids = list(
        Document.objects.filter(
            status=Document.Status.PUBLISHED,
            published_at__date__gte=period_start,
        ).values_list("id", flat=True)
    )
    if not doc_ids:
        return []

    rows = list(DocumentEmbedding.objects.filter(document_id__in=doc_ids))
    if len(rows) < max(4, n_clusters):
        return []

    if not _has_api_key():
        trend = ResearchTrend.objects.create(
            period_start=period_start,
            period_end=end_date,
            kind=ResearchTrend.Kind.CLUSTER,
            title="Trend analysis pending (no OPENAI_API_KEY configured)",
            description="Configure OPENAI_API_KEY to enable LangChain-powered trend analysis.",
            score=0.0,
        )
        return [trend]

    try:
        import numpy as np
        from sklearn.cluster import KMeans
    except ImportError:
        logger.warning("scikit-learn not installed; skipping trend analysis")
        return []

    def _row_vector(row):
        return row.embedding if _has_pgvector() else row.embedding_raw

    raw = [_row_vector(row) for row in rows]
    raw = [v for v in raw if v is not None and len(v)]
    if not raw:
        return []
    vectors = np.array(raw)

    k = min(n_clusters, max(2, len(vectors) // 4))
    km = KMeans(n_clusters=k, random_state=42, n_init="auto").fit(vectors)

    trends: list[ResearchTrend] = []
    for cluster_id in range(k):
        idxs = [i for i, lbl in enumerate(km.labels_) if lbl == cluster_id]
        if not idxs:
            continue
        sample_chunks = [rows[i].chunk_text for i in idxs[:6]]
        response = _chat_deep().invoke(
            [
                {
                    "role": "system",
                    "content": (
                        "You are labelling a cluster of research excerpts. Produce a short title "
                        "(<= 8 words) and a 2–3 sentence description of the shared theme. "
                        "Return as TITLE: ...\\nDESCRIPTION: ..."
                    ),
                },
                {"role": "user", "content": "\n---\n".join(sample_chunks)},
            ]
        )
        content = response.content
        title = "Emerging cluster"
        description = content.strip()
        for line in content.splitlines():
            if line.lower().startswith("title:"):
                title = line.split(":", 1)[1].strip()[:255]
            elif line.lower().startswith("description:"):
                description = line.split(":", 1)[1].strip()
        trend = ResearchTrend.objects.create(
            period_start=period_start,
            period_end=end_date,
            kind=ResearchTrend.Kind.CLUSTER,
            title=title,
            description=description,
            score=float(len(idxs)) / len(vectors),
        )
        doc_ids_in_cluster = {rows[i].document_id for i in idxs}
        trend.supporting_documents.set(
            Document.objects.filter(pk__in=doc_ids_in_cluster)
        )
        trends.append(trend)
    return trends


def _keyword_doc_counts(period_start=None, period_end=None):
    """AGROVOC-approved keyword -> published-document-count within an optional window."""
    from apps.repository.documents.models import Document, Keyword

    doc_filter = Q(documents__status=Document.Status.PUBLISHED)
    if period_start is not None:
        doc_filter &= Q(documents__published_at__date__gte=period_start)
    if period_end is not None:
        doc_filter &= Q(documents__published_at__date__lt=period_end)

    return {
        row["slug"]: row["total"]
        for row in Keyword.objects.filter(
            kind=Keyword.Kind.AGROVOC, status=Keyword.Status.APPROVED
        )
        .annotate(total=Count("documents", filter=doc_filter))
        .filter(total__gt=0)
        .values("slug", "total")
    }


def _detect_growth_areas(period_start, period_end, top_n: int = 5) -> list[Any]:
    """FRREP-AR010 — flag subject keywords whose share of published output grew
    significantly between the prior window (same length, immediately before
    ``period_start``) and the current window (``period_start``..``period_end``).

    "Grew significantly" = at least doubled its share of tagged output *and*
    has at least 2 documents this period (guards against a single new
    document creating a spurious 100% "growth" reading).
    """
    from apps.repository.documents.models import Keyword, Document
    from apps.repository.analytics.models import ResearchTrend

    window_len = max((period_end - period_start).days, 1)
    prior_start = period_start - timedelta(days=window_len)

    prior_counts = _keyword_doc_counts(prior_start, period_start)
    current_counts = _keyword_doc_counts(period_start, period_end)
    prior_total = sum(prior_counts.values()) or 1
    current_total = sum(current_counts.values()) or 1

    candidates = []
    for slug, current_n in current_counts.items():
        if current_n < 2:
            continue
        current_share = current_n / current_total
        prior_share = prior_counts.get(slug, 0) / prior_total
        if current_share >= prior_share * 2:
            candidates.append((slug, current_n, current_share, prior_share))

    candidates.sort(key=lambda c: c[2] - c[3], reverse=True)

    trends = []
    for slug, current_n, current_share, prior_share in candidates[:top_n]:
        kw = Keyword.objects.filter(slug=slug).first()
        if kw is None:
            continue
        doc_ids = list(
            Document.objects.filter(
                keywords=kw,
                status=Document.Status.PUBLISHED,
                published_at__date__gte=period_start,
                published_at__date__lt=period_end,
            ).values_list("id", flat=True)[:50]
        )
        trend = ResearchTrend.objects.create(
            period_start=period_start,
            period_end=period_end,
            kind=ResearchTrend.Kind.GROWTH_AREA,
            title=f"Growing interest: {kw.label}",
            description=(
                f'"{kw.label}" grew from {round(prior_share * 100, 1)}% to '
                f"{round(current_share * 100, 1)}% of published output between the "
                f"prior and current period ({current_n} document"
                f"{'s' if current_n != 1 else ''} this period)."
            ),
            score=round(current_share - prior_share, 4),
        )
        trend.supporting_documents.set(Document.objects.filter(pk__in=doc_ids))
        trends.append(trend)
    return trends


def _detect_gaps(period_start, period_end, top_n: int = 5) -> list[Any]:
    """FRREP-AR010 — flag under-represented subject keywords as content gaps.

    Heuristic: this codebase has no curated "RUFORUM priority research areas"
    list to diff coverage against, so the corpus's own AGROVOC subject
    taxonomy is used as the expected-coverage baseline instead — keywords
    already curated/approved as relevant subjects (i.e. someone tagged a
    document with them) but sitting at <= 25% of the corpus-wide median
    document count are flagged as under-covered relative to their peers.
    Requires at least 5 tagged subjects for a median to be meaningful.
    """
    from apps.repository.documents.models import Document, Keyword
    from apps.repository.analytics.models import ResearchTrend

    counts = _keyword_doc_counts()
    if len(counts) < 5:
        return []

    median = statistics.median(counts.values())
    if median <= 0:
        return []
    threshold = max(1, median * 0.25)

    gap_slugs = sorted(
        (slug for slug, total in counts.items() if total <= threshold),
        key=lambda slug: counts[slug],
    )

    trends = []
    for slug in gap_slugs[:top_n]:
        kw = Keyword.objects.filter(slug=slug).first()
        if kw is None:
            continue
        total = counts[slug]
        doc_ids = list(
            Document.objects.filter(keywords=kw, status=Document.Status.PUBLISHED).values_list(
                "id", flat=True
            )[:20]
        )
        trend = ResearchTrend.objects.create(
            period_start=period_start,
            period_end=period_end,
            kind=ResearchTrend.Kind.GAP,
            title=f"Coverage gap: {kw.label}",
            description=(
                f'"{kw.label}" has only {total} document{"s" if total != 1 else ""} in the '
                f"corpus, well below the median of {median:g} across tagged subjects — a "
                "candidate content gap for future submissions or outreach."
            ),
            score=round(total / median, 4),
        )
        trend.supporting_documents.set(Document.objects.filter(pk__in=doc_ids))
        trends.append(trend)
    return trends


def _has_pgvector() -> bool:
    from apps.repository.analytics.models import _PGVECTOR_AVAILABLE

    return _PGVECTOR_AVAILABLE