"""
Helper functions to download official statistics from IBGE SIDRA
"""

import ssl  # Builtin
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Literal
from urllib.error import HTTPError
from zipfile import ZipFile

import pandas as pd
import requests
import urllib3
from pysus.data.local import ParquetSet
from pysus.ftp.databases.ibge_datasus import IBGEDATASUS

# requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS = 'ALL:@SECLEVEL=1'


APIBASE = "https://servicodados.ibge.gov.br/api/v3/"

ibge = IBGEDATASUS().load()


def get_sidra_table(
    table_id,
    territorial_level,
    geocode="all",
    period=None,
    variables=None,
    classification=None,
    categories=None,
    format=None,
    decimals=None,
    headers=None,
):
    """
    Wrapper for the SIDRA API.
    More information here: http://apisidra.ibge.gov.br/home/ajuda
    :param table_id: código da tabela de onde se deseja extrair os dados.
        código pode ser obtido aqui: https://sidra.ibge.gov.br/acervo#/S/Q
    :param territorial_level:
        1 – Brasil,
        2 – Grande Região,
        3 – Unidade da Federação,
        6 – Município, etc
    :param geocode: geocódigo do IBGE: 3304557,3550308 – especifica os
        municípios do Rio de Janeiro e São Paulo.
    all – especifica todos os municípios. in n3 11,12 - especifica os
        municípios contidos nas Unidades da Federação Rondônia e Acre.

    :param period:  Os períodos podem ser especificados de forma avulsa,
        separados por vírgula (,), em faixas, separados por traço (-), ou de
        ambas as formas. Um período pode ter o formato AAAA, de 4 dígitos, que
        representa um ano, ou o formato AAAADD, de 6 dígitos, onde AAAA
        representa um ano e DD seu correspondente mês (01 a 12),
        trimestre (01 a 04), semestre (01 a 02), etc, de acordo com a
        periodicidade de divulgação dos dados da tabela.

        Exemplo 1: /p/2008,2010-2012 – 2008, e 2010 a 2012.

        Exemplo 2: /p/201101-201112,201204,201208 – janeiro a dezembro de 2012,
            abril de 2012 e agosto de 2012.

        O parâmetro p pode ser seguido pela constante all para especificar
        todos os períodos disponíveis.

        Exemplo 3: /p/all

        O parâmetro p pode ser seguido pela constante first e um número de
        períodos, indicando os primeiros períodos da lista de períodos
        disponíveis (períodos mais antigos).
        O número de períodos pode ser omitido quando se tratar de apenas um
        aperíodo.

        Exemplo 4: /p/first 12

        Exemplo 5: /p/first

        O parâmetro p pode ser seguido pela constante last e um número de
        períodos, indicando os últimos períodos da série (períodos mais
        recentes). O número de períodos pode ser omitido quando se tratar de
        apenas um período.

        Exemplo 6: /p/last 12

        Exemplo 7: /p/last (default, quando não especificado o parâmetro p)
    :param variables: As variáveis são especificadas através de seus códigos,
        separados por vírgula (,).
        A lista de variáveis pode incluir também as variáveis de percentual
        geradas automaticamente pelo Sidra (são variáveis cujos códigos são
        superiores a 1.000.000).

        Exemplo 1: /v/63,69 – especifica o percentual no mês e o percentual
        acumulado no ano do IPCA.

        O parâmetro v pode ser seguido pela constante all para especificar
        todas as variáveis da tabela, inclusive as variáveis de percentual
        geradas automaticamente pelo Sidra.

        Exemplo 2: /v/all

        O parâmetro v pode ser seguido pela constante allxp para especificar
        todas as variáveis da tabela, exceto as variáveis de percentual
        geradas automaticamente pelo Sidra.

        Exemplo 3: /v/allxp (default, quando não especificado o parâmetro v)
    :param classification: informa o código de uma das classificações da tabela
        Como exemplos, temos 1 – Situação do domicílio, 2 – Sexo, 81 – Produto
        da lavoura temporária, etc.
    :param categories:  As categorias são especificadas através de seus
        códigos, de forma individual ou para compor uma soma, separadas por
        vírgula (,).
        As categorias que compõem a soma devem ser separadas por espaço.

        Exemplo 1: /c81/2692,2702,2694 2695 – especifica os produtos da lavoura
        temporária arroz, feijão e (batata doce + batata inglesa)
    :param format:
    :param decimals:
    :param headers: `y` para receber o header (valor default, caso o parâmetro
        h não seja especificado). `n` para não receber o header.
    :return:
    """
    base_url = "https://apisidra.ibge.gov.br/values"
    query = f"/t/{table_id}/n{territorial_level}/{geocode}"
    if period is not None:
        query += f"/p/{period}"
    if variables is not None:
        query += f"/v/{variables}"
    if classification is not None:
        query += f"/c{classification}"
    if categories is not None:
        query += f"/{categories}"
    if format is not None:
        query += f"/f/{format}"
    if decimals is not None:
        query += f"/d/{decimals}"
    if headers is not None:
        query += f"/h/{headers}"

    url = base_url + query

    print(f"Requesting data from {url}")
    try:
        with get_legacy_session() as s, s.get(url) as response:
            df = pd.DataFrame(response.json())
    except HTTPError:
        response = requests.get(url)
        print(f"Consulta falhou: {response.text}")
        return None
    return df


