#!/usr/bin/python3

import sys
from typing import Any, Dict, List

import requests
import yaml

# Define the base URL for the Wikidata API
WIKIDATA_API_URL = "https://www.wikidata.org/w/api.php"


def get_entity_label(qid: str) -> str | None:
    """
    Fetches the English label for a given Wikidata entity QID.

    Args:
        qid (str): The Wikidata entity ID (e.g., "Q6106").

    Returns:
        Optional[str]: The English label of the entity, or None if not found.
    """
    params: Dict[str, str] = {
        "action": "wbgetentities",
        "ids": qid,
        "format": "json",
        "props": "labels",
        "languages": "en",
    }
    try:
        response = requests.get(WIKIDATA_API_URL, params=params)
        response.raise_for_status()
        entity = response.json().get("entities", {}).get(qid, {})
        return entity.get("labels", {}).get("en", {}).get("value")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching label for QID {qid}: {e}", file=sys.stderr)
        return None


def get_entity_details(qid: str) -> dict[str, Any] | None:
    """
    Fetches and processes detailed information for a given airport QID.

    Args:
        qid (str): The QID of the airport Wikidata entity.

    Returns:
        Optional[Dict[str, Any]]: A dictionary containing the detailed airport data.
    """
    params: Dict[str, str] = {
        "action": "wbgetentities",
        "ids": qid,
        "format": "json",
        "props": "claims|labels",
    }
    try:
        response = requests.get(WIKIDATA_API_URL, params=params)
        response.raise_for_status()
        entity = response.json().get("entities", {}).get(qid, {})
        if not entity:
            return None

        claims = entity.get("claims", {})

        # Helper to safely extract claim values
        def get_simple_claim_value(prop_id: str) -> str | None:
            claim = claims.get(prop_id)
            if not claim:
                return None
            v = claim[0].get("mainsnak", {}).get("datavalue", {}).get("value")
            assert isinstance(v, str) or v is None
            return v

        # Get IATA code, name, and website
        iata = get_simple_claim_value("P238")
        name = entity.get("labels", {}).get("en", {}).get("value")
        website = get_simple_claim_value("P856")

        # Get City Name by resolving its QID
        city_qid_claim = claims.get("P131")
        city_name = None
        if city_qid_claim:
            city_qid = (
                city_qid_claim[0]
                .get("mainsnak", {})
                .get("datavalue", {})
                .get("value", {})
                .get("id")
            )
            if city_qid:
                city_name = get_entity_label(city_qid)

        # Get coordinates
        coords_claim = claims.get("P625")
        latitude, longitude = None, None
        if coords_claim:
            coords = (
                coords_claim[0]
                .get("mainsnak", {})
                .get("datavalue", {})
                .get("value", {})
            )
            latitude = coords.get("latitude")
            longitude = coords.get("longitude")

        # Get elevation
        elevation_claim = claims.get("P2044")
        elevation = None
        if elevation_claim:
            amount_str = (
                elevation_claim[0]
                .get("mainsnak", {})
                .get("datavalue", {})
                .get("value", {})
                .get("amount")
            )
            if amount_str:
                elevation = float(amount_str) if "." in amount_str else int(amount_str)

        # Get Country Code
        country_claim = claims.get("P17")
        country_code = None
        if country_claim:
            country_qid = (
                country_claim[0]
                .get("mainsnak", {})
                .get("datavalue", {})
                .get("value", {})
                .get("id")
            )
            if country_qid:
                # Fetch the ISO 3166-1 alpha-2 code (P297) for the country entity
                country_code_params = {
                    "action": "wbgetclaims",
                    "entity": country_qid,
                    "property": "P297",
                    "format": "json",
                }
                country_res = requests.get(WIKIDATA_API_URL, params=country_code_params)
                country_res.raise_for_status()
                country_claims = country_res.json().get("claims", {}).get("P297")
                if country_claims:
                    code = (
                        country_claims[0]
                        .get("mainsnak", {})
                        .get("datavalue", {})
                        .get("value")
                    )
                    if code:
                        country_code = code.lower()

        data = {
            "iata": iata,
            "name": name,
            "city": city_name,
            "qid": qid,
            "latitude": latitude,
            "longitude": longitude,
            "elevation": elevation,
            "website": website,
            "country": country_code,
        }

        # Return the final structure, filtering out null values for cleaner output
        return {iata: {k: v for k, v in data.items() if v is not None}}

    except requests.exceptions.RequestException as e:
        print(f"Error fetching entity details for QID {qid}: {e}", file=sys.stderr)
        return None


def find_airport_by_iata(iata_code: str) -> dict[str, Any] | None:
    """
    Finds an airport by its IATA code using Wikidata's search API.

    Args:
        iata_code (str): The IATA code of the airport (e.g., "PDX").

    Returns:
        Optional[Dict[str, Any]]: A dictionary with the airport data or None.
    """
    params: Dict[str, str] = {
        "action": "query",
        "list": "search",
        "srsearch": f"haswbstatement:P238={iata_code.upper()}",
        "format": "json",
    }
    try:
        response = requests.get(WIKIDATA_API_URL, params=params)
        response.raise_for_status()
        search_results: List[Dict[str, Any]] = (
            response.json().get("query", {}).get("search", [])
        )

        if not search_results:
            print(f"No airport found with IATA code: {iata_code}", file=sys.stderr)
            return None

        qid = search_results[0]["title"]
        return get_entity_details(qid)

    except requests.exceptions.RequestException as e:
        print(f"Error searching on Wikidata API: {e}", file=sys.stderr)
        return None


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python airport_lookup.py <IATA_CODE>", file=sys.stderr)
        sys.exit(1)

    iata_code_arg = sys.argv[1]
    airport_data = find_airport_by_iata(iata_code_arg)

    if airport_data:
        print(
            yaml.safe_dump(
                airport_data,
                default_flow_style=False,
                allow_unicode=True,
                sort_keys=False,
            ),
            end="",
        )