src/olostep_client.py

from __future__ import annotations

import concurrent.futures
import os
from typing import Any
from urllib.parse import urlparse

from olostep import Olostep, Olostep_BaseError, RetryStrategy

from src.schemas import SourcePage, UserProfile


class OlostepError(RuntimeError):
    pass


BLOCKED_HOST_PARTS = (
    "facebook.com",
    "instagram.com",
    "linkedin.com",
    "pinterest.",
    "reddit.com",
    "tiktok.com",
    "x.com",
    "twitter.com",
    "youtube.com",
    "medium.com",
    "mastersportal.com",
    "scholarshipportal.com",
    "scholars4dev.com",
    "applykite.com",
    "profellow.com",
    "topuniversities.com",
    "wemakescholars.com",
    "computerscience.org",
    "timesconsultant.com",
)

LIST_PHRASES = (
    "top scholarships",
    "best scholarships",
    "scholarships list",
    "study scholarships in",
    "scholarships for students",
    "scholarships for pakistani",
)

IMPORTANT_WORDS = (
    "apply",
    "application",
    "eligibility",
    "deadline",
    "award",
)


class OlostepClient:
    """Two-stage research through the official Olostep Python SDK."""

    def __init__(self, api_key: str | None = None, timeout: int = 30) -> None:
        self.api_key = os.getenv("OLOSTEP_API_KEY", "") if api_key is None else api_key
        self.timeout = timeout
        self._sdk: Olostep | None = None
        self._sdk_key = ""

    @property
    def enabled(self) -> bool:
        return bool(self.api_key)

    def search(self, profile: UserProfile, limit: int = 15) -> list[SourcePage]:
        """Run short profile-based searches, then scrape only the best eight links."""
        if not self.enabled:
            return []

        queries = _build_queries(profile)
        raw_links = self._search_many(queries, per_query=8)
        selected = _select_important_links(raw_links, profile, limit=min(limit, 8))
        return self._scrape_many(selected, total_limit=limit)

    def _get_sdk(self) -> Olostep:
        if self._sdk is None or self._sdk_key != self.api_key:
            os.environ["OLOSTEP_API_TIMEOUT"] = str(self.timeout)
            self._sdk = Olostep(
                api_key=self.api_key,
                retry_strategy=RetryStrategy(max_retries=1, initial_delay=0.5),
            )
            self._sdk_key = self.api_key
        return self._sdk

    def _search_one(self, query: str, limit: int) -> list[dict[str, Any]]:
        try:
            result = self._get_sdk().searches.create(
                query=query,
                limit=min(max(limit, 1), 10),
            )
        except Olostep_BaseError as exc:
            raise OlostepError(f"Olostep search failed: {exc}") from exc
        except Exception as exc:
            raise OlostepError(f"Olostep search failed: {exc}") from exc
        return list(result.links or [])[:limit]

    def _search_many(self, queries: list[str], per_query: int) -> list[dict[str, Any]]:
        found: dict[str, dict[str, Any]] = {}
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(queries)) as executor:
            futures = [executor.submit(self._search_one, query, per_query) for query in queries]
            for future in concurrent.futures.as_completed(futures):
                try:
                    for item in future.result():
                        url = item.get("url") or item.get("link") or ""
                        if url and not _is_blocked_result(item):
                            found.setdefault(_dedupe_key(url), item)
                except OlostepError:
                    continue
        return list(found.values())

    def _scrape_one(self, item: dict[str, Any]) -> SourcePage | None:
        url = item.get("url") or item.get("link") or ""
        if not url:
            return None
        try:
            scraped = self._get_sdk().scrapes.create(
                url_to_scrape=url,
                formats=["markdown"],
                remove_images=True,
            )
        except Olostep_BaseError as exc:
            raise OlostepError(f"Olostep scrape failed for {url}: {exc}") from exc
        except Exception as exc:
            raise OlostepError(f"Olostep scrape failed for {url}: {exc}") from exc

        content = scraped.markdown_content or ""
        if not content:
            return None
        try:
            return SourcePage(
                title=item.get("title") or "Funding source",
                url=url,
                snippet=item.get("description") or item.get("snippet") or "",
                content=str(content)[:8000],
            )
        except ValueError:
            return None

    def _scrape_many(
        self, items: list[dict[str, Any]], total_limit: int
    ) -> list[SourcePage]:
        pages: list[SourcePage] = []
        if not items:
            return pages
        with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(items), 8)) as executor:
            futures = [executor.submit(self._scrape_one, item) for item in items]
            for future in concurrent.futures.as_completed(futures):
                try:
                    page = future.result()
                    if page is not None:
                        pages.append(page)
                except OlostepError:
                    continue
        return pages[:total_limit]


def _is_blocked_url(url: str) -> bool:
    normalized = url.lower()
    return any(host_part in normalized for host_part in BLOCKED_HOST_PARTS)


def _is_blocked_result(item: dict[str, Any]) -> bool:
    """Reject known aggregators and generic commercial scholarship articles."""
    url = item.get("url") or item.get("link") or ""
    if _is_blocked_url(url):
        return True
    parsed = urlparse(url.lower())
    host = parsed.netloc.removeprefix("www.")
    title = (item.get("title") or "").lower()
    institutional = any(marker in host for marker in (".edu", ".gov", ".ac."))
    return not institutional and any(phrase in title for phrase in LIST_PHRASES)


def _importance_score(item: dict[str, Any], profile: UserProfile) -> int:
    url = (item.get("url") or item.get("link") or "").lower()
    text = " ".join(
        [
            item.get("title") or "",
            item.get("description") or item.get("snippet") or "",
            url,
        ]
    ).lower()
    score = 4 if any(domain in url for domain in (".edu", ".gov", ".org")) else 0
    score += 3 if any(term in text for term in ("scholarship", "grant", "fellowship", "funding", "financial aid")) else 0
    score += 2 if any(word in text for word in IMPORTANT_WORDS) else 0
    score -= 5 if any(phrase in text for phrase in LIST_PHRASES) else 0
    return score


def _select_important_links(
    items: list[dict[str, Any]], profile: UserProfile, limit: int
) -> list[dict[str, Any]]:
    ranked = sorted(
        items,
        key=lambda item: _importance_score(item, profile),
        reverse=True,
    )
    return ranked[:limit]


def _build_queries(profile: UserProfile) -> list[str]:
    preferences = profile.preferences or "international applicants"
    return [
        f"site:.edu {profile.applying_for} {profile.field_of_study} scholarship international",
        f"site:.gov {profile.origin} students {profile.applying_for} scholarship",
        f"official university fully funded {profile.applying_for} {profile.field_of_study}",
        f"official research council fellowship grant {profile.field_of_study}",
        f"official foundation {profile.field_of_study} funding {preferences}",
    ]


def _dedupe_key(url: Any) -> str:
    parsed = urlparse(str(url))
    host = (parsed.netloc or "").lower()
    host = host[4:] if host.startswith("www.") else host
    return f"{host}{(parsed.path or '').rstrip('/')}".lower()