def list_agregados(**kwargs):
    """
    Lista de agregados agrupados por pesquisa. Veja
    https://servicodados.ibge.gov.br/api/docs/agregados?versao=3#api-Agregados-agregadosGet # noqa
    para maiores detalhes
    :param kwargs: parâmetros válidos: período, assunto, classificacao, periodicidade,nivel.
    :return: Dataframe
    """
    url = APIBASE + "agregados?"
    url += "&".join([f"{k}={v}" for k, v in kwargs.items()])
    print(f"Fetching Data groupings from {url}")
    try:
        with get_legacy_session() as s, s.get(url) as response:
            table = pd.DataFrame(response.json())
    except requests.exceptions.SSLError as e:
        print(f"Failed fetching aggregates: {e}")
        return pd.DataFrame()
    return table


def localidades_por_agregado(agregado: int, nivel: str):
    """
    Obtém as localidades associadas ao agregado de acordo com um ou mais níveis
    geográficos.
    :param agregado: codigo numérico do agregado
    :param nivel: Identificador do nível geográfico ao qual pertence as
    localidades. Pode conter um ou mais níveis
    delimitados pelo caracter | (pipe). p.ex. N7|N6
    :return:
    """
    url = APIBASE + f"agregados/{agregado}/localidades/{nivel}"
    try:
        with get_legacy_session() as s, s.get(url) as response:
            table = pd.DataFrame(response.json())
    except Exception as e:
        print(f"Could not download from {url}\n{e}")
        return None
    return table


def metadados(agregado: int):
    """
    Obtém os metadados associados ao agregado

    :param agregado: Identificador do agregado
    """
    url = APIBASE + f"agregados/{agregado}/metadados"
    try:
        with get_legacy_session() as s, s.get(url) as response:
            data = response.json()
    except Exception as e:
        print(f"Could not download from {url}\n{e}")
        return None
    return data


def lista_periodos(agregado: int):
    """
    Obtém os períodos associados ao agregado
    :param agregado:
    :return: pd.DataFrame com os períodos de atualização
    """
    url = APIBASE + f"agregados/{agregado}/periodos"
    try:
        with get_legacy_session() as s, s.get(url) as response:
            table = pd.DataFrame(response.json())
    except Exception:
        return None
    return table


