ingest.py

from __future__ import annotations

from pathlib import Path

import httpx
from bs4 import BeautifulSoup
from rich.console import Console

from backend import Backend

DATA_PATH = Path(__file__).resolve().parent / "data" / "store.db"
KINKLIST_RAW_URL = "https://raw.githubusercontent.com/Goctionni/KinkList/master/v1.0.2.html"
OPENVERSE_URL = "https://api.openverse.org/v1/images/"
ALLOWED_LICENSES = {"cc0", "by", "by-sa", "pdm"}
console = Console()

CLUSTER_MAP = {
    "restrictive": "restraint",
    "domination": "power",
    "no consent": "edge_no_consent",
    "taboo": "extreme",
    "surrealism": "theater",
    "fluids": "body_fluids",
    "degradation": "power",
    "touch & stimulation": "body_focus",
    "misc. fetish": "fetish",
    "pain": "sensation",
    "ass play": "body_focus",
    "toys": "gear",
}


def normalize_cluster(raw_category: str, backend: Backend) -> str:
    normalized = raw_category.strip().lower()
    return CLUSTER_MAP.get(normalized, backend.slugify(raw_category))


def ensure_sources(backend: Backend) -> None:
    backend.upsert_source(
        source_id="wikimedia_commons",
        name="Wikimedia Commons",
        source_type="image",
        base_url="https://commons.wikimedia.org/",
        license_model="per-asset-open-license",
        notes="Use only with stored per-asset attribution and license.",
    )
    backend.upsert_source(
        source_id="openverse",
        name="Openverse",
        source_type="image",
        base_url="https://api.openverse.org/",
        license_model="per-asset-open-license",
        notes="Openverse image search results filtered by license and attribution.",
    )
    backend.upsert_source(
        source_id="kinklist_github",
        name="KinkList GitHub",
        source_type="taxonomy",
        base_url="https://github.com/Goctionni/KinkList",
        license_model="repository-license",
        notes="Imported taxonomy terms from the public KinkList project.",
    )
    backend.upsert_source(
        source_id="ncsf",
        name="National Coalition for Sexual Freedom",
        source_type="definition",
        base_url="https://ncsfreedom.org/",
        license_model="mixed",
        notes="Primary target for safety and consent definitions.",
    )


def fetch_kinklist_text(client: httpx.Client) -> str:
    response = client.get(KINKLIST_RAW_URL, timeout=30.0)
    response.raise_for_status()
    soup = BeautifulSoup(response.text, "html.parser")
    textarea = soup.find("textarea")
    if not textarea or not textarea.text.strip():
        raise RuntimeError("Could not find KinkList textarea content.")
    return textarea.text


def import_kinklist(backend: Backend, client: httpx.Client) -> None:
    text = fetch_kinklist_text(client)
    category = ""
    roles = ""
    imported = 0
    merged = 0
    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if line.startswith("#"):
            category = line.lstrip("#").strip()
            roles = ""
            continue
        if line.startswith("(") and line.endswith(")"):
            roles = line
            continue
        if not line.startswith("* "):
            continue
        term = line[2:].strip()
        matched_id = backend.match_kink_id(term)
        cluster = normalize_cluster(category or "imported", backend)
        notes = f"Imported from KinkList category '{category}'. Roles {roles}." if roles else f"Imported from KinkList category '{category}'."
        if matched_id:
            backend.add_alias(matched_id, term, "kinklist_github")
            backend.add_definition(
                matched_id,
                "kinklist_github",
                f"KinkList categorizes '{term}' under {category} {roles}".strip(),
                tone="taxonomy",
                source_url=KINKLIST_RAW_URL,
            )
            merged += 1
        else:
            kink_id = backend.slugify(term)
            backend.ensure_kink(
                kink_id=kink_id,
                name=term,
                cluster=cluster,
                short_definition=f"Imported taxonomy term from KinkList category '{category}'.",
                risk_level="extreme" if cluster in {"extreme", "edge_no_consent"} else "general",
                is_extreme=cluster in {"extreme", "edge_no_consent"},
                notes=notes,
            )
            backend.add_definition(
                kink_id,
                "kinklist_github",
                f"KinkList categorizes '{term}' under {category} {roles}".strip(),
                tone="taxonomy",
                source_url=KINKLIST_RAW_URL,
            )
            backend.add_example(kink_id, f"Category: {category}", "example")
            if roles:
                backend.add_example(kink_id, roles, "example")
            imported += 1
    console.print(f"Imported new KinkList terms: [bold]{imported}[/bold], merged as aliases/definitions: [bold]{merged}[/bold]")


def openverse_results(client: httpx.Client, query: str, page_size: int = 2) -> list[dict]:
    response = client.get(OPENVERSE_URL, params={"q": query, "page_size": page_size}, timeout=30.0)
    response.raise_for_status()
    payload = response.json()
    return payload.get("results", [])


def import_openverse_assets(backend: Backend, client: httpx.Client, per_kink_limit: int = 2) -> None:
    inserted = 0
    for kink in backend.list_kinks():
        queries = [kink["name"], *kink["aliases"][:1]]
        seen_urls: set[str] = set()
        for query in queries:
            if len(seen_urls) >= per_kink_limit:
                break
            try:
                results = openverse_results(client, query)
            except Exception:
                continue
            for result in results:
                if len(seen_urls) >= per_kink_limit:
                    break
                license_code = (result.get("license") or "").lower()
                asset_url = result.get("url") or ""
                if not asset_url or asset_url in seen_urls:
                    continue
                if license_code not in ALLOWED_LICENSES:
                    continue
                backend.add_asset(
                    kink_id=kink["id"],
                    source_id="openverse",
                    asset_url=asset_url,
                    thumbnail_url=result.get("thumbnail") or "",
                    mime_type=result.get("filetype") or "image/jpeg",
                    license=result.get("license_url") or result.get("license") or "",
                    creator=result.get("creator") or "",
                    attribution_text=result.get("attribution") or "",
                    is_explicit=bool(result.get("mature")),
                    is_illustration=False,
                    click_to_reveal=True,
                )
                seen_urls.add(asset_url)
                inserted += 1
                if result.get("title"):
                    backend.add_example(kink["id"], result["title"], "visual_cue")
    console.print(f"Imported Openverse assets: [bold]{inserted}[/bold]")


def main() -> None:
    backend = Backend(DATA_PATH)
    ensure_sources(backend)
    with httpx.Client(follow_redirects=True, headers={"User-Agent": "kink-cli-prototype/0.1"}) as client:
        import_kinklist(backend, client)
        import_openverse_assets(backend, client)
    backend.rebuild_similarity_edges()
    stats = backend.catalog_stats()
    console.print("[bold]Catalog stats after ingestion[/bold]")
    for key, value in stats.items():
        console.print(f"- {key}: {value}")


if __name__ == "__main__":
    main()