from __future__ import annotations

import hashlib
import json
import os
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path

import pandas as pd
from tqdm.auto import tqdm

from .normalize import prepare_dataframe
from .taxonomy import SOURCE_LABEL_MAP

NORMALIZATION_VERSION = "2.2"


@dataclass(frozen=True)
class FetchedDataset:
    source: str
    dataset_id: str
    split: str
    rows: int
    cache_path: str
    cached: bool


def _extract_label(row: dict, label_column: str):
    if "." not in label_column:
        return row.get(label_column)
    current = row.get(label_column.split(".")[0])
    for part in label_column.split(".")[1:]:
        current = current.get(part) if isinstance(current, dict) else None
    return current


def _extract_finee_chatml(row: dict) -> tuple[str, str | None]:
    messages = row.get("messages")
    if not isinstance(messages, list):
        raise TypeError("FinEE ChatML row must contain a list-valued 'messages' field.")

    user_content = None
    assistant_content = None
    for message in messages:
        if not isinstance(message, dict):
            continue
        role = message.get("role")
        if role == "user" and user_content is None:
            user_content = message.get("content")
        elif role == "assistant" and assistant_content is None:
            assistant_content = message.get("content")

    if not isinstance(user_content, str) or not user_content.strip():
        raise ValueError("FinEE ChatML row has no non-empty user message content.")
    if not isinstance(assistant_content, str) or not assistant_content.strip():
        raise ValueError("FinEE ChatML row has no non-empty assistant message content.")

    transaction_text = re.sub(
        r"^\s*extract financial entities from\s*:\s*",
        "",
        user_content.strip(),
        flags=re.IGNORECASE,
    )

    payload_text = assistant_content.strip()
    if payload_text.startswith("```"):
        lines = payload_text.splitlines()
        if lines and lines[0].strip().startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        payload_text = "\n".join(lines).strip()
        if payload_text.lower().startswith("json"):
            payload_text = payload_text[4:].lstrip()

    try:
        payload = json.loads(payload_text)
    except json.JSONDecodeError as exc:
        raise ValueError("FinEE assistant content is not valid JSON.") from exc

    if not isinstance(payload, dict):
        raise TypeError("FinEE assistant JSON must be an object.")
    category = payload.get("category")
    if category is None or not str(category).strip():
        return transaction_text, None

    return transaction_text, str(category)


def _extract_training_fields(
    row: dict,
    *,
    dataset_id: str,
    text_column: str,
    label_column: str,
    dataset_format: str | None,
) -> tuple[object, object]:
    if dataset_format == "finee-chatml":
        return _extract_finee_chatml(row)

    if text_column not in row:
        raise KeyError(f"Dataset '{dataset_id}' does not contain text column '{text_column}'.")
    return row[text_column], _extract_label(row, label_column)


def _batch_to_rows(batch: object) -> list[dict]:
    """Normalize Hugging Face columnar batches and lightweight row-mapping test doubles."""
    if isinstance(batch, Mapping):
        keys = list(batch)
        if not keys:
            return []
        columns = {key: batch[key] for key in keys}
        try:
            lengths = {len(values) for values in columns.values()}
        except TypeError as exc:
            raise TypeError("Dataset batch columns must be sized sequences.") from exc
        if len(lengths) != 1:
            raise ValueError("Dataset batch columns have inconsistent lengths.")
        row_count = next(iter(lengths))
        return [{key: columns[key][index] for key in keys} for index in range(row_count)]

    if isinstance(batch, Sequence) and not isinstance(batch, (str, bytes, bytearray)):
        rows: list[dict] = []
        for index, value in enumerate(batch):
            if not isinstance(value, Mapping):
                raise TypeError(f"Dataset row {index} must be a mapping, got {type(value).__name__}.")
            rows.append(dict(value))
        return rows

    raise TypeError(
        "Dataset slice must be either a mapping of column sequences or a sequence of row mappings."
    )


