from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any, Dict, Iterable, List

import logging
import time
import pyarrow as pa
from deltalake import DeltaTable
from deltalake.writer import write_deltalake

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SourceInfo:
    accountID: str
    projectID: str
    tenantID: str
    applicationID: str
    biztype: str


def _ensure_list_of_dicts(v: Any) -> List[Dict[str, Any]]:
    if v is None:
        return []
    if isinstance(v, list):
        return [x for x in v if isinstance(x, dict)]
    return []


def _normalize_account_id(account_id: Any) -> str:
    """Coerce VM account ID to string for stable Parquet schema (Glue/Athena varchar)."""
    if account_id is None or account_id == "":
        raise ValueError("accountID is required")
    return str(account_id)


def _to_arrow_table(rows: List[Dict[str, Any]]) -> pa.Table:
    normalized: List[Dict[str, Any]] = []
    for row in rows:
        item = dict(row)
        if "accountID" in item:
            item["accountID"] = _normalize_account_id(item["accountID"])
        normalized.append(item)
    return pa.Table.from_pylist(normalized)


def _delta_table_exists(path: str) -> bool:
    """Return True if a Delta table exists at the given path."""
    try:
        DeltaTable(path)
        return True
    except Exception:
        return False


def delete_date_from_deltalake(
    *,
    delta_root: str,
    date: str,
) -> None:
    """
    Delete all data for a given date from all Delta Lake tables.
    This should be called before processing a date to ensure clean data.
    """
    t0 = time.time()
    table_names = [
        "sources",
        "series_count_by_metric_name",
        "series_count_by_label_name",
        "series_count_by_label_value_pair",
    ]

    for table_name in table_names:
        table_path = f"{delta_root.rstrip('/')}/{table_name}"
        dt = DeltaTable(table_path)
        if dt is not None:
            try:
                delete_predicate = f"date = '{date}'"
                dt.delete(delete_predicate)
                logger.info("Deleted date %s from table %s", date, table_name)
            except Exception as e:
                logger.warning("Failed to delete date %s from table %s: %s", date, table_name, e)

    logger.info(
        "delete_date_from_deltalake done for date=%s (%dms)",
        date,
        int((time.time() - t0) * 1000),
    )


def write_tsdb_status_to_deltalake(
    *,
    delta_root: str,
    date: str,
    source: SourceInfo,
    tsdb_status: Dict[str, Any],
) -> None:
    """
    Split the response from /api/v1/status/tsdb into 4 Delta Lake tables:
      - sources: source information (accountID/projectID/tenantID/applicationID/biztype + date)
      - series_count_by_metric_name
      - series_count_by_label_name
      - series_count_by_label_value_pair

    Each detail table includes source + date for easy joining and tracing.
    """
    t0 = time.time()
    base_cols = {
        "date": date,
        **asdict(source),
    }

    data = (tsdb_status or {}).get("data", {}) if isinstance(tsdb_status, dict) else {}

    # 1) sources table: one record per account/project fetch per day
    t_sources = time.time()
    sources_row = {
        **base_cols,
        "isPartial": tsdb_status.get("isPartial"),
        "totalSeries": data.get("totalSeries"),
        "totalLabelValuePairs": data.get("totalLabelValuePairs"),
    }
    sources_path = f"{delta_root.rstrip('/')}/sources"
    write_deltalake(
        sources_path,
        _to_arrow_table([sources_row]),
        mode="append",
        partition_by=["biztype", "applicationID", "date"],
        schema_mode="merge",
    )
    logger.info(
        "deltalake write sources done: %s (%dms)",
        sources_path,
        int((time.time() - t_sources) * 1000),
    )

    # 2) Three detail tables: generic expansion
    def write_name_value_table(table_name: str, items: Iterable[Dict[str, Any]]) -> None:
        t_table = time.time()
        rows: List[Dict[str, Any]] = []
        for it in items:
            name = it.get("name")
            value = it.get("value")
            if name is None or value is None:
                continue
            row = {**base_cols, "name": name, "value": value}

            # seriesCountByMetricName may include extra fields in newer API responses
            if table_name == "series_count_by_metric_name":
                if "requestsCount" in it:
                    row["requestsCount"] = it.get("requestsCount")
                if "lastRequestTimestamp" in it:
                    row["lastRequestTimestamp"] = it.get("lastRequestTimestamp")

            rows.append(row)
        if not rows:
            return
        table_path = f"{delta_root.rstrip('/')}/{table_name}"
        write_deltalake(
            table_path,
            _to_arrow_table(rows),
            mode="append",
            partition_by=["biztype", "applicationID", "date"],
            schema_mode="merge",
        )
        logger.info(
            "deltalake write %s rows=%d (%dms)",
            table_name,
            len(rows),
            int((time.time() - t_table) * 1000),
        )

    write_name_value_table(
        "series_count_by_metric_name",
        _ensure_list_of_dicts(data.get("seriesCountByMetricName")),
    )
    write_name_value_table(
        "series_count_by_label_name",
        _ensure_list_of_dicts(data.get("seriesCountByLabelName")),
    )
    write_name_value_table(
        "series_count_by_label_value_pair",
        _ensure_list_of_dicts(data.get("seriesCountByLabelValuePair")),
    )
    logger.info(
        "deltalake write total accountID=%s projectID=%s date=%s (%dms)",
        source.accountID,
        source.projectID,
        date,
        int((time.time() - t0) * 1000),
    )