class FetchData:
    """
    Obtém o conjunto de variáveis a partir do identificador do agregado,
    períodos pesquisados e identificador das variáveis.

    :param agregado: Identificador do agregado.
    :param periodos: Período do qual se deseja obter os resultados. Consulte os
        identificadores dos períodos na Base de Identificadores. Informe
        valores negativos para obter os últimos resultados. Pode conter um ou
        mais períodos delimitados pelo caractere | (pipe).
    :param variavel: Um ou mais identificadores de variável separados pelo
        caractere | (pipe). Caso omitido, assume o valor allxp, que retorna
        quaisquer variáveis relacionadas ao agregado. Para saber mais sobre as
        variáveis de cada agregado, acesse seus respectivos metadados.
    :kwargs: Parâmetros adicionais:
        - **localidades**: Uma ou mais localidades delimitadas pelo caractere |
            (pipe). No caso do Brasil, o identificador é BR. Para qualquer
            outra localidade que NÃO seja Brasil, essa deve seguir o padrão
            N<NIVEL_GEOGRAFICO>[<LOCALIDADE>], em que <LOCALIDADE> pode ser uma
            ou mais localidades separadas por vírgula. É possível ainda
            generalizar o resultado, informando a classe da localidade,
            conforme os exemplos a seguir:

            https://servicodados.ibge.gov.br/api/v3/agregados/1705/variaveis?localidades=N7  # noqa

            Obtém os resultados referentes às variáveis do agregado 1705 cujas
            localidades sejam regiões metropolitanas (N7).

            https://servicodados.ibge.gov.br/api/v3/agregados/1705/variaveis?localidades=N7[3501,3301]  # noqa

            Obtém os resultados referentes às variáveis do agregado 1705 cujas
            localidades sejam as regiões metropolitanas (N7) de São Paulo e Rio
            de Janeiro (3501,3301). Observe que 3501 e 3301 são os
            identificadores das regiões metropolitanas de São Paulo e Rio de
            Janeiro. Não podem ser confundidos com os identificadores dos
            municípios de São Paulo/SP e Rio de Janeiro/RJ, que são 3550308
            e 3304557, respectivamente.

        - **classificacao**: Além de estar relacionado a uma dada localidade e
            a um determinado período, os resultados das variáveis podem estar
            relacionados a outros conjuntos de dados, que na nomenclatura do
            SIDRA recebem o nome de classificação. Como exemplo, considere o
            agregado Produção, venda, valor da produção e área colhida da
            lavoura temporária nos estabelecimentos agropecuários. Além da
            localidade e do período, os resultados produzidos por esse agregado
            referem-se aos produtos produzidos, condição do produtor, grupos de
            atividades econômicas, grupos de área, grupos de área colhida e
            pronafiano, que são as classificações do agregado. Para conhecer as
            classificações de cada agregado, acesse seus respectivos metadados.
            Aos componentes da classificação, dá-se o nome de categoria. Na
            prática, você fará uso das classificações para restringir a
            consulta, conforme os exemplos a seguir:

            https://servicodados.ibge.gov.br/api/v3/agregados/1712/variaveis?classificacao=226[4844]&localidades=BR  # noqa

            Obtém os resultados referentes às variáveis do agregado 1712 cujo
            produto produzido (226) seja abacaxi (4844) no Brasil (BR).

            https://servicodados.ibge.gov.br/api/v3/agregados/1712/variaveis?classificacao=226[4844]|218[4780]&localidades=BR  # noqa

            Obtém os resultados referentes às variáveis do agregado 1712 cujo
            produto produzido (226) seja abacaxi (4844) e cuja condição do
            produtor (218) seja proprietário (4780) no Brasil (BR).

        - **view**: Modo de visualização. Caso deseje que a resposta seja
            renderizada usando notação OLAP, configure esse parâmetro com o
            valor OLAP -
            https://servicodados.ibge.gov.br/api/v3/agregados/1705/variaveis?view=OLAP&localidades=BR.  # noqa
            A outra opção é configurar esse parâmetro com o valor flat. No modo
            flat, o primeiro elemento do array são metadados, de forma que os
            resultados vêm a partir do segundo elemento.
    """

    def __init__(
        self, agregado: int, periodos: str, variavel: str = "allxp", **kwargs
    ):
        self.url = (
            APIBASE
            + f"agregados/{agregado}/periodos/{periodos}/variaveis/{variavel}?"
        )
        self.url += "&".join([f"{k}={v}" for k, v in kwargs.items()])
        self.JSON = None
        self._fetch_JSON()

    def _fetch_JSON(self):
        try:
            print(f"Fetching {self.url}")
            with get_legacy_session() as s, s.get(self.url) as response:
                self.JSON = response.json()
        except Exception as e:
            print("Couldn't download data:", e, sep="\n")

    def to_dataframe(self):
        return pd.DataFrame(self.JSON)