def _cache_signature(
    *,
    dataset_id: str,
    split: str,
    source: str,
    text_column: str,
    label_column: str,
    dataset_format: str | None,
    country_column: str | None,
    currency_column: str | None,
    language_column: str | None,
    default_country: str,
    default_currency: str,
    default_language: str,
) -> dict:
    mapping = SOURCE_LABEL_MAP.get(source)
    if mapping is None:
        raise ValueError(f"No canonical taxonomy mapping is registered for source '{source}'.")
    payload = {
        "normalization_version": NORMALIZATION_VERSION,
        "dataset_id": dataset_id,
        "split": split,
        "source": source,
        "text_column": text_column,
        "label_column": label_column,
        "dataset_format": dataset_format,
        "country_column": country_column,
        "currency_column": currency_column,
        "language_column": language_column,
        "default_country": default_country,
        "default_currency": default_currency,
        "default_language": default_language,
        "taxonomy_mapping": dict(sorted(mapping.items())),
    }
    fingerprint = hashlib.sha256(
        json.dumps(payload, sort_keys=True, ensure_ascii=False).encode("utf-8")
    ).hexdigest()
    return {"fingerprint": fingerprint, **payload}


def _cache_is_current(metadata_path: Path, expected_signature: dict) -> bool:
    try:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return False
    return metadata.get("cache_signature") == expected_signature


