#!/usr/bin/env python3
"""Build and upload a flattened who-clicked dataset from local collected_data."""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any

import pandas as pd
from huggingface_hub import HfApi


PROJECT_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_HUMAN_DIR = PROJECT_ROOT / "experiments" / "collected_data"
DEFAULT_AGENT_DIR = PROJECT_ROOT / "experiments" / "agents" / "collected_data"
DEFAULT_OUTPUT = PROJECT_ROOT / "experiments" / "exports" / "whoclicked.csv"
DEFAULT_REPO = os.getenv("HF_WHOCLICKED_REPO", "velocitatem/whoclickedit")

BASE_COLUMNS = [
    "actor_type",
    "is_agent",
    "record_type",
    "topic",
    "source_session_dir",
    "source_file",
    "source_row_index",
    "ingest_format",
    "sessionId",
    "experimentId",
    "storeMode",
    "ts",
    "eventName",
    "page",
    "productId",
    "price",
    "userAgent",
    "kafka_partition_id",
    "kafka_offset",
    "kafka_timestamp_ms",
    "kafka_compression",
    "kafka_is_transactional",
    "kafka_headers",
    "kafka_key_payload",
    "kafka_key_encoding",
    "kafka_key_schema_id",
    "kafka_value_encoding",
    "kafka_value_schema_id",
    "kafka_value_size",
]


def _token() -> str | None:
    return os.getenv("HF_TOKEN") or None


def _exception_details(exc: Exception) -> str:
    parts = [str(exc).strip()]
    response = getattr(exc, "response", None)
    if response is not None:
        status = getattr(response, "status_code", None)
        if status is not None:
            parts.append(f"HTTP {status}")
        text = getattr(response, "text", "")
        if text:
            text = text.strip()
            if text:
                parts.append(text[:500])
    return " | ".join(p for p in parts if p)


def _flatten_dict(data: dict[str, Any], prefix: str = "") -> dict[str, Any]:
    flat: dict[str, Any] = {}
    for key, value in data.items():
        normalized_key = str(key).strip().replace(" ", "_")
        next_key = f"{prefix}_{normalized_key}" if prefix else normalized_key
        if isinstance(value, dict):
            flat.update(_flatten_dict(value, next_key))
        else:
            flat[next_key] = value
    return flat


def _as_scalar(value: Any) -> Any:
    if isinstance(value, (dict, list, tuple)):
        return json.dumps(value, ensure_ascii=True, sort_keys=True)
    return value


def _empty_envelope() -> dict[str, Any]:
    return {
        "kafka_partition_id": None,
        "kafka_offset": None,
        "kafka_timestamp_ms": None,
        "kafka_compression": None,
        "kafka_is_transactional": None,
        "kafka_headers": None,
        "kafka_key_payload": None,
        "kafka_key_encoding": None,
        "kafka_key_schema_id": None,
        "kafka_value_encoding": None,
        "kafka_value_schema_id": None,
        "kafka_value_size": None,
    }


def _extract_payload_and_envelope(
    record: Any,
) -> tuple[dict[str, Any], dict[str, Any], str]:
    if (
        isinstance(record, dict)
        and isinstance(record.get("value"), dict)
        and isinstance(record["value"].get("payload"), dict)
    ):
        key = record.get("key") if isinstance(record.get("key"), dict) else {}
        value = record["value"]
        envelope = {
            "kafka_partition_id": record.get("partitionID"),
            "kafka_offset": record.get("offset"),
            "kafka_timestamp_ms": record.get("timestamp"),
            "kafka_compression": record.get("compression"),
            "kafka_is_transactional": record.get("isTransactional"),
            "kafka_headers": _as_scalar(record.get("headers")),
            "kafka_key_payload": key.get("payload"),
            "kafka_key_encoding": key.get("encoding"),
            "kafka_key_schema_id": key.get("schemaId"),
            "kafka_value_encoding": value.get("encoding"),
            "kafka_value_schema_id": value.get("schemaId"),
            "kafka_value_size": value.get("size"),
        }
        return dict(value["payload"]), envelope, "kafka_envelope"

    if isinstance(record, dict):
        return dict(record), _empty_envelope(), "flat_payload"

    return {}, _empty_envelope(), "unknown"


