from __future__ import annotations

from datetime import date
from io import BytesIO
import os
from pathlib import Path
from urllib.request import urlopen

import pandas as pd
import plotly.express as px
import streamlit as st


DATASET_REPO = "SDataPro/maycee-retail-dataset"
REMOTE_DATA_ROOT = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/data"
LOCAL_DATA_ROOT = "../output/huggingface/dry_run/data"
PRODUCT_URL = "https://sdatapro.com"
DATASET_URL = f"https://huggingface.co/datasets/{DATASET_REPO}"

TABLE_FILES = {
    "brands": "train-0.parquet",
    "categories": "train-0.parquet",
    "customers": "train-0.parquet",
    "date_dim": "train-0.parquet",
    "districts": "train-0.parquet",
    "items": "train-00000.parquet",
    "products": "train-0.parquet",
    "promotions": "train-0.parquet",
    "regions": "train-0.parquet",
    "returns": "train-00000.parquet",
    "stores": "train-0.parquet",
    "suppliers": "train-0.parquet",
    "transactions": "train-00000.parquet",
}

NOISE_SCENARIOS = pd.DataFrame(
    [
        {
            "scenario": "POS outage",
            "window": "2018-07-01 to 2018-09-30",
            "visible_in_free_tier": "Yes",
            "signal": "High null-phone rate for loyalty signups in the outage window.",
        },
        {
            "scenario": "Holiday return backlog",
            "window": "2019-11-01 to 2019-12-31 purchases",
            "visible_in_free_tier": "Partly",
            "signal": "Purchases are visible; linked 2020 returns sit outside the public free range.",
        },
        {
            "scenario": "COVID return reason gap",
            "window": "2020-03-01 to 2020-06-30",
            "visible_in_free_tier": "No",
            "signal": "Premium-range scenario; listed here so public users understand the broader model.",
        },
    ]
)


st.set_page_config(
    page_title="Maycee Retail Dataset Explorer",
    page_icon="M",
    layout="wide",
)


def _data_root() -> str:
    return os.environ.get("MAYCEE_DATA_ROOT", REMOTE_DATA_ROOT).rstrip("/")


def _table_location(table: str) -> str | Path:
    root = _data_root()
    filename = TABLE_FILES[table]
    if root.startswith("http://") or root.startswith("https://"):
        return f"{root}/{table}/{filename}"
    return (Path(__file__).resolve().parent / root / table / filename).resolve()


@st.cache_data(show_spinner=False)
def read_table(table: str) -> pd.DataFrame:
    location = _table_location(table)
    if isinstance(location, Path):
        return pd.read_parquet(location)
    with urlopen(location) as response:
        return pd.read_parquet(BytesIO(response.read()))


@st.cache_data(show_spinner="Loading Maycee Retail free-tier tables...")
def load_model() -> dict[str, pd.DataFrame]:
    transactions = read_table("transactions")
    items = read_table("items")
    promotions = read_table("promotions")
    returns = read_table("returns")
    products = read_table("products")
    categories = read_table("categories")
    stores = read_table("stores")
    districts = read_table("districts")
    regions = read_table("regions")

    transactions["partition_date"] = pd.to_datetime(transactions["partition_date"]).dt.date
    items["partition_date"] = pd.to_datetime(items["partition_date"]).dt.date
    returns["partition_date"] = pd.to_datetime(returns["partition_date"]).dt.date

    category_names = categories.set_index("category_id")["name"]
    category_model = categories.assign(
        category_name=categories["name"],
        category_group=categories["parent_category_id"].map(category_names).fillna(categories["name"]),
    )[["category_id", "category_name", "category_group"]]

    product_model = products.merge(category_model, on="category_id", how="left")
    item_model = items.merge(
        product_model[
            [
                "product_id",
                "name",
                "category_name",
                "category_group",
                "brand_id",
                "supplier_id",
            ]
        ],
        on="product_id",
        how="left",
    )

    district_model = districts.rename(columns={"name": "district_name"})
    region_model = regions.rename(columns={"name": "region_name"})
    store_model = (
        stores.merge(district_model[["district_id", "region_id", "district_name"]], on="district_id", how="left")
        .merge(region_model[["region_id", "region_name"]], on="region_id", how="left")
        .rename(columns={"name": "store_name"})
    )

    transaction_model = transactions.merge(
        store_model[["store_id", "store_name", "city", "store_type", "district_name", "region_name"]],
        on="store_id",
        how="left",
    )

    return {
        "transactions": transaction_model,
        "items": item_model,
        "promotions": promotions,
        "returns": returns,
        "products": product_model,
        "stores": store_model,
        "categories": categories,
        "districts": districts,
        "regions": regions,
    }


def money(value: float) -> str:
    if abs(value) >= 1_000_000:
        return f"${value / 1_000_000:.2f}M"
    if abs(value) >= 1_000:
        return f"${value / 1_000:.1f}K"
    return f"${value:,.0f}"


def percent(value: float) -> str:
    return f"{value:.1f}%"


data = load_model()
transactions = data["transactions"]
items = data["items"]
returns = data["returns"]

min_date = transactions["partition_date"].min()
max_date = transactions["partition_date"].max()