def fetch_huggingface_dataset(
    dataset_id: str,
    text_column: str,
    label_column: str,
    source: str,
    split: str = "train",
    cache_dir: str | Path | None = None,
    progress: bool = True,
    force: bool = False,
    dataset_format: str | None = None,
    country_column: str | None = None,
    currency_column: str | None = None,
    language_column: str | None = None,
    default_country: str = "unknown",
    default_currency: str = "unknown",
    default_language: str = "unknown",
    normalize_chunk_size: int = 250_000,
) -> tuple[pd.DataFrame, FetchedDataset]:
    """Fetch, validate, normalize, and cache one dataset with configuration provenance."""
    if default_country == "unknown" and source in {"finee-india", "synthetic-indian-transactions"}:
        default_country = "India"
    if default_currency == "unknown" and source in {"finee-india", "synthetic-indian-transactions"}:
        default_currency = "INR"

    from datasets import load_dataset
    from datasets.exceptions import DatasetNotFoundError

    target_dir = Path(cache_dir or "data/cache") / "normalized"
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / f"{source}.parquet"
    metadata_path = target.with_suffix(".json")
    signature = _cache_signature(
        dataset_id=dataset_id,
        split=split,
        source=source,
        text_column=text_column,
        label_column=label_column,
        dataset_format=dataset_format,
        country_column=country_column,
        currency_column=currency_column,
        language_column=language_column,
        default_country=default_country,
        default_currency=default_currency,
        default_language=default_language,
    )

    if target.exists() and metadata_path.exists() and not force and _cache_is_current(metadata_path, signature):
        frame = pd.read_parquet(target)
        return frame, FetchedDataset(source, dataset_id, split, len(frame), str(target), True)

    hf_token = os.getenv("HF_TOKEN") or None
    try:
        dataset = load_dataset(
            dataset_id,
            split=split,
            cache_dir=str(cache_dir) if cache_dir else None,
            token=hf_token,
        )
    except DatasetNotFoundError as exc:
        message = str(exc).lower()
        if "gated dataset" in message:
            if not hf_token:
                raise RuntimeError(
                    f"Hugging Face dataset '{dataset_id}' is gated and requires authentication. "
                    "In Kaggle, attach a Secret named 'HF_TOKEN' containing a Hugging Face "
                    "access token with permission to this dataset."
                ) from exc
            raise RuntimeError(
                f"Hugging Face authentication was provided, but access to gated dataset "
                f"'{dataset_id}' was not granted for that token."
            ) from exc
        raise

    shard_dir = target_dir / f"{source}.parts"
    if shard_dir.exists():
        for part in shard_dir.glob("*.parquet"):
            part.unlink()
    shard_dir.mkdir(parents=True, exist_ok=True)

    total = len(dataset)
    outer = tqdm(
        range(0, total, normalize_chunk_size),
        total=(total + normalize_chunk_size - 1) // normalize_chunk_size,
        desc=f"Fetching {source}",
        unit="chunks",
        disable=not progress,
    )
    shard_paths: list[Path] = []
    for shard_number, start in enumerate(outer):
        stop = min(start + normalize_chunk_size, total)
        rows = _batch_to_rows(dataset[start:stop])
        if len(rows) != stop - start:
            raise ValueError(
                f"Dataset '{dataset_id}' returned {len(rows)} rows for requested slice "
                f"[{start}:{stop}], expected {stop - start}."
            )

        batch_rows = []
        for offset, row in enumerate(rows):
            text, label = _extract_training_fields(
                row,
                dataset_id=dataset_id,
                text_column=text_column,
                label_column=label_column,
                dataset_format=dataset_format,
            )
            country = row.get(country_column) if country_column else default_country
            currency = row.get(currency_column) if currency_column else default_currency
            language = row.get(language_column) if language_column else default_language
            raw_id = f"{source}:{start + offset}:{text}:{label}"
            record_id = hashlib.sha1(raw_id.encode("utf-8")).hexdigest()
            batch_rows.append(
                {
                    "text": text,
                    "label": label,
                    "source": source,
                    "source_label": label,
                    "country": country,
                    "currency": currency,
                    "language": language,
                    "record_id": record_id,
                }
            )
        batch_frame = prepare_dataframe(
            pd.DataFrame(batch_rows),
            progress=progress,
            chunk_size=normalize_chunk_size,
            canonicalize=True,
            strict_taxonomy=True,
        )
        part_path = shard_dir / f"part-{shard_number:05d}.parquet"
        batch_frame.to_parquet(part_path, index=False)
        shard_paths.append(part_path)
        outer.set_postfix(rows=stop)

    if shard_paths:
        frame = pd.concat((pd.read_parquet(path) for path in shard_paths), ignore_index=True)
    else:
        frame = pd.DataFrame(
            columns=[
                "text", "label", "source", "source_label", "country", "currency", "language", "record_id"
            ]
        )
    frame = prepare_dataframe(frame, progress=progress, chunk_size=normalize_chunk_size)
    frame.to_parquet(target, index=False)
    metadata_path.write_text(
        json.dumps(
            {
                "cache_signature": signature,
                "source": source,
                "dataset_id": dataset_id,
                "split": split,
                "rows": len(frame),
                "dataset_format": dataset_format,
                "normalization_version": NORMALIZATION_VERSION,
                "countries": sorted(frame["country"].dropna().unique().tolist()) if "country" in frame else [],
                "categories": sorted(frame["label"].dropna().unique().tolist()),
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    for part_path in shard_paths:
        part_path.unlink(missing_ok=True)
    try:
        shard_dir.rmdir()
    except OSError:
        pass
    return frame, FetchedDataset(source, dataset_id, split, len(frame), str(target), False)


def fetch_configured_datasets(
    config: list[dict],
    cache_dir: str | Path,
    progress: bool = True,
    force: bool = False,
) -> tuple[pd.DataFrame, list[dict]]:
    """Fetch every Hugging Face dataset declared in the training manifest."""
    frames: list[pd.DataFrame] = []
    manifests: list[dict] = []

    sources = [item for item in config if item.get("dataset_id")]
    for item in tqdm(sources, desc="Dataset sources", unit="dataset", disable=not progress):
        frame, fetched = fetch_huggingface_dataset(
            dataset_id=item["dataset_id"],
            text_column=item.get("text_column", ""),
            label_column=item.get("label_column", ""),
            source=item["source"],
            split=item.get("split", "train"),
            cache_dir=cache_dir,
            progress=progress,
            force=force,
            dataset_format=item.get("format"),
            country_column=item.get("country_column"),
            currency_column=item.get("currency_column"),
            language_column=item.get("language_column"),
            default_country=item.get("default_country", "unknown"),
            default_currency=item.get("default_currency", "unknown"),
            default_language=item.get("default_language", "unknown"),
            normalize_chunk_size=item.get("normalize_chunk_size", 250_000),
        )
        frames.append(frame)
        manifests.append(
            {
                "source": fetched.source,
                "dataset_id": fetched.dataset_id,
                "split": fetched.split,
                "rows": fetched.rows,
                "cache_path": fetched.cache_path,
                "cached": fetched.cached,
            }
        )

    if not frames:
        raise ValueError("No Hugging Face datasets with dataset_id were configured.")
    combined = prepare_dataframe(pd.concat(frames, ignore_index=True), progress=progress, chunk_size=250_000)
    return combined, manifests