"""Dataset loading and analysis."""

import hashlib
import json
from collections import Counter
from pathlib import Path

from bilgeai.analyzer.models import DatasetProfile
from bilgeai.core.exceptions import DataError


def analyze_dataset(path: Path, tokenizer_name: str | None = None) -> DatasetProfile:
    """Analyze a dataset file or directory and return its profile.

    Args:
        path: Path to dataset file or directory.
        tokenizer_name: HuggingFace model/tokenizer name for accurate token counting.
            Falls back to char/4 heuristic if not provided or loading fails.
    """
    if not path.exists():
        raise DataError(f"Dataset not found: {path}")

    # Directory support: find dataset files inside
    if path.is_dir():
        return _analyze_directory(path, tokenizer_name)

    suffix = path.suffix.lower()
    if suffix in (".jsonl", ".json"):
        return _analyze_jsonl(path, tokenizer_name)
    elif suffix == ".csv":
        return _analyze_csv(path, tokenizer_name)
    elif suffix == ".parquet":
        return _analyze_parquet(path, tokenizer_name)
    else:
        raise DataError(
            f"Unsupported file format: {suffix}",
            suggestion="Supported formats: .jsonl, .json, .csv, .parquet, or a directory",
        )


# ---------------------------------------------------------------------------
# Tokenizer loading
# ---------------------------------------------------------------------------


def _load_tokenizer(tokenizer_name: str | None):
    """Try to load a HuggingFace tokenizer. Returns (tokenizer, name) or (None, None)."""
    if not tokenizer_name:
        return None, None
    try:
        from transformers import AutoTokenizer

        tok = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
        return tok, tokenizer_name
    except Exception:
        return None, None


# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------


