from pathlib import Path

import polars as pl


def load_data(parquet_path: Path) -> pl.DataFrame:
    return pl.read_parquet(parquet_path)


def get_total_cost(df: pl.DataFrame) -> float:
    return df["cost"].sum()


def get_cost_by_provider(df: pl.DataFrame) -> pl.DataFrame:
    return df.group_by("cloud_provider").agg(pl.col("cost").sum().alias("total_cost"))


def with_month_col(df: pl.DataFrame) -> pl.DataFrame:
    if "month" not in df.columns:
        df = df.with_columns(
            pl.col("timestamp")
            .str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S")
            .dt.strftime("%Y-%m")
            .alias("month")
        )
    return df


def get_cost_by_month(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)
    return df.group_by("month").agg(pl.col("cost").sum().alias("total_cost"))


def get_cost_by_provider_month(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)
    return df.group_by(["cloud_provider", "month"]).agg(
        pl.col("cost").sum().alias("total_cost")
    )


def get_min_max_cost_month_by_provider(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)

    cost_by_month = df.group_by(["cloud_provider", "month"]).agg(
        pl.col("cost").sum().alias("total_cost")
    )

    min_cost = (
        cost_by_month.sort("total_cost")
        .group_by("cloud_provider")
        .first()
        .rename({"month": "min_month", "total_cost": "min_total_cost"})
    )

    max_cost = (
        cost_by_month.sort("total_cost", descending=True)
        .group_by("cloud_provider")
        .first()
        .rename({"month": "max_month", "total_cost": "max_total_cost"})
    )

    return min_cost.join(max_cost, on="cloud_provider")


def get_line_items_by_provider_month(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)
    return df.group_by(["cloud_provider", "month"]).agg(pl.len().alias("line_items"))


def get_min_max_line_items_month_by_provider(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)

    line_items_by_month = df.group_by(["cloud_provider", "month"]).agg(
        pl.len().alias("line_items")
    )

    min_items = (
        line_items_by_month.sort("line_items")
        .group_by("cloud_provider")
        .first()
        .rename({"month": "min_month", "line_items": "min_line_items"})
    )

    max_items = (
        line_items_by_month.sort("line_items", descending=True)
        .group_by("cloud_provider")
        .first()
        .rename({"month": "max_month", "line_items": "max_line_items"})
    )

    return min_items.join(max_items, on="cloud_provider")


def get_max_payloads_and_connections(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)
    grouped = df.group_by(["cloud_provider", "month"]).agg(
        [pl.len().alias("line_items")]
    )

    grouped = grouped.with_columns(
        (pl.col("line_items") * 1000).alias("estimated_bytes")
    )

    grouped = grouped.with_columns(
        (pl.col("estimated_bytes") / (5 * 1024 * 1024))
        .ceil()
        .cast(pl.Int64)
        .alias("estimated_payloads")
    )

    grouped = grouped.with_columns(
        (pl.col("estimated_payloads") / 15)
        .ceil()
        .cast(pl.Int64)
        .alias("required_connections")
    )

    return grouped.sort(by=["cloud_provider", "month"])


def get_peak_payload_month_by_provider(df: pl.DataFrame) -> pl.DataFrame:
    df = with_month_col(df)

    grouped = (
        df.group_by(["cloud_provider", "month"])
        .agg(pl.len().alias("line_items"))
        .with_columns(
            (pl.col("line_items") * 1000 / (1024 * 1024)).alias("estimated_mb")
        )
        .with_columns(
            (pl.col("estimated_mb") / 5)
            .ceil()
            .cast(pl.Int64)
            .alias("estimated_payloads")
        )
        .with_columns(
            (pl.col("estimated_payloads") / 15)
            .ceil()
            .cast(pl.Int64)
            .alias("required_connections")
        )
    )

    peak = (
        grouped.sort("line_items", descending=True).group_by("cloud_provider").first()
    )

    return peak


def get_null_counts(df: pl.DataFrame) -> dict:
    return {col: df[col].null_count() for col in df.columns}


def get_service_usage_resources(df: pl.DataFrame) -> pl.DataFrame:
    filtered = df.filter(pl.col("resource_type") == "service-usage")
    return filtered.select(["cloud_provider", "resource_id"]).unique()