def _load_json_list(path: Path) -> list[Any]:
    raw = json.loads(path.read_text())
    if not isinstance(raw, list):
        raise ValueError(f"Expected list in {path}, got {type(raw).__name__}")
    return raw


def _normalize_file_rows(
    actor_type: str,
    is_agent: int,
    session_dir_name: str,
    source_file: str,
    records: list[Any],
) -> list[dict[str, Any]]:
    record_type = "interaction" if source_file == "int.json" else "price_log"
    topic = "user-interactions" if record_type == "interaction" else "price-logs"

    rows: list[dict[str, Any]] = []
    for idx, raw_record in enumerate(records):
        payload, envelope, ingest_format = _extract_payload_and_envelope(raw_record)
        metadata = payload.pop("metadata", None)

        payload_flat = _flatten_dict(payload)
        row: dict[str, Any] = {
            "actor_type": actor_type,
            "is_agent": is_agent,
            "record_type": record_type,
            "topic": topic,
            "source_session_dir": session_dir_name,
            "source_file": source_file,
            "source_row_index": idx,
            "ingest_format": ingest_format,
            **envelope,
        }
        row.update({k: _as_scalar(v) for k, v in payload_flat.items()})

        if isinstance(metadata, dict):
            metadata_flat = _flatten_dict(metadata, "metadata")
            row.update({k: _as_scalar(v) for k, v in metadata_flat.items()})
        elif metadata is not None:
            row["metadata_raw"] = _as_scalar(metadata)

        rows.append(row)

    return rows


def _collect_rows_for_actor(
    actor_type: str, is_agent: int, base_dir: Path
) -> list[dict[str, Any]]:
    if not base_dir.exists():
        raise FileNotFoundError(f"Directory not found: {base_dir}")

    rows: list[dict[str, Any]] = []
    for session_dir in sorted(
        (p for p in base_dir.iterdir() if p.is_dir()), key=lambda p: p.name
    ):
        for source_file in ("int.json", "price.json"):
            file_path = session_dir / source_file
            if not file_path.exists():
                continue
            records = _load_json_list(file_path)
            rows.extend(
                _normalize_file_rows(
                    actor_type=actor_type,
                    is_agent=is_agent,
                    session_dir_name=session_dir.name,
                    source_file=source_file,
                    records=records,
                )
            )
    return rows


def build_dataframe(human_dir: Path, agent_dir: Path) -> pd.DataFrame:
    rows = [
        *_collect_rows_for_actor("human", 0, human_dir),
        *_collect_rows_for_actor("agent", 1, agent_dir),
    ]
    if not rows:
        return pd.DataFrame(columns=BASE_COLUMNS)

    df = pd.DataFrame(rows)
    ordered_columns = [
        *BASE_COLUMNS,
        *sorted(c for c in df.columns if c not in BASE_COLUMNS),
    ]
    return df[ordered_columns]


def _print_summary(df: pd.DataFrame, output_path: Path) -> None:
    print(f"wrote {len(df)} rows and {len(df.columns)} columns to {output_path}")
    if df.empty:
        return

    print("rows by actor/record_type:")
    grouped = (
        df.groupby(["actor_type", "record_type"], dropna=False)
        .size()
        .reset_index(name="count")
        .sort_values(["actor_type", "record_type"])
    )
    for _, row in grouped.iterrows():
        print(f"  - {row['actor_type']} / {row['record_type']}: {int(row['count'])}")

    required = ["actor_type", "is_agent", "record_type", "sessionId", "ts"]
    missing = {col: int(df[col].isna().sum()) for col in required if col in df.columns}
    print(f"missing in required columns: {missing}")


def build_csv(human_dir: Path, agent_dir: Path, output: Path) -> pd.DataFrame:
    df = build_dataframe(human_dir=human_dir, agent_dir=agent_dir)
    output.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(output, index=False)
    _print_summary(df, output)
    return df


def _resolve_repo_id(api: HfApi, repo_id: str) -> str:
    if "/" in repo_id:
        return repo_id
    try:
        me = api.whoami(token=_token())
        username = me.get("name")
        if username:
            return f"{username}/{repo_id}"
    except Exception:
        pass
    return repo_id