def _estimate_tokens(rows: list[dict], columns: list[str], tokenizer=None) -> dict:
    """Estimate token lengths from rows.

    Uses real tokenizer if provided, otherwise char/4 heuristic.
    Returns dict with avg_tokens, min_tokens, max_tokens, p50/p90/p99_tokens,
    text_column, and token_method.
    """
    text_cols = _find_text_columns(rows, columns)

    text_lengths = []
    use_tokenizer = tokenizer is not None

    for row in rows:
        if text_cols:
            text = " ".join(str(row.get(c, "")) for c in text_cols)
        else:
            text = " ".join(str(v) for v in row.values())

        if use_tokenizer:
            try:
                text_lengths.append(len(tokenizer.encode(text, add_special_tokens=False)))
            except Exception:
                text_lengths.append(len(text) // 4)
                use_tokenizer = False
        else:
            text_lengths.append(len(text) // 4)

    text_lengths.sort()
    n = len(text_lengths)

    if n == 0:
        return {}

    return {
        "avg_tokens": sum(text_lengths) / n,
        "min_tokens": text_lengths[0],
        "max_tokens": text_lengths[-1],
        "p50_tokens": text_lengths[n // 2],
        "p90_tokens": text_lengths[int(n * 0.9)],
        "p99_tokens": text_lengths[int(n * 0.99)],
        "text_column": ", ".join(text_cols) if text_cols else None,
        "token_method": "tokenizer" if use_tokenizer else "char_estimate",
    }


# ---------------------------------------------------------------------------
# Content analysis
# ---------------------------------------------------------------------------


def _analyze_content(rows: list[dict], columns: list[str]) -> dict:
    """Analyze dataset content: language, duplicates, empty rows.

    Returns dict with language, language_distribution, duplicate_ratio, empty_ratio.
    """
    text_cols = _find_text_columns(rows, columns)
    if not text_cols:
        return {}

    # Extract text samples
    texts = []
    empty_count = 0
    for row in rows:
        text = " ".join(str(row.get(c, "")) for c in text_cols).strip()
        if not text:
            empty_count += 1
        texts.append(text)

    n = len(texts)
    if n == 0:
        return {}

    # Empty ratio
    empty_ratio = empty_count / n

    # Duplicate detection (hash-based, fast)
    hashes = [hashlib.md5(t.encode("utf-8", errors="ignore")).hexdigest() for t in texts]
    unique_count = len(set(hashes))
    duplicate_ratio = 1.0 - (unique_count / n) if n > 0 else 0.0

    # Language detection (sample-based)
    lang_stats = _detect_language(texts[:500])

    return {
        "empty_ratio": round(empty_ratio, 4),
        "duplicate_ratio": round(duplicate_ratio, 4),
        **lang_stats,
    }


def _detect_language(texts: list[str]) -> dict:
    """Detect dominant language from text samples.

    Uses character-range heuristic (no external dependency).
    Returns dict with language and language_distribution.
    """
    lang_counts: Counter = Counter()

    for text in texts:
        if not text.strip():
            continue
        lang = _classify_text_language(text)
        lang_counts[lang] += 1

    total = sum(lang_counts.values())
    if total == 0:
        return {}

    distribution = {lang: round(count / total, 3) for lang, count in lang_counts.most_common(5)}
    dominant = lang_counts.most_common(1)[0][0]

    return {
        "language": dominant,
        "language_distribution": distribution,
    }


def _classify_text_language(text: str) -> str:
    """Classify a single text's language using Unicode character ranges.

    Simple but effective for distinguishing major language families.
    """
    sample = text[:500]

    # Count character types
    latin = 0
    turkish = 0
    cjk = 0
    cyrillic = 0
    arabic = 0
    total_alpha = 0

    for ch in sample:
        if not ch.isalpha():
            continue
        total_alpha += 1
        cp = ord(ch)

        if cp in (0xC7, 0xE7, 0xD6, 0xF6, 0xDC, 0xFC, 0x11E, 0x11F, 0x130, 0x131, 0x15E, 0x15F):
            turkish += 1
        elif 0x0400 <= cp <= 0x04FF:
            cyrillic += 1
        elif 0x0600 <= cp <= 0x06FF:
            arabic += 1
        elif 0x4E00 <= cp <= 0x9FFF or 0x3040 <= cp <= 0x30FF:
            cjk += 1
        elif cp < 0x0250:
            latin += 1

    if total_alpha == 0:
        return "unknown"

    # Turkish detection: Latin base + Turkish-specific characters
    if turkish / total_alpha > 0.02 and latin / total_alpha > 0.3:
        return "tr"
    if cyrillic / total_alpha > 0.3:
        return "ru"
    if arabic / total_alpha > 0.3:
        return "ar"
    if cjk / total_alpha > 0.3:
        return "zh"
    if latin / total_alpha > 0.5:
        return "en"

    return "other"


# ---------------------------------------------------------------------------
# Text column detection
# ---------------------------------------------------------------------------


def _find_text_columns(rows: list[dict], columns: list[str]) -> list[str]:
    """Identify columns that contain training-relevant text content."""
    # Known text column names in priority order
    known_text = [
        "text",
        "content",
        "instruction",
        "input",
        "output",
        "response",
        "question",
        "answer",
        "prompt",
        "completion",
        "conversations",
        "raw",
        "message",
        "body",
        "sentence",
    ]
    found = [c for c in columns if c.lower() in known_text]
    if found:
        return found

    # Fallback: find string columns with avg length > 50 chars (sample first 100 rows)
    sample = rows[:100]
    text_cols = []
    for col in columns:
        lengths = [len(str(r.get(col, ""))) for r in sample if r.get(col) is not None]
        if lengths and sum(lengths) / len(lengths) > 50:
            text_cols.append(col)
    return text_cols


# ---------------------------------------------------------------------------
# File-type analyzers
# ---------------------------------------------------------------------------


def _analyze_jsonl(path: Path, tokenizer_name: str | None = None) -> DatasetProfile:
    """Analyze a JSONL/JSON dataset."""
    rows = []
    with open(path, encoding="utf-8") as f:
        content = f.read().strip()

    # Handle both JSON array and JSONL formats
    if content.startswith("["):
        rows = json.loads(content)
    else:
        for line in content.split("\n"):
            line = line.strip()
            if line:
                rows.append(json.loads(line))

    if not rows:
        raise DataError("Dataset is empty", suggestion="Provide a non-empty dataset file.")

    columns = list(rows[0].keys()) if rows else []
    fmt = _detect_format(columns)

    tokenizer, tok_name = _load_tokenizer(tokenizer_name)
    token_stats = _estimate_tokens(rows, columns, tokenizer)
    content_stats = _analyze_content(rows, columns)

    return DatasetProfile(
        path=str(path),
        format=fmt,
        num_rows=len(rows),
        num_columns=len(columns),
        columns=columns,
        tokenizer_name=tok_name,
        estimated_size_mb=round(path.stat().st_size / (1024 * 1024), 2),
        **token_stats,
        **content_stats,
    )


def _analyze_csv(path: Path, tokenizer_name: str | None = None) -> DatasetProfile:
    """Analyze a CSV dataset."""
    import csv

    with open(path, encoding="utf-8") as f:
        reader = csv.DictReader(f)
        columns = list(reader.fieldnames or [])
        rows = list(reader)

    if not rows:
        raise DataError("Dataset is empty")

    fmt = _detect_format(columns)

    tokenizer, tok_name = _load_tokenizer(tokenizer_name)
    token_stats = _estimate_tokens(rows, columns, tokenizer)
    content_stats = _analyze_content(rows, columns)

    return DatasetProfile(
        path=str(path),
        format=fmt,
        num_rows=len(rows),
        num_columns=len(columns),
        columns=columns,
        tokenizer_name=tok_name,
        estimated_size_mb=round(path.stat().st_size / (1024 * 1024), 2),
        **token_stats,
        **content_stats,
    )


def _analyze_parquet(path: Path, tokenizer_name: str | None = None) -> DatasetProfile:
    """Analyze a Parquet dataset, sampling rows for token estimation."""
    try:
        import pyarrow.parquet as pq
    except ImportError as err:
        raise DataError(
            "pyarrow is required to read Parquet files",
            suggestion="pip install pyarrow",
        ) from err

    pf = pq.ParquetFile(path)
    metadata = pf.metadata
    total_rows = metadata.num_rows
    columns = pf.schema_arrow.names

    # Sample up to 10k rows for token estimation (read first row groups)
    max_sample = 10_000
    sampled_rows = 0
    batches = []
    for i in range(pf.num_row_groups):
        rg = pf.read_row_group(i)
        batches.append(rg)
        sampled_rows += rg.num_rows
        if sampled_rows >= max_sample:
            break

    import pyarrow as pa

    sample_table = pa.concat_tables(batches)
    rows = sample_table.to_pylist()

    tokenizer, tok_name = _load_tokenizer(tokenizer_name)
    token_stats = _estimate_tokens(rows, columns, tokenizer)
    content_stats = _analyze_content(rows, columns)

    return DatasetProfile(
        path=str(path),
        format=_detect_format(columns),
        num_rows=total_rows,
        num_columns=len(columns),
        columns=columns,
        tokenizer_name=tok_name,
        estimated_size_mb=round(path.stat().st_size / (1024 * 1024), 2),
        **token_stats,
        **content_stats,
    )


def _analyze_directory(path: Path, tokenizer_name: str | None = None) -> DatasetProfile:
    """Analyze a directory of dataset files by sampling the first few files."""
    supported = (".parquet", ".jsonl", ".json", ".csv")
    files = sorted(f for f in path.iterdir() if f.is_file() and f.suffix.lower() in supported)

    if not files:
        raise DataError(
            f"No dataset files found in {path}",
            suggestion=f"Directory should contain {', '.join(supported)} files.",
        )

    first_ext = files[0].suffix.lower()
    sample_files = files[: min(3, len(files))]

    total_rows = 0
    total_size_bytes = sum(f.stat().st_size for f in files)
    columns: list[str] = []

    # Collect per-file profiles and aggregate
    weighted_avg_sum = 0.0
    weighted_count = 0
    global_min = float("inf")
    global_max = 0
    p50_samples = []
    p90_samples = []
    last_profile: DatasetProfile | None = None

    for sample_file in sample_files:
        profile = analyze_dataset(sample_file, tokenizer_name)
        last_profile = profile
        if not columns:
            columns = profile.columns
        total_rows += profile.num_rows
        if profile.avg_tokens > 0:
            weighted_avg_sum += profile.avg_tokens * profile.num_rows
            weighted_count += profile.num_rows
            global_min = min(global_min, profile.min_tokens)
            global_max = max(global_max, profile.max_tokens)
            p50_samples.append(profile.p50_tokens)
            p90_samples.append(profile.p90_tokens)

    if sample_files:
        avg_rows_per_file = total_rows / len(sample_files)
        estimated_total_rows = int(avg_rows_per_file * len(files))
    else:
        estimated_total_rows = 0

    token_stats = {}
    if weighted_count > 0:
        avg_tokens = weighted_avg_sum / weighted_count
        p50_sorted = sorted(p50_samples)
        p90_sorted = sorted(p90_samples)
        token_stats = {
            "avg_tokens": avg_tokens,
            "min_tokens": int(global_min),
            "max_tokens": global_max,
            "p50_tokens": p50_sorted[len(p50_sorted) // 2],
            "p90_tokens": p90_sorted[len(p90_sorted) // 2],
            "p99_tokens": global_max,
        }

    # Carry over content analysis from last sampled file
    content_stats = {}
    if last_profile:
        content_stats = {
            "language": last_profile.language,
            "language_distribution": last_profile.language_distribution,
            "duplicate_ratio": last_profile.duplicate_ratio,
            "empty_ratio": last_profile.empty_ratio,
            "text_column": last_profile.text_column,
            "token_method": last_profile.token_method,
            "tokenizer_name": last_profile.tokenizer_name,
        }

    fmt = _detect_format(columns)
    file_label = f"{len(files)} {first_ext} files"

    return DatasetProfile(
        path=f"{path} ({file_label})",
        format=fmt,
        num_rows=estimated_total_rows,
        num_columns=len(columns),
        columns=columns,
        estimated_size_mb=round(total_size_bytes / (1024 * 1024), 2),
        **token_stats,
        **content_stats,
    )


# ---------------------------------------------------------------------------
# Format detection
# ---------------------------------------------------------------------------


def _detect_format(columns: list[str]) -> str:
    """Detect dataset format from column names."""
    col_set = {c.lower() for c in columns}

    # Alpaca format: instruction, input, output
    if {"instruction", "output"} <= col_set:
        return "alpaca"

    # ShareGPT format: conversations
    if "conversations" in col_set:
        return "sharegpt"

    # Custom mapping possible
    if col_set:
        return "custom"

    return "unknown"
