"""RAM estimation for record linkage jobs."""

from __future__ import annotations

import os
import platform
import subprocess
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List, Optional

if TYPE_CHECKING:
    from blazematch.config import (
        BlockingRule,
        EmbeddingBlockingConfig,
        FieldComparison,
    )

from blazematch.config import NUMERIC_METRICS

# Bytes per element for numeric columns in a pandas DataFrame (float64)
_NUMERIC_COLUMN_BYTES = 8


@dataclass
class RAMEstimate:
    """Result of a RAM estimation."""

    blocking_bytes: int
    feature_bytes: int
    model_train_bytes: int
    model_inference_bytes: int
    dataframe_bytes: int
    peak_bytes: int
    warnings: List[str] = field(default_factory=list)
    breakdown: dict = field(default_factory=dict)
    system_ram_gb: float = 0.0

    @property
    def minimum_ram_gb(self) -> float:
        """Peak memory * 1.2x overhead, in GB."""
        return self.peak_bytes * 1.2 / (1024**3)

    @property
    def recommended_ram_gb(self) -> float:
        """Minimum * 2x for OS/GC headroom, in GB."""
        return self.minimum_ram_gb * 2.0

    @property
    def can_proceed(self) -> bool:
        """Whether system RAM >= minimum required."""
        return self.system_ram_gb >= self.minimum_ram_gb

    def summary(self) -> str:
        """Human-readable summary of the estimation."""
        lines = [
            "=== Blazematch RAM Estimate ===",
            "",
            f"Estimated pairs:       {self.breakdown.get('estimated_pairs', '?'):,}",
            f"Feature columns:       {self.breakdown.get('feature_count', '?')}",
            f"Model type:            {self.breakdown.get('model', '?')}",
            "",
            "Memory by stage:",
            f"  DataFrames:          {_fmt_bytes(self.dataframe_bytes)}",
            f"  Blocking:            {_fmt_bytes(self.blocking_bytes)}",
            f"  Feature computation: {_fmt_bytes(self.feature_bytes)}",
            f"  Model training:      {_fmt_bytes(self.model_train_bytes)}",
            f"  Model inference:     {_fmt_bytes(self.model_inference_bytes)}",
            "",
            f"Peak memory:           {_fmt_bytes(self.peak_bytes)}",
            f"Minimum RAM:           {self.minimum_ram_gb:.2f} GB (peak * 1.2x)",
            f"Recommended RAM:       {self.recommended_ram_gb:.2f} GB (minimum * 2x)",
            f"System RAM:            {self.system_ram_gb:.2f} GB",
            "",
        ]

        if self.can_proceed:
            lines.append("Status: OK — system has enough RAM")
        else:
            lines.append(
                f"Status: WARNING — system RAM ({self.system_ram_gb:.2f} GB) "
                f"< minimum ({self.minimum_ram_gb:.2f} GB)"
            )

        if self.warnings:
            lines.append("")
            lines.append("Warnings:")
            for w in self.warnings:
                lines.append(f"  - {w}")

        return "\n".join(lines)


def _fmt_bytes(b: int) -> str:
    """Format bytes as human-readable string."""
    if b < 1024:
        return f"{b} B"
    elif b < 1024**2:
        return f"{b / 1024:.1f} KB"
    elif b < 1024**3:
        return f"{b / 1024**2:.1f} MB"
    else:
        return f"{b / 1024**3:.2f} GB"


def _get_system_ram_bytes() -> int:
    """Detect physical RAM in bytes (cross-platform, no psutil).

    Returns 0 if detection fails on all methods.
    """
    system = platform.system()

    if system == "Linux":
        try:
            with open("/proc/meminfo") as f:
                for line in f:
                    if line.startswith("MemTotal:"):
                        # MemTotal: value in kibibytes (1024 bytes)
                        return int(line.split()[1]) * 1024
        except (OSError, ValueError, IndexError):
            pass

    if system == "Darwin":
        try:
            out = subprocess.check_output(
                ["sysctl", "-n", "hw.memsize"], text=True
            )
            return int(out.strip())
        except (subprocess.SubprocessError, ValueError):
            pass

    if system == "Windows":
        try:
            import ctypes

            class MEMORYSTATUSEX(ctypes.Structure):
                _fields_ = [
                    ("dwLength", ctypes.c_ulong),
                    ("dwMemoryLoad", ctypes.c_ulong),
                    ("ullTotalPhys", ctypes.c_ulonglong),
                    ("ullAvailPhys", ctypes.c_ulonglong),
                    ("ullTotalPageFile", ctypes.c_ulonglong),
                    ("ullAvailPageFile", ctypes.c_ulonglong),
                    ("ullTotalVirtual", ctypes.c_ulonglong),
                    ("ullAvailVirtual", ctypes.c_ulonglong),
                    ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
                ]

            stat = MEMORYSTATUSEX()
            stat.dwLength = ctypes.sizeof(stat)
            ok = ctypes.windll.kernel32.GlobalMemoryStatusEx(
                ctypes.byref(stat)
            )
            if ok and stat.ullTotalPhys > 0:
                return stat.ullTotalPhys
        except (AttributeError, OSError):
            pass

    # Fallback via os.sysconf
    try:
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
        return pages * page_size
    except (ValueError, OSError, AttributeError):
        pass

    return 0