st.title("Maycee Retail Dataset Explorer")
st.caption("Public free tier: 2017-01-01 through 2019-12-31. Synthetic retail data under CC BY 4.0.")

with st.sidebar:
    selected_range = st.date_input(
        "Date range",
        value=(min_date, max_date),
        min_value=min_date,
        max_value=max_date,
    )
    if isinstance(selected_range, tuple) and len(selected_range) == 2:
        start_date, end_date = selected_range
    else:
        start_date, end_date = min_date, max_date

    show_noise = st.toggle("Noise scenarios", value=True)

    st.link_button("Dataset", DATASET_URL, width="stretch")
    st.link_button("SDataPro", PRODUCT_URL, width="stretch")

start_date = date.fromisoformat(str(start_date))
end_date = date.fromisoformat(str(end_date))
if start_date > end_date:
    start_date, end_date = end_date, start_date

tx = transactions[
    (transactions["partition_date"] >= start_date)
    & (transactions["partition_date"] <= end_date)
].copy()
it = items[
    (items["partition_date"] >= start_date)
    & (items["partition_date"] <= end_date)
].copy()
rt = returns[
    (returns["partition_date"] >= start_date)
    & (returns["partition_date"] <= end_date)
].copy()

revenue = float(tx["total_amount"].sum())
transaction_count = int(len(tx))
line_count = int(len(it))
return_count = int(len(rt))
avg_basket = float(tx["total_amount"].mean()) if transaction_count else 0.0
gross_profit = float(it["gross_profit"].sum()) if line_count else 0.0
line_total = float(it["line_total"].sum()) if line_count else 0.0
margin = (gross_profit / line_total * 100) if line_total else 0.0

kpi_cols = st.columns(6)
kpi_cols[0].metric("Revenue", money(revenue))
kpi_cols[1].metric("Transactions", f"{transaction_count:,}")
kpi_cols[2].metric("Line Items", f"{line_count:,}")
kpi_cols[3].metric("Returns", f"{return_count:,}")
kpi_cols[4].metric("Avg Basket", money(avg_basket))
kpi_cols[5].metric("Gross Margin", percent(margin))

overview_tab, stores_tab, schema_tab = st.tabs(["Overview", "Stores", "Schema"])

with overview_tab:
    left, right = st.columns(2)

    monthly = (
        tx.assign(month=lambda frame: pd.to_datetime(frame["partition_date"]).dt.to_period("M").astype(str))
        .groupby("month", as_index=False)
        .agg(revenue=("total_amount", "sum"), transactions=("transaction_id", "count"))
    )
    fig = px.line(monthly, x="month", y="revenue", markers=True, labels={"month": "Month", "revenue": "Revenue"})
    fig.update_layout(margin=dict(l=0, r=0, t=30, b=0))
    left.plotly_chart(fig, width="stretch")

    category = (
        it.groupby("category_group", dropna=False, as_index=False)
        .agg(revenue=("line_total", "sum"))
        .sort_values("revenue", ascending=False)
        .head(12)
    )
    fig = px.bar(category, x="revenue", y="category_group", orientation="h", labels={"category_group": "Category"})
    fig.update_layout(margin=dict(l=0, r=0, t=30, b=0), yaxis={"categoryorder": "total ascending"})
    right.plotly_chart(fig, width="stretch")

    if show_noise:
        st.dataframe(NOISE_SCENARIOS, width="stretch", hide_index=True)

with stores_tab:
    left, right = st.columns(2)

    region = (
        tx.groupby("region_name", dropna=False, as_index=False)
        .agg(revenue=("total_amount", "sum"), transactions=("transaction_id", "count"))
        .sort_values("revenue", ascending=False)
    )
    fig = px.bar(region, x="region_name", y="revenue", labels={"region_name": "Region", "revenue": "Revenue"})
    fig.update_layout(margin=dict(l=0, r=0, t=30, b=0))
    left.plotly_chart(fig, width="stretch")

    store = (
        tx.groupby(["store_name", "city", "region_name"], dropna=False, as_index=False)
        .agg(transactions=("transaction_id", "count"), revenue=("total_amount", "sum"))
        .sort_values("transactions", ascending=False)
        .head(20)
    )
    fig = px.bar(store, x="transactions", y="store_name", color="region_name", orientation="h")
    fig.update_layout(margin=dict(l=0, r=0, t=30, b=0), yaxis={"categoryorder": "total ascending"})
    right.plotly_chart(fig, width="stretch")

    channel = (
        tx.groupby(["channel", "payment_method"], dropna=False, as_index=False)
        .agg(transactions=("transaction_id", "count"), revenue=("total_amount", "sum"))
        .sort_values("transactions", ascending=False)
    )
    st.dataframe(channel, width="stretch", hide_index=True)

with schema_tab:
    table = st.selectbox("Table", sorted(TABLE_FILES))
    sample = read_table(table)
    schema = pd.DataFrame(
        {
            "column": sample.columns,
            "dtype": [str(dtype) for dtype in sample.dtypes],
            "non_null": [int(sample[column].notna().sum()) for column in sample.columns],
        }
    )
    st.dataframe(schema, width="stretch", hide_index=True)
    st.dataframe(sample.head(25), width="stretch", hide_index=True)