from __future__ import annotations

import time
from typing import Literal, overload

import pandas as pd
import structlog

from agrobr.models import MetaInfo
from agrobr.utils.result import build_source_meta, finalize_result

from . import client
from .models import UF_CODES, normalize_safra_sicor, resolve_produto_sicor
from .parser import PARSER_VERSION, agregar_por_programa, agregar_por_uf, parse_credito_rural

logger = structlog.get_logger()


@overload
async def credito_rural(
    produto: str,
    safra: str | None = None,
    finalidade: str = "custeio",
    uf: str | None = None,
    agregacao: str = "municipio",
    programa: str | None = None,
    tipo_seguro: str | None = None,
    as_polars: bool = False,
    *,
    return_meta: Literal[False] = False,
) -> pd.DataFrame: ...


@overload
async def credito_rural(
    produto: str,
    safra: str | None = None,
    finalidade: str = "custeio",
    uf: str | None = None,
    agregacao: str = "municipio",
    programa: str | None = None,
    tipo_seguro: str | None = None,
    as_polars: bool = False,
    *,
    return_meta: Literal[True],
) -> tuple[pd.DataFrame, MetaInfo]: ...


async def credito_rural(
    produto: str,
    safra: str | None = None,
    finalidade: str = "custeio",
    uf: str | None = None,
    agregacao: str = "municipio",
    programa: str | None = None,
    tipo_seguro: str | None = None,
    as_polars: bool = False,
    return_meta: bool = False,
) -> pd.DataFrame | tuple[pd.DataFrame, MetaInfo]:
    t0 = time.monotonic()

    produto_sicor = resolve_produto_sicor(produto)
    safra_sicor = normalize_safra_sicor(safra) if safra else None
    cd_uf = UF_CODES.get(uf.upper()) if uf else None

    logger.info(
        "bcb_credito_rural_request",
        produto=produto,
        produto_sicor=produto_sicor,
        safra=safra,
        safra_sicor=safra_sicor,
        finalidade=finalidade,
        uf=uf,
        programa=programa,
        tipo_seguro=tipo_seguro,
    )

    dados, source_used = await client.fetch_credito_rural_with_fallback(
        finalidade=finalidade,
        produto_sicor=produto_sicor,
        safra_sicor=safra_sicor,
        cd_uf=cd_uf,
    )

    fetch_ms = int((time.monotonic() - t0) * 1000)

    attempted_sources = ["bcb_odata"]
    if source_used == "bigquery":
        attempted_sources.append("bcb_bigquery")

    t1 = time.monotonic()
    df = parse_credito_rural(dados, finalidade=finalidade)

    if uf and "uf" in df.columns:
        df = df[df["uf"] == uf.upper()].reset_index(drop=True)

    if programa and "programa" in df.columns:
        df = df[df["programa"].str.lower() == programa.lower()].reset_index(drop=True)

    if tipo_seguro and "tipo_seguro" in df.columns:
        df = df[df["tipo_seguro"].str.lower() == tipo_seguro.lower()].reset_index(drop=True)

    if agregacao == "uf":
        df = agregar_por_uf(df)
    elif agregacao == "programa":
        df = agregar_por_programa(df)

    parse_ms = int((time.monotonic() - t1) * 1000)

    source_method = "httpx" if source_used == "odata" else "bigquery"

    logger.info(
        "bcb_credito_rural_ok",
        produto=produto,
        safra=safra,
        records=len(df),
        source_used=source_used,
    )

    meta = build_source_meta(
        "bcb",
        f"{client.BASE_URL}/{client.ENDPOINT_MAP.get(finalidade.lower(), 'CusteioMunicipio')}",
        source_method,
        fetch_ms,
        parse_ms,
        df,
        PARSER_VERSION,
        schema_version="1.1",
        attempted_sources=attempted_sources,
        selected_source=f"bcb_{source_used}",
    )
    return finalize_result(df, meta, as_polars=as_polars, return_meta=return_meta)