def _estimate_n_pairs(
    n_left: int,
    n_right: int,
    n_pairs: Optional[int] = None,
    blocking_reduction_factor: Optional[float] = None,
    blocking_rules: Optional[List[BlockingRule]] = None,
    embedding_blocking: Optional[EmbeddingBlockingConfig] = None,
) -> tuple[int, list[str]]:
    """Estimate the number of candidate pairs after blocking.

    Returns (estimated_pairs, warnings).
    """
    warnings: list[str] = []

    if n_pairs is not None:
        return n_pairs, warnings

    if blocking_reduction_factor is not None:
        if blocking_reduction_factor <= 0:
            raise ValueError(
                f"blocking_reduction_factor must be positive, "
                f"got {blocking_reduction_factor}"
            )
        est = int(n_left * n_right / blocking_reduction_factor)
        return est, warnings

    if embedding_blocking is not None:
        est = n_left * embedding_blocking.top_k
        return est, warnings

    if blocking_rules:
        n_rules = len(blocking_rules)
        est = int(n_left * n_right / (100 * n_rules))
        warnings.append(
            f"Pair count is a heuristic estimate ({est:,} pairs). "
            "Pass n_pairs or blocking_reduction_factor for accuracy."
        )
        return est, warnings

    # No blocking info at all
    est = int(n_left * n_right / 100)
    warnings.append(
        f"Pair count is a rough heuristic ({est:,} pairs). "
        "Pass n_pairs or blocking_reduction_factor for accuracy."
    )
    return est, warnings


