src/dip_client.py

"""Small, explicit client for the German Bundestag DIP API.

This module intentionally keeps every field close to the source API.  It does
not infer political promise fulfilment or create status labels.  It only fetches
and normalises Bundestag DIP records so that they can be used as evidence in the
promise-tracker dashboard.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import pandas as pd
import requests

API_BASE = "https://search.dip.bundestag.de/api/v1"

RESOURCE_TYPES: Tuple[str, ...] = (
    "vorgang",
    "vorgangsposition",
    "drucksache",
    "drucksache-text",
    "plenarprotokoll",
    "plenarprotokoll-text",
    "aktivitaet",
    "person",
)

RESOURCE_FILTER_SUPPORT = {
    "aktivitaet": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
    "drucksache": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
    "drucksache-text": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
    "person": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end"},
    "plenarprotokoll": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
    "plenarprotokoll-text": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
    "vorgang": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end"},
    "vorgangsposition": {"f.wahlperiode", "f.datum.start", "f.datum.end", "f.aktualisiert.start", "f.aktualisiert.end", "f.zuordnung"},
}

KB_COLUMNS: Tuple[str, ...] = (
    "resource_type",
    "dip_id",
    "typ",
    "title",
    "abstract",
    "date",
    "updated",
    "election_period",
    "assignment",
    "document_type",
    "document_number",
    "print_type",
    "procedure_type",
    "procedure_position",
    "consultation_status",
    "initiative",
    "subject_area",
    "descriptors",
    "gesta",
    "announcement",
    "entry_into_force",
    "pdf_url",
    "api_url",
    "frontend_url_guess",
    "related_vorgang_ids",
    "related_drucksache_ids",
    "related_plenarprotokoll_ids",
    "related_person_ids",
    "text_excerpt",
    "retrieved_at",
)


def now_utc_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _json_dumps(value: Any) -> str:
    if value is None or value == "":
        return ""
    return json.dumps(value, ensure_ascii=False, sort_keys=True)


def _as_list(value: Any) -> List[Any]:
    if value is None or value == "":
        return []
    if isinstance(value, list):
        return value
    return [value]


def _join_strings(value: Any) -> str:
    parts: List[str] = []
    for item in _as_list(value):
        if item is None:
            continue
        if isinstance(item, str):
            parts.append(item)
        elif isinstance(item, dict):
            # Prefer common human-readable keys if present; otherwise preserve JSON.
            for key in ("name", "titel", "bezeichnung", "id"):
                if key in item and item[key]:
                    parts.append(str(item[key]))
                    break
            else:
                parts.append(_json_dumps(item))
        else:
            parts.append(str(item))
    return "; ".join(parts)


def _extract_ids(value: Any) -> str:
    ids: List[str] = []
    for item in _as_list(value):
        if isinstance(item, dict):
            val = item.get("id") or item.get("drucksache_id") or item.get("plenarprotokoll_id")
            if val is not None:
                ids.append(str(val))
        elif item is not None:
            ids.append(str(item))
    return "; ".join(dict.fromkeys(ids))


def _first_pdf_url(doc: Dict[str, Any]) -> str:
    fundstelle = doc.get("fundstelle")
    if isinstance(fundstelle, dict):
        return str(fundstelle.get("pdf_url") or "")
    return ""


def _extract_text_excerpt(doc: Dict[str, Any], max_chars: int = 1200) -> str:
    for key in ("text", "volltext", "inhalt", "abstract"):
        value = doc.get(key)
        if isinstance(value, str) and value.strip():
            cleaned = " ".join(value.split())
            return cleaned[:max_chars]
    return ""


def _extract_related(doc: Dict[str, Any], keys: Sequence[str]) -> str:
    collected: List[str] = []
    for key in keys:
        if key in doc:
            text = _extract_ids(doc[key])
            if text:
                collected.extend([x.strip() for x in text.split(";") if x.strip()])
    return "; ".join(dict.fromkeys(collected))


def normalise_document(resource_type: str, doc: Dict[str, Any], retrieved_at: Optional[str] = None) -> Dict[str, str]:
    """Flatten one DIP JSON document to the dashboard knowledge-base schema."""
    retrieved_at = retrieved_at or now_utc_iso()
    dip_id = str(doc.get("id", ""))

    row: Dict[str, str] = {
        "resource_type": resource_type,
        "dip_id": dip_id,
        "typ": str(doc.get("typ", "")),
        "title": str(doc.get("titel", "")),
        "abstract": str(doc.get("abstract", "")),
        "date": str(doc.get("datum", "")),
        "updated": str(doc.get("aktualisiert", "")),
        "election_period": str(doc.get("wahlperiode", "")),
        "assignment": str(doc.get("zuordnung", doc.get("herausgeber", ""))),
        "document_type": str(doc.get("dokumentart", "")),
        "document_number": str(doc.get("dokumentnummer", "")),
        "print_type": str(doc.get("drucksachetyp", doc.get("drucksachtyp", ""))),
        "procedure_type": str(doc.get("vorgangstyp", "")),
        "procedure_position": str(doc.get("vorgangsposition", "")),
        "consultation_status": str(doc.get("beratungsstand", "")),
        "initiative": _join_strings(doc.get("initiative")),
        "subject_area": _join_strings(doc.get("sachgebiet")),
        "descriptors": _join_strings(doc.get("deskriptor")),
        "gesta": str(doc.get("gesta", "")),
        "announcement": _json_dumps(doc.get("verkuendung")),
        "entry_into_force": _json_dumps(doc.get("inkrafttreten")),
        "pdf_url": _first_pdf_url(doc),
        "api_url": f"{API_BASE}/{resource_type}/{dip_id}" if dip_id else "",
        # DIP frontend URLs differ by entity and can change. This is only a convenient guess.
        "frontend_url_guess": f"https://dip.bundestag.de/{resource_type}/{dip_id}" if dip_id else "",
        "related_vorgang_ids": _extract_related(doc, ("vorgangsbezug", "vorgang", "vorgang_verlinkung")),
        "related_drucksache_ids": _extract_related(doc, ("drucksache", "drucksachen", "drucksachenbezug")),
        "related_plenarprotokoll_ids": _extract_related(doc, ("plenarprotokoll", "plenarprotokollbezug")),
        "related_person_ids": _extract_related(doc, ("person", "urheber", "redner")),
        "text_excerpt": _extract_text_excerpt(doc),
        "retrieved_at": retrieved_at,
    }

    for col in KB_COLUMNS:
        row.setdefault(col, "")
    return row


@dataclass
class DipFetchResult:
    resource_type: str
    documents: List[Dict[str, Any]]
    num_found: Optional[int]
    pages_fetched: int
    final_cursor: Optional[str]


class DipApiClient:
    """REST client for the DIP API with cursor pagination."""

    def __init__(self, api_key: str, api_base: str = API_BASE, timeout: int = 30):
        if not api_key or not api_key.strip():
            raise ValueError("A DIP API key is required. Set DIP_API_KEY or enter it in the dashboard.")
        self.api_key = api_key.strip()
        self.api_base = api_base.rstrip("/")
        self.timeout = timeout
        self.session = requests.Session()
        self.session.headers.update({"Authorization": f"ApiKey {self.api_key}"})

    def get(self, resource_type: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        if resource_type not in RESOURCE_TYPES:
            raise ValueError(f"Unsupported DIP resource type: {resource_type}")
        query = dict(params or {})
        query.setdefault("format", "json")
        url = f"{self.api_base}/{resource_type}"
        response = self.session.get(url, params=query, timeout=self.timeout)
        if response.status_code == 401:
            raise PermissionError("The DIP API rejected the API key. Check the Hugging Face secret DIP_API_KEY or the sidebar input.")
        response.raise_for_status()
        return response.json()

    def fetch_pages(
        self,
        resource_type: str,
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = 1,
        sleep_seconds: float = 0.2,
    ) -> DipFetchResult:
        allowed = RESOURCE_FILTER_SUPPORT.get(resource_type, set())
        query = {
            k: v
            for k, v in dict(params or {}).items()
            if v not in (None, "") and (not k.startswith("f.") or k in allowed)
        }
        documents: List[Dict[str, Any]] = []
        seen_cursors = set()
        cursor: Optional[str] = None
        num_found: Optional[int] = None
        pages = 0

        for _ in range(max_pages):
            if cursor:
                query["cursor"] = cursor
            payload = self.get(resource_type, query)
            pages += 1
            num_found = payload.get("numFound", num_found)
            batch = payload.get("documents", [])
            if isinstance(batch, list):
                documents.extend(batch)

            new_cursor = payload.get("cursor")
            if not new_cursor or new_cursor == cursor or new_cursor in seen_cursors:
                cursor = new_cursor
                break
            seen_cursors.add(str(new_cursor))
            cursor = str(new_cursor)
            time.sleep(sleep_seconds)

        return DipFetchResult(resource_type, documents, num_found, pages, cursor)


def build_query_params(
    wahlperiode: Optional[int] = None,
    date_start: Optional[str] = None,
    date_end: Optional[str] = None,
    updated_start: Optional[str] = None,
    updated_end: Optional[str] = None,
    zuordnung: Optional[str] = None,
) -> Dict[str, Any]:
    params: Dict[str, Any] = {}
    if wahlperiode:
        params["f.wahlperiode"] = int(wahlperiode)
    if date_start:
        params["f.datum.start"] = date_start
    if date_end:
        params["f.datum.end"] = date_end
    if updated_start:
        params["f.aktualisiert.start"] = updated_start
    if updated_end:
        params["f.aktualisiert.end"] = updated_end
    if zuordnung:
        params["f.zuordnung"] = zuordnung
    return params


def build_knowledge_base(
    api_key: str,
    resources: Sequence[str],
    params: Optional[Dict[str, Any]] = None,
    max_pages_per_resource: int = 1,
) -> Tuple[pd.DataFrame, List[Dict[str, Any]], Dict[str, Any]]:
    """Fetch several DIP endpoints and return a normalised DataFrame plus raw docs."""
    client = DipApiClient(api_key)
    retrieved_at = now_utc_iso()
    rows: List[Dict[str, str]] = []
    raw_docs: List[Dict[str, Any]] = []
    metadata: Dict[str, Any] = {
        "retrieved_at": retrieved_at,
        "params": params or {},
        "resources": list(resources),
        "resource_results": {},
    }

    for resource in resources:
        result = client.fetch_pages(resource, params=params, max_pages=max_pages_per_resource)
        metadata["resource_results"][resource] = {
            "num_found": result.num_found,
            "pages_fetched": result.pages_fetched,
            "documents_returned": len(result.documents),
            "final_cursor": result.final_cursor,
        }
        for doc in result.documents:
            raw_docs.append({"resource_type": resource, "document": doc})
            rows.append(normalise_document(resource, doc, retrieved_at=retrieved_at))

    df = pd.DataFrame(rows, columns=list(KB_COLUMNS))
    if not df.empty:
        df = df.drop_duplicates(subset=["resource_type", "dip_id"], keep="first")
    return df, raw_docs, metadata


def save_knowledge_base(df: pd.DataFrame, raw_docs: List[Dict[str, Any]], metadata: Dict[str, Any], output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_dir / "dip_knowledge_base.csv", index=False)
    with (output_dir / "dip_raw_documents.jsonl").open("w", encoding="utf-8") as f:
        for item in raw_docs:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    (output_dir / "dip_fetch_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")


def empty_knowledge_base() -> pd.DataFrame:
    return pd.DataFrame(columns=list(KB_COLUMNS))