from mm_stats.auth import PostgresDb
from mm_stats.definitions import logger


def get_audacious_country_ids():
    """Get country ids for audacious countries."""
    db = PostgresDb()
    sql = "set schema 'data_preparation';"
    sql += """
        select id
        from ne_10m_admin_0_countries
        where region_id > 0
        order by id
    """

    data = db.retr_query(sql)
    country_ids = [item[0] for item in data]

    logger.info(f"got audacious country ids: {country_ids}")
    return country_ids


def get_tm_country_ids():
    """Get country ids for all countries with tm projects."""
    db = PostgresDb()
    sql = "set schema 'data_preparation';"
    sql += """
        select
          ne_id
        from data_preparation.projects p
        where ne_id is not null
        group by ne_id, ne_name
        order by ne_id
    """

    data = db.retr_query(sql)
    country_ids = [item[0] for item in data]

    logger.info(f"got country ids for all tm projects: {country_ids}")
    return country_ids


def get_country_ids():
    """Get country ids for all countries with tm projects or HOT Audacious."""
    tm_countries = get_tm_country_ids()
    hot_countries = get_audacious_country_ids()

    country_ids = list(set(tm_countries) | set(hot_countries))
    return country_ids


def get_unprocessed_country_ids():
    """Get country ids for audacious countries."""
    db = PostgresDb()
    sql = "set schema 'data_preparation';"
    sql += """
        with countries_with_results as (
            select
                ne_id as id
            from osm_user_contributions_per_country_per_month
            group by ne_id
        ),

        all_countries as (
            select id
            from ne_10m_admin_0_countries
            where region_id IS NOT NULL
        )

        select
          a.id
        from all_countries a
        left join countries_with_results b
            on a.id = b.id
        where b.id is NULL
    """

    data = db.retr_query(sql)
    country_ids = [item[0] for item in data]

    logger.info(f"got country ids: {country_ids}")
    return country_ids