def estimate_ram(
    n_left: int,
    n_right: int,
    comparisons: List[FieldComparison],
    blocking_rules: Optional[List[BlockingRule]] = None,
    embedding_blocking: Optional[EmbeddingBlockingConfig] = None,
    model: str = "logistic",
    n_pairs: Optional[int] = None,
    blocking_reduction_factor: Optional[float] = None,
    avg_string_length: int = 20,
    system_ram_gb: Optional[float] = None,
    model_kwargs: Optional[dict] = None,
    print_summary: bool = True,
) -> RAMEstimate:
    """Estimate RAM requirements for a record linkage job.

    Args:
        n_left: Number of records in the left dataset.
        n_right: Number of records in the right dataset.
        comparisons: List of FieldComparison objects (same as LinkConfig).
        blocking_rules: Optional blocking rules.
        embedding_blocking: Optional embedding blocking config.
        model: Model type ("logistic", "random_forest", "gradient_boosting",
               "svm", "xgboost", "lightgbm", "fellegi_sunter").
        n_pairs: Exact number of candidate pairs (overrides heuristic).
        blocking_reduction_factor: E.g. 50 for state-level blocking.
        avg_string_length: Average string field length in bytes (default 20).
        system_ram_gb: Override detected system RAM (in GB).
        model_kwargs: Optional model parameters (n_estimators, max_depth, etc.).
        print_summary: If True, print human-readable summary.

    Returns:
        RAMEstimate with per-stage memory estimates.
    """
    if not comparisons:
        raise ValueError("comparisons must not be empty")

    model_kwargs = model_kwargs or {}
    warnings: list[str] = []

    # Estimate pair count
    est_pairs, pair_warnings = _estimate_n_pairs(
        n_left, n_right,
        n_pairs=n_pairs,
        blocking_reduction_factor=blocking_reduction_factor,
        blocking_rules=blocking_rules,
        embedding_blocking=embedding_blocking,
    )
    warnings.extend(pair_warnings)

    if est_pairs > 100_000_000:
        warnings.append(
            f"Estimated {est_pairs:,} pairs — this is very large and may be slow."
        )

    # Count features
    n_features = len(comparisons)
    n_string = sum(1 for c in comparisons if c.metric not in NUMERIC_METRICS)
    n_numeric = n_features - n_string

    # --- DataFrame memory ---
    # String columns: avg_string_length per cell; numeric columns: 8 bytes (float64)
    n_rows = n_left + n_right
    dataframe_bytes = int(
        n_rows * n_string * avg_string_length
        + n_rows * n_numeric * _NUMERIC_COLUMN_BYTES
    )

    # --- Blocking memory ---
    if embedding_blocking is not None:
        dim = embedding_blocking.max_features
        # TF-IDF matrices are sparse; typical density ~1% for char n-grams
        sparsity = 0.01
        tfidf_bytes = int(
            # CSR format: ~12 bytes per nonzero (8 float + 4 index)
            12 * (n_left + n_right) * dim * sparsity
        )
        blocking_bytes = int(
            tfidf_bytes
            + 4 * n_right * dim * sparsity  # FAISS flat index (from sparse vecs)
            + 4 * n_right * dim             # FAISS reconstructed dense index
            + 2.5 * 1024 * 1024            # FAISS overhead (~2.5 MB)
            + 16 * est_pairs                # pair indices
        )
    else:
        # Rule-based: DuckDB copies + pair output
        blocking_bytes = int(
            n_rows * n_features * avg_string_length * 2
            + 16 * est_pairs
        )

    # --- Feature computation memory ---
    # String inputs: pairs * n_string * avg_string_length * 2 (left + right)
    string_input = int(est_pairs * n_string * avg_string_length * 2)
    # Numeric inputs: pairs * n_numeric * 8 * 2
    numeric_input = int(est_pairs * n_numeric * 8 * 2)
    # Output matrix + intermediate: 3 * 8 * pairs * features
    feature_output = int(3 * 8 * est_pairs * n_features)
    feature_bytes = string_input + numeric_input + feature_output

    # --- Model training memory ---
    model_train_bytes = 0
    if model in ("random_forest", "gradient_boosting", "lightgbm", "xgboost"):
        n_trees = model_kwargs.get("n_estimators", 100)
        max_depth = model_kwargs.get("max_depth", 6)
        # Tree structure: ~80 bytes per node
        tree_structure = int(80 * n_trees * (2 ** (max_depth + 1) - 1))
        # sklearn training: sorts feature arrays per tree split
        # Dominant cost is O(n_samples * n_features * 8) for sort buffers
        train_size = int(est_pairs * 0.3)  # typical labelled subset
        sort_buffers = int(train_size * n_features * 8)
        model_train_bytes = tree_structure + sort_buffers
    elif model == "svm":
        # SVM needs kernel matrix for training data
        train_size = int(est_pairs * 0.3)  # typical labelled subset
        model_train_bytes = int(8 * train_size * n_features)
        if est_pairs > 500_000:
            warnings.append(
                "SVM with >500K pairs can be very slow and memory-intensive. "
                "Consider logistic regression or tree-based models."
            )
    elif model == "dbscan":
        # DBSCAN computes pairwise distance matrix: O(n^2 * 8) bytes
        model_train_bytes = int(8 * est_pairs * est_pairs)
        if est_pairs > 1_000_000:
            warnings.append(
                "DBSCAN with >1M pairs requires O(n^2) distance matrix. "
                "Consider other clustering methods."
            )

    # --- Model inference memory ---
    model_inference_bytes = int(8 * est_pairs)  # score array
    if model == "svm":
        # SVM inference scales with number of support vectors;
        # estimate ~10% of training pairs as support vectors
        n_sv_est = int(est_pairs * 0.1)
        model_inference_bytes += int(8 * est_pairs * min(n_sv_est, 10000))

    # --- Peak memory ---
    # Peak is max of concurrent stages
    pairs_array = 16 * est_pairs  # left_idx + right_idx
    feature_matrix = 8 * est_pairs * n_features

    blocking_peak = dataframe_bytes + blocking_bytes
    feature_peak = dataframe_bytes + pairs_array + feature_bytes
    inference_peak = (
        dataframe_bytes + pairs_array + feature_matrix + model_inference_bytes
    )
    training_peak = (
        dataframe_bytes + pairs_array + feature_matrix + model_train_bytes
    )

    peak_bytes = max(blocking_peak, feature_peak, inference_peak, training_peak)

    # System RAM
    if system_ram_gb is not None:
        sys_ram_gb = system_ram_gb
    else:
        sys_ram_bytes = _get_system_ram_bytes()
        if sys_ram_bytes > 0:
            sys_ram_gb = sys_ram_bytes / (1024**3)
        else:
            sys_ram_gb = 0.0
            warnings.append(
                "Could not detect system RAM. "
                "Pass system_ram_gb to enable the can_proceed check."
            )

    result = RAMEstimate(
        blocking_bytes=blocking_bytes,
        feature_bytes=feature_bytes,
        model_train_bytes=model_train_bytes,
        model_inference_bytes=model_inference_bytes,
        dataframe_bytes=dataframe_bytes,
        peak_bytes=peak_bytes,
        warnings=warnings,
        system_ram_gb=sys_ram_gb,
        breakdown={
            "estimated_pairs": est_pairs,
            "feature_count": n_features,
            "model": model,
            "n_string_comparisons": n_string,
            "n_numeric_comparisons": n_numeric,
        },
    )

    if print_summary:
        print(result.summary())

    return result
