import requests

SPARQL_URL = "https://query.wikidata.org/sparql"
HEADERS = {"User-Agent": "WildDex/1.0 (wildlife app; contact via github)", "Accept": "application/sparql-results+json"}


def fetch_wikidata_traits(scientific_name: str) -> dict:
    """
    Query Wikidata by scientific name and return available trait data.
    Returns a dict with keys matching species_detail columns (values may be None).
    """
    query = f"""
    SELECT
      ?mass ?massUnit ?massUnitLabel
      ?lifespan ?lifespanUnit ?lifespanUnitLabel
      ?dietLabel
      ?habitatLabel
      ?activityLabel
    WHERE {{
      ?item wdt:P225 "{scientific_name}" .

      OPTIONAL {{
        ?item p:P2067 ?massStmt .
        ?massStmt psv:P2067 ?massVal .
        ?massVal wikibase:quantityAmount ?mass .
        ?massVal wikibase:quantityUnit ?massUnit .
      }}
      OPTIONAL {{
        ?item p:P2250 ?lifespanStmt .
        ?lifespanStmt psv:P2250 ?lifespanVal .
        ?lifespanVal wikibase:quantityAmount ?lifespan .
        ?lifespanVal wikibase:quantityUnit ?lifespanUnit .
      }}
      OPTIONAL {{ ?item wdt:P1034 ?diet . }}
      OPTIONAL {{ ?item wdt:P2974 ?habitat . }}
      OPTIONAL {{ ?item wdt:P5204 ?activity . }}

      SERVICE wikibase:label {{ bd:serviceParam wikibase:language "en" . }}
    }}
    LIMIT 5
    """

    resp = requests.get(SPARQL_URL, params={"query": query, "format": "json"}, headers=HEADERS)
    resp.raise_for_status()
    bindings = resp.json()["results"]["bindings"]

    if not bindings:
        print(f"  No Wikidata results for '{scientific_name}'")
        return {}

    return _extract_traits(bindings)


def _extract_traits(bindings: list[dict]) -> dict:
    result = {
        "size_range": None,
        "weight_range": None,
        "lifespan": None,
        "diet": None,
        "activity_period": None,
        "habitat_type": None,
    }

    diets, habitats, activities = set(), set(), set()
    masses, lifespans = [], []

    for row in bindings:
        if "mass" in row and "massUnitLabel" in row:
            mass = float(row["mass"]["value"])
            unit = row["massUnitLabel"]["value"]
            masses.append(f"{mass:g} {unit}")

        if "lifespan" in row and "lifespanUnitLabel" in row:
            ls = float(row["lifespan"]["value"])
            unit = row["lifespanUnitLabel"]["value"]
            lifespans.append(f"{ls:g} {unit}")

        if "dietLabel" in row:
            diets.add(row["dietLabel"]["value"])

        if "habitatLabel" in row:
            habitats.add(row["habitatLabel"]["value"])

        if "activityLabel" in row:
            activities.add(row["activityLabel"]["value"])

    if masses:
        result["weight_range"] = " / ".join(sorted(set(masses)))
    if lifespans:
        result["lifespan"] = " / ".join(sorted(set(lifespans)))
    if diets:
        result["diet"] = ", ".join(sorted(diets))
    if habitats:
        result["habitat_type"] = ", ".join(sorted(habitats))
    if activities:
        result["activity_period"] = ", ".join(sorted(activities))

    return result


if __name__ == "__main__":
    test_species = "Thamnophis sirtalis"
    print(f"Fetching Wikidata traits for '{test_species}'...")
    traits = fetch_wikidata_traits(test_species)
    for key, value in traits.items():
        print(f"  {key}: {value}")