from urllib.parse import urlparse

from axon.models import SearchResult

_HIGH_TRUST_TLDS = {".gov", ".edu"}
_KNOWN_HIGH_TRUST = {"wikipedia.org", "github.com", "arxiv.org", "nature.com", "who.int"}


def score_source(result: SearchResult) -> float:
    parsed = urlparse(result.url)
    host = (parsed.hostname or "").lower().removeprefix("www.")
    score = 0.5
    if any(host.endswith(tld) for tld in _HIGH_TRUST_TLDS):
        score += 0.25
    if any(host == domain or host.endswith(f".{domain}") for domain in _KNOWN_HIGH_TRUST):
        score += 0.2
    if result.published_date:
        score += 0.05
    if len(result.content) > 200:
        score += 0.05
    return round(min(score, 1.0), 3)


def annotate_source_quality(results: list[SearchResult]) -> list[SearchResult]:
    for result in results:
        result.source_quality = score_source(result)
    return results