"""
HTTPSConnectionPool(host='servicodados.ibge.gov.br', port=443):
    Max retries exceeded with url:
    /api/v3/agregados/{...}
    Caused by SSLError(
        SSLError(1, '[SSL: UNSAFE_LEGACY_RENEGOTIATION_DISABLED]
        unsafe legacy renegotiation disabled (_ssl.c:1129)'

SOLUTION: https://github.com/scrapy/scrapy/issues/5491#issuecomment-1241862323
"""


class CustomHttpAdapter(requests.sessions.HTTPAdapter):
    # "Transport adapter" that allows us to use custom ssl_context.

    def __init__(self, ssl_context=None, **kwargs):
        self.ssl_context = ssl_context
        super().__init__(**kwargs)

    def init_poolmanager(self, connections, maxsize, block=False):
        self.poolmanager = urllib3.poolmanager.PoolManager(
            num_pools=connections,
            maxsize=maxsize,
            block=block,
            ssl_context=self.ssl_context,
        )


def get_legacy_session():
    ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
    ctx.options |= 0x4  # OP_LEGACY_SERVER_CONNECT
    session = requests.session()
    session.mount("https://", CustomHttpAdapter(ctx))
    return session


def get_population(
    year: int,
    source: Literal["POP", "censo", "POPTCU", "projpop"] = "POPTCU",
    censo_data: Literal["ALF", "ESCA", "ESCB", "IDOSO", "RENDA"] = "ALF",
) -> pd.DataFrame:
    """
    Get population data from IBGE as shared by DATASUS
    :param year: year of the data
    :param source:
         "POP" - 1992-presente: Estimativas populacionais estratificadas por
                 idade e sexo.
         "censo" - 1991, 2000 e 2010: Censos Demográficos
         "POPTCU" - 1992-presente: Estimativas populacionais enviadas para o
                    TCU, estratificadas por idade e sexo pelo MS/SGEP/Datasus.
         "projpop": Estimativas preliminares para os anos intercensitários dos
                    totais populacionais, estratificadas por idade e sexo pelo
                    MS/SGEP/Datasus.
    :param censo_data:
        "ALF": Censo Demográfico
        "ESCA": Censo Escolar da Educação Básica
        "ESCB": Censo Escolar da Educação Superior
        "IDOSO": População de pessoas com 65 anos ou mais
        "RENDA": População de pessoas de acordo com a renda familiar
    :return: DataFrame with population data
    """

    files = ibge.get_files(year=int(year), source=source)

    if files == []:
        return pd.DataFrame()

    if source == "censo":
        opts = ["ALF", "ESCA", "ESCB", "IDOSO", "RENDA"]
        if not censo_data or censo_data not in opts:
            raise ValueError(
                f"Incorrect 'censo_data' parameter. Options: {opts}"
            )
        file = [f for f in files if censo_data in f.name][0].download()
    else:
        file = files[0].download()

    if isinstance(file, ParquetSet):
        return file.to_dataframe()

    file = Path(str(file))

    if file.suffix.lower() == ".zip":
        return _unzip_to_dataframe(str(file))
    else:
        raise NotImplementedError(f"Unkown file type '{file.suffix}'")


def _unzip_to_dataframe(file: str) -> pd.DataFrame:
    zip_file = ZipFile(file)  # pyright: ignore
    with TemporaryDirectory() as tempdir:
        for file in zip_file.namelist():
            if file.lower().endswith(".csv"):
                return pd.read_csv(zip_file.extract(file, tempdir))

            if file.lower().endswith((".dbf", ".dbc")):
                return ParquetSet(
                    zip_file.extract(file, tempdir)
                ).to_dataframe()

        raise ValueError(f"No data found in {zip_file}")