def upload_csv(
    input_path: Path,
    repo_id: str,
    path_in_repo: str,
    commit_message: str,
    create_if_missing: bool = False,
) -> None:
    if not input_path.exists():
        raise FileNotFoundError(f"Input CSV not found: {input_path}")

    api = HfApi(token=_token())

    try:
        me = api.whoami(token=_token())
    except Exception as exc:
        detail = _exception_details(exc)
        hint = "Set HF_TOKEN with write access or run huggingface-cli login."
        raise RuntimeError(
            f"Hugging Face auth failed. {hint} Details: {detail}"
        ) from exc

    user_name = me.get("name") or me.get("fullname") or "unknown"
    print(f"authenticated to HF as: {user_name}")

    resolved_repo_id = _resolve_repo_id(api, repo_id)
    if create_if_missing:
        api.create_repo(repo_id=resolved_repo_id, repo_type="dataset", exist_ok=True)
    else:
        try:
            api.repo_info(repo_id=resolved_repo_id, repo_type="dataset")
        except Exception as exc:
            detail = _exception_details(exc)
            hint = (
                "Check owner/repo spelling, ensure it is a dataset repo, "
                "or pass --create-if-missing."
            )
            raise RuntimeError(
                f"Dataset repo '{resolved_repo_id}' is not accessible. {hint} Details: {detail}"
            ) from exc

    try:
        commit = api.upload_file(
            path_or_fileobj=str(input_path),
            path_in_repo=path_in_repo,
            repo_id=resolved_repo_id,
            repo_type="dataset",
            commit_message=commit_message,
        )
    except Exception as exc:
        detail = _exception_details(exc)
        hint = (
            "Pass --repo <owner>/whoclickedit and ensure HF_TOKEN is set "
            "(or run huggingface-cli login)."
        )
        raise RuntimeError(
            f"Upload failed for '{resolved_repo_id}'. {hint} Details: {detail}"
        ) from exc

    print(
        f"uploaded {input_path} to https://huggingface.co/datasets/{resolved_repo_id}"
    )
    print(f"commit: {commit}")


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="ETL for whoclickedit: flatten local collected_data and upload to HF"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    build = sub.add_parser("build", help="build flattened CSV locally")
    build.add_argument("--human-dir", type=Path, default=DEFAULT_HUMAN_DIR)
    build.add_argument("--agent-dir", type=Path, default=DEFAULT_AGENT_DIR)
    build.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)

    upload = sub.add_parser("upload", help="upload an existing CSV to HF dataset")
    upload.add_argument("--input", type=Path, default=DEFAULT_OUTPUT)
    upload.add_argument("--repo", default=DEFAULT_REPO)
    upload.add_argument("--path-in-repo", default="whoclicked.csv")
    upload.add_argument("--message", default="Update flattened whoclickedit dataset")
    upload.add_argument("--create-if-missing", action="store_true")

    build_upload = sub.add_parser(
        "build-upload", help="build CSV and upload to HF dataset"
    )
    build_upload.add_argument("--human-dir", type=Path, default=DEFAULT_HUMAN_DIR)
    build_upload.add_argument("--agent-dir", type=Path, default=DEFAULT_AGENT_DIR)
    build_upload.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    build_upload.add_argument("--repo", default=DEFAULT_REPO)
    build_upload.add_argument("--path-in-repo", default="whoclicked.csv")
    build_upload.add_argument(
        "--message", default="Update flattened whoclickedit dataset"
    )
    build_upload.add_argument("--create-if-missing", action="store_true")

    return parser.parse_args()


def main() -> int:
    args = _parse_args()

    try:
        if args.command == "build":
            build_csv(
                human_dir=args.human_dir, agent_dir=args.agent_dir, output=args.output
            )
            return 0

        if args.command == "upload":
            upload_csv(
                input_path=args.input,
                repo_id=args.repo,
                path_in_repo=args.path_in_repo,
                commit_message=args.message,
                create_if_missing=args.create_if_missing,
            )
            return 0

        if args.command == "build-upload":
            build_csv(
                human_dir=args.human_dir, agent_dir=args.agent_dir, output=args.output
            )
            upload_csv(
                input_path=args.output,
                repo_id=args.repo,
                path_in_repo=args.path_in_repo,
                commit_message=args.message,
                create_if_missing=args.create_if_missing,
            )
            return 0

        raise ValueError(f"Unknown command: {args.command}")

    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())