import io
import zipfile
import datetime
from pathlib import Path
from typing import Literal

import requests
import pandas as pd
import polars as pl


# metadata
# https://github.com/codigoquant/b3fileparser/blob/main/b3fileparser/b3_meta_data.py
# thanks to codigoquant https://github.com/codigoquant
FIELD_SIZES = {
    "TIPO_DE_REGISTRO": 2,
    "DATA_DO_PREGAO": 8,
    "CODIGO_BDI": 2,
    "CODIGO_DE_NEGOCIACAO": 12,
    "TIPO_DE_MERCADO": 3,
    "NOME_DA_EMPRESA": 12,
    "ESPECIFICACAO_DO_PAPEL": 10,
    "PRAZO_EM_DIAS_DO_MERCADO_A_TERMO": 3,
    "MOEDA_DE_REFERENCIA": 4,
    "PRECO_DE_ABERTURA": 13,
    "PRECO_MAXIMO": 13,
    "PRECO_MINIMO": 13,
    "PRECO_MEDIO": 13,
    "PRECO_ULTIMO_NEGOCIO": 13,
    "PRECO_MELHOR_OFERTA_DE_COMPRA": 13,
    "PRECO_MELHOR_OFERTA_DE_VENDAS": 13,
    "NUMERO_DE_NEGOCIOS": 5,
    "QUANTIDADE_NEGOCIADA": 18,
    "VOLUME_TOTAL_NEGOCIADO": 18,
    "PRECO_DE_EXERCICIO": 13,
    "INDICADOR_DE_CORRECAO_DE_PRECOS": 1,
    "DATA_DE_VENCIMENTO": 8,
    "FATOR_DE_COTACAO": 7,
    "PRECO_DE_EXERCICIO_EM_PONTOS": 13,
    "CODIGO_ISIN": 12,
    "NUMERO_DE_DISTRIBUICAO": 3,
}

DATE_COLUMNS = (
    "DATA_DO_PREGAO",
    "DATA_DE_VENCIMENTO",
)

FLOAT32_COLUMNS = (
    "PRECO_DE_ABERTURA",
    "PRECO_MAXIMO",
    "PRECO_MINIMO",
    "PRECO_MEDIO",
    "PRECO_ULTIMO_NEGOCIO",
    "PRECO_MELHOR_OFERTA_DE_COMPRA",
    "PRECO_MELHOR_OFERTA_DE_VENDAS",
    "PRECO_DE_EXERCICIO",
    "PRECO_DE_EXERCICIO_EM_PONTOS",
)

FLOAT64_COLUMNS = (
    "VOLUME_TOTAL_NEGOCIADO",
    "QUANTIDADE_NEGOCIADA",
)

UINT32_COLUMNS = (
    "FATOR_DE_COTACAO",
    "PRAZO_EM_DIAS_DO_MERCADO_A_TERMO",
    "NUMERO_DE_NEGOCIOS",
    "NUMERO_DE_DISTRIBUICAO",
    "TIPO_DE_REGISTRO",
)

STRING_COLUMNS = (
    "CODIGO_DE_NEGOCIACAO",
    "NOME_DA_EMPRESA",
    "CODIGO_ISIN",
)

CATEGORY_COLUMNS = (
    "INDICADOR_DE_CORRECAO_DE_PRECOS",
    "TIPO_DE_MERCADO",
    "CODIGO_BDI",
    "MOEDA_DE_REFERENCIA",
    "ESPECIFICACAO_DO_PAPEL",
)

MARKETS = {
    "010": "VISTA",
    "012": "EXERCICIO_DE_OPCOES_DE_COMPRA",
    "013": "EXERCÍCIO_DE_OPCOES_DE_VENDA",
    "017": "LEILAO",
    "020": "FRACIONARIO",
    "030": "TERMO",
    "050": "FUTURO_COM_RETENCAO_DE_GANHO",
    "060": "FUTURO_COM_MOVIMENTACAO_CONTINUA",
    "070": "OPCOES_DE_COMPRA",
    "080": "OPCOES_DE_VENDA",
}

INDOPC = {
    "0": "0",
    "1": "US$",
    "2": "TJLP",
    "8": "IGPM",
    "9": "URV",
}

CODBDI = {
    "00": "0",
    "02": "LOTE_PADRAO",
    "05": "SANCIONADAS PELOS REGULAMENTOS BMFBOVESPA",
    "06": "CONCORDATARIAS",
    "07": "RECUPERACAO_EXTRAJUDICIAL",
    "08": "RECUPERAÇÃO_JUDICIAL",
    "09": "REGIME_DE_ADMINISTRACAO_ESPECIAL_TEMPORARIA",
    "10": "DIREITOS_E_RECIBOS",
    "11": "INTERVENCAO",
    "12": "FUNDOS_IMOBILIARIOS",
    "13": "13",
    "14": "CERT.INVEST/TIT.DIV.PUBLICA",
    "18": "OBRIGACÕES",
    "22": "BÔNUS(PRIVADOS)",
    "26": "APOLICES/BÔNUS/TITULOS PUBLICOS",
    "32": "EXERCICIO_DE_OPCOES_DE_COMPRA_DE_INDICES",
    "33": "EXERCICIO_DE_OPCOES_DE_VENDA_DE_INDICES",
    "34": "34",
    "35": "35",
    "36": "36",
    "37": "37",
    "38": "EXERCICIO_DE_OPCOES_DE_COMPRA",
    "42": "EXERCICIO_DE_OPCOES_DE_VENDA",
    "46": "LEILAO_DE_NAO_COTADOS",
    "48": "LEILAO_DE_PRIVATIZACAO",
    "49": "LEILAO_DO_FUNDO_RECUPERACAO_ECONOMICA_ESPIRITO_SANTO",
    "50": "LEILAO",
    "51": "LEILAO_FINOR",
    "52": "LEILAO_FINAM",
    "53": "LEILAO_FISET",
    "54": "LEILAO_DE_ACÕES_EM_MORA",
    "56": "VENDAS_POR_ALVARA_JUDICIAL",
    "58": "OUTROS",
    "60": "PERMUTA_POR_ACÕES",
    "61": "META",
    "62": "MERCADO_A_TERMO",
    "66": "DEBENTURES_COM_DATA_DE_VENCIMENTO_ATE_3_ANOS",
    "68": "DEBENTURES_COM_DATA_DE_VENCIMENTO_MAIOR_QUE_3_ANOS",
    "70": "FUTURO_COM_RETENCAO_DE_GANHOS",
    "71": "MERCADO_DE_FUTURO",
    "74": "OPCOES_DE_COMPRA_DE_INDICES",
    "75": "OPCOES_DE_VENDA_DE_INDICES",
    "78": "OPCOES_DE_COMPRA",
    "82": "OPCOES_DE_VENDA",
    "83": "BOVESPAFIX",
    "84": "SOMA_FIX",
    "90": "TERMO_VISTA_REGISTRADO",
    "96": "MERCADO_FRACIONARIO",
    "99": "TOTAL_GERAL",
}


def _get_txt_from_zip(zip_file: zipfile.ZipFile) -> io.BytesIO:
    file_name = zip_file.namelist()[0]
    with zip_file.open(file_name) as f:
        return f.read()


def _requests_get_txt(date: datetime.date, raise_ssl_error: bool = False) -> io.BytesIO:
    url = f'https://bvmf.bmfbovespa.com.br/InstDados/SerHist/COTAHIST_D{date.strftime("%d%m%Y")}.ZIP'
    r = requests.get(url, verify=raise_ssl_error)
    r.raise_for_status()
    with zipfile.ZipFile(io.BytesIO(r.content)) as thezip:
        return _get_txt_from_zip(thezip)


def _read_bytes(bytes_data: io.BytesIO) -> pl.DataFrame:
    df = pl.read_csv(
        bytes_data,
        has_header=False,
        new_columns=["full_str"],
        encoding="latin1",
        truncate_ragged_lines=True,
    )

    slices = {}
    start = 0
    for col_name, width in FIELD_SIZES.items():
        slices[col_name] = (start, width)
        start += width

    df = (
        df.with_columns(
            [
                pl.col("full_str")
                .str.slice(slice_tuple[0], slice_tuple[1])
                .str.strip_chars()
                .alias(col)
                for col, slice_tuple in slices.items()
            ]
        )
        .drop(["full_str"])
        .slice(1, -1)  # dropa cabeçalho erodapé
        .with_columns(
            pl.col(DATE_COLUMNS).str.to_date(format="%Y%m%d"),
            pl.col(FLOAT32_COLUMNS).cast(pl.Float32) / 100,
            pl.col(FLOAT64_COLUMNS).cast(pl.Float64),
            pl.col(UINT32_COLUMNS).cast(pl.UInt32, strict=False),
            pl.col("CODIGO_BDI").map_elements(
                lambda x: CODBDI.get(x, x), return_dtype=str
            ),
            pl.col("TIPO_DE_MERCADO").map_elements(
                lambda x: MARKETS.get(x, x), return_dtype=str
            ),
            pl.col("INDICADOR_DE_CORRECAO_DE_PRECOS").map_elements(
                lambda x: INDOPC.get(x, x), return_dtype=str
            ),
        )
    )

    return df


def _check_engine(
    engine: Literal["pandas", "polars"] = "pandas",
) -> Literal["pandas", "polars"]:
    if engine not in ["pandas", "polars"]:
        raise ValueError("engine must be either 'pandas' or 'polars'")
    return engine


def get(
    date: datetime.date,
    engine: Literal["pandas", "polars"] = "pandas",
    raise_ssl_error: bool = False,
) -> pd.DataFrame | pl.DataFrame:
    """Obtém dados históricos da B3 para uma data específica via download.

    Params
    ----------
    date : datetime.date
        Data para a qual os dados devem ser obtidos.
    engine : {"pandas", "polars"}, default="pandas"
        Engine de processamento a ser utilizado.
    raise_ssl_error : bool, default=False
        Se True, levanta erros de SSL durante o download.

    Returns
    -------
    pandas.DataFrame ou polars.DataFrame
        DataFrame contendo os dados históricos da B3.
    """
    _check_engine(engine)
    bytes_data = _requests_get_txt(date, raise_ssl_error=raise_ssl_error)
    df_polars = _read_bytes(bytes_data)

    if engine == "pandas":
        return df_polars.to_pandas()
    return df_polars


def read_bytes(
    data: bytes | io.BytesIO,
    engine: Literal["pandas", "polars"] = "pandas",
) -> pd.DataFrame | pl.DataFrame:
    """Lê dados históricos da B3 a partir de bytes ou BytesIO.

    Params
    ----------
    data : bytes ou io.BytesIO
        Dados em formato bytes ou BytesIO contendo o arquivo da B3.
    engine : {"pandas", "polars"}, default="pandas"
        Engine de processamento a ser utilizado.

    Returns
    -------
    pandas.DataFrame ou polars.DataFrame
        DataFrame contendo os dados históricos da B3.
    """
    _check_engine(engine)
    df_polars = _read_bytes(data)

    if engine == "pandas":
        return df_polars.to_pandas()
    return df_polars


def read_zip(
    path: str | Path,
    engine: Literal["pandas", "polars"] = "pandas",
) -> pd.DataFrame | pl.DataFrame:
    """Lê dados históricos da B3 a partir de um arquivo ZIP.

    Params
    ----------
    path : str ou Path
        Caminho para o arquivo ZIP contendo os dados da B3.
    engine : {"pandas", "polars"}, default="pandas"
        Engine de processamento a ser utilizado.

    Returns
    -------
    pandas.DataFrame ou polars.DataFrame
        DataFrame contendo os dados históricos da B3.
    """
    _check_engine(engine)
    with zipfile.ZipFile(path) as thezip:
        df_polars = _read_bytes(io.BytesIO(_get_txt_from_zip(thezip)))

    if engine == "pandas":
        return df_polars.to_pandas()
    return df_polars


def read_txt(
    path: str | Path,
    engine: Literal["pandas", "polars"] = "pandas",
) -> pd.DataFrame | pl.DataFrame:
    """Lê dados históricos da B3 a partir de um arquivo TXT.

    Params
    ----------
    path : str ou Path
        Caminho para o arquivo TXT contendo os dados da B3.
    engine : {"pandas", "polars"}, default="pandas"
        Engine de processamento a ser utilizado.

    Returns
    -------
    pandas.DataFrame ou polars.DataFrame
        DataFrame contendo os dados históricos da B3.
    """
    _check_engine(engine)
    if isinstance(path, str):
        path = Path(path)

    if not path.suffix.lower() == ".txt":
        raise ValueError("path must be a .txt file")

    with open(path, "rb") as f:
        df_polars = _read_bytes(io.BytesIO(f.read()))

    if engine == "pandas":
        return df_polars.to_pandas()
    return df_polars
