"""DataBridge AI — Turbo Engine.

Polars + DuckDB acceleration layer with Pandas fallback.
All functions return Pandas-compatible types for tool output compatibility.
"""
import hashlib
import json
import os
from pathlib import Path

import pandas as pd

# ---------------------------------------------------------------------------
# Availability flags (same pattern as RAPIDFUZZ_AVAILABLE in server.py)
# ---------------------------------------------------------------------------

try:
    import polars as pl
    POLARS_AVAILABLE = True
except ImportError:
    POLARS_AVAILABLE = False

try:
    import duckdb
    DUCKDB_AVAILABLE = True
except ImportError:
    DUCKDB_AVAILABLE = False


def engine_status() -> dict:
    """Return availability of acceleration engines."""
    return {
        "polars": POLARS_AVAILABLE,
        "duckdb": DUCKDB_AVAILABLE,
        "polars_version": pl.__version__ if POLARS_AVAILABLE else None,
        "duckdb_version": duckdb.__version__ if DUCKDB_AVAILABLE else None,
    }


# ---------------------------------------------------------------------------
# Shared DuckDB connection (session-scoped singleton)
# ---------------------------------------------------------------------------

_duckdb_conn = None


def get_duckdb():
    """Get or create the shared in-memory DuckDB connection."""
    global _duckdb_conn
    if not DUCKDB_AVAILABLE:
        raise RuntimeError("DuckDB is not installed. Run: pip install duckdb")
    if _duckdb_conn is None:
        _duckdb_conn = duckdb.connect(":memory:")
    return _duckdb_conn


def reset_duckdb():
    """Close and reset the shared DuckDB connection (for testing)."""
    global _duckdb_conn
    if _duckdb_conn is not None:
        _duckdb_conn.close()
        _duckdb_conn = None


# ---------------------------------------------------------------------------
# CSV / Data loading
# ---------------------------------------------------------------------------

def read_csv(path: str) -> pd.DataFrame:
    """Read CSV with Polars (fast) then convert to Pandas for tool compatibility.

    Falls back to Pandas if Polars is not installed.
    """
    if POLARS_AVAILABLE:
        try:
            return pl.read_csv(path, infer_schema_length=10000).to_pandas()
        except Exception:
            # Polars can fail on edge-case encodings; fall back silently
            pass
    return pd.read_csv(path)


def read_csv_polars(path: str):
    """Read CSV as a native Polars DataFrame. Returns None if Polars unavailable."""
    if POLARS_AVAILABLE:
        return pl.read_csv(path, infer_schema_length=10000)
    return None


def dataframe_from_records(records: list) -> pd.DataFrame:
    """Build a DataFrame from a list of dicts, using Polars if available."""
    if POLARS_AVAILABLE and records:
        try:
            return pl.DataFrame(records).to_pandas()
        except Exception:
            pass
    return pd.DataFrame(records)


# ---------------------------------------------------------------------------
# Profiling
# ---------------------------------------------------------------------------

def profile_stats(path: str) -> dict:
    """Fast profiling using Polars. Returns dict matching profile_data output format.

    Falls back to Pandas if Polars is unavailable.
    """
    if POLARS_AVAILABLE:
        try:
            return _profile_polars(path)
        except Exception:
            pass
    return _profile_pandas(path)


def _profile_polars(path: str) -> dict:
    """Polars-native profiling — vectorized null_count, n_unique, describe."""
    df = pl.read_csv(path, infer_schema_length=10000)
    n_rows = df.height
    n_cols = df.width

    # Column types
    col_types = {col: str(dtype) for col, dtype in zip(df.columns, df.dtypes)}

    # Cardinality (n_unique / n_rows) — computed in parallel by Polars
    cardinality = {}
    for col in df.columns:
        n_unique = df[col].n_unique()
        cardinality[col] = n_unique / n_rows if n_rows > 0 else 0.0

    # Key / cardinality classification
    potential_keys = [c for c, v in cardinality.items() if v > 0.99]
    high_cardinality = [c for c, v in cardinality.items() if v > 0.9]
    low_cardinality = [c for c, v in cardinality.items() if v < 0.1]

    # Structure type heuristic
    date_cols = {"date", "datetime", "timestamp", "created_at", "updated_at"}
    has_date = any(c.lower() in date_cols for c in df.columns)
    structure_type = "Transactional/Fact" if n_rows > 1000 and has_date else "Dimension/Reference"

    # Null percentage
    null_pct = {}
    for col in df.columns:
        nc = df[col].null_count()
        null_pct[col] = round(nc / n_rows * 100, 2) if n_rows > 0 else 0.0

    # Duplicate rows
    duplicate_rows = n_rows - df.unique().height

    # Statistics via describe() — convert to Pandas JSON for compatibility
    stats_json = json.loads(df.to_pandas().describe(include="all").to_json())

    return {
        "file": path,
        "rows": n_rows,
        "columns": n_cols,
        "structure_type": structure_type,
        "column_types": col_types,
        "potential_key_columns": potential_keys,
        "high_cardinality_cols": high_cardinality,
        "low_cardinality_cols": low_cardinality,
        "data_quality": {
            "null_percentage": null_pct,
            "duplicate_rows": duplicate_rows,
            "duplicate_percentage": round(duplicate_rows / n_rows * 100, 2) if n_rows > 0 else 0.0,
        },
        "statistics": stats_json,
        "engine": "polars",
    }


def _profile_pandas(path: str) -> dict:
    """Pandas fallback profiling (same logic as original profile_data)."""
    df = pd.read_csv(path)
    n_rows = len(df)

    cardinality = df.nunique() / n_rows if n_rows > 0 else df.nunique() * 0

    date_cols = {"date", "datetime", "timestamp", "created_at", "updated_at"}
    has_date = any(c.lower() in date_cols for c in df.columns)
    structure_type = "Transactional/Fact" if n_rows > 1000 and has_date else "Dimension/Reference"

    potential_keys = list(cardinality[cardinality > 0.99].index)
    null_pct = (df.isnull().sum() / n_rows * 100).round(2).to_dict() if n_rows > 0 else {}
    duplicate_rows = df.duplicated().sum()

    return {
        "file": path,
        "rows": n_rows,
        "columns": len(df.columns),
        "structure_type": structure_type,
        "column_types": {col: str(dtype) for col, dtype in df.dtypes.items()},
        "potential_key_columns": potential_keys,
        "high_cardinality_cols": list(cardinality[cardinality > 0.9].index),
        "low_cardinality_cols": list(cardinality[cardinality < 0.1].index),
        "data_quality": {
            "null_percentage": null_pct,
            "duplicate_rows": int(duplicate_rows),
            "duplicate_percentage": round(duplicate_rows / n_rows * 100, 2) if n_rows > 0 else 0.0,
        },
        "statistics": json.loads(df.describe(include="all").to_json()),
        "engine": "pandas",
    }


# ---------------------------------------------------------------------------
# Vectorized row hashing
# ---------------------------------------------------------------------------

def _pandas_row_hash(row: pd.Series, columns: list) -> str:
    """Exact replica of server.py compute_row_hash for backward compatibility."""
    values = "|".join(str(row[col]) for col in columns)
    return hashlib.sha256(values.encode()).hexdigest()[:16]


def compute_row_hashes_vectorized(df: pd.DataFrame, columns: list) -> pd.Series:
    """Vectorized row hashing — replaces df.apply(lambda row: ..., axis=1).

    Uses DuckDB SQL for vectorized SHA-256, falling back to Pandas apply.
    Produces hashes identical to compute_row_hash() in server.py.
    """
    if DuckDB_AVAILABLE_FOR_HASHING and len(df) > 0:
        try:
            return _hash_duckdb(df, columns)
        except Exception:
            pass

    if POLARS_AVAILABLE and len(df) > 0:
        try:
            return _hash_polars(df, columns)
        except Exception:
            pass

    # Pandas fallback — identical to original compute_row_hash
    return df.apply(lambda row: _pandas_row_hash(row, columns), axis=1)


# We gate DuckDB hashing separately because it must produce identical output
DuckDB_AVAILABLE_FOR_HASHING = DUCKDB_AVAILABLE


def _hash_duckdb(df: pd.DataFrame, columns: list) -> pd.Series:
    """DuckDB vectorized SHA-256 hashing matching compute_row_hash output.

    DuckDB's sha256() returns hex strings. We concatenate columns with '|' and
    truncate to 16 chars to match the Python implementation.
    """
    conn = get_duckdb()
    # Register the DataFrame as a temp table
    conn.register("_hash_input", df)
    try:
        # Build concatenation expression: CAST(col1 AS VARCHAR) || '|' || CAST(col2 AS VARCHAR) ...
        parts = [f"CAST(\"{col}\" AS VARCHAR)" for col in columns]
        concat_expr = " || '|' || ".join(parts)
        sql = f"SELECT SUBSTRING(sha256({concat_expr}), 1, 16) AS _h FROM _hash_input"
        result = conn.execute(sql).fetchdf()
        return result["_h"].reset_index(drop=True)
    finally:
        conn.unregister("_hash_input")


def _hash_polars(df: pd.DataFrame, columns: list) -> pd.Series:
    """Polars-based hashing fallback — still uses SHA-256 for compatibility.

    Polars concat_str + map to hashlib for exact output match.
    """
    pldf = pl.from_pandas(df[columns].astype(str))
    # Concatenate all columns with '|'
    concat_col = pl.concat_str([pl.col(c) for c in columns], separator="|")
    hashed = pldf.select(
        concat_col.alias("_concat")
    ).with_columns(
        pl.col("_concat").map_elements(
            lambda v: hashlib.sha256(v.encode()).hexdigest()[:16],
            return_dtype=pl.Utf8,
        ).alias("_hash")
    )["_hash"]
    return hashed.to_pandas().reset_index(drop=True)


# ---------------------------------------------------------------------------
# DuckDB local SQL
# ---------------------------------------------------------------------------

def query_local_sql(sql: str, register_files: dict | None = None) -> pd.DataFrame:
    """Execute SQL against local files via DuckDB.

    DuckDB handles read_csv_auto(), read_parquet(), read_json_auto() natively.
    Optionally register file paths as named tables first.
    """
    conn = get_duckdb()

    # Register files as named views if provided
    if register_files:
        for table_name, file_path in register_files.items():
            _register_file(conn, file_path, table_name)

    return conn.execute(sql).fetchdf()


def register_table(file_path: str, table_name: str) -> dict:
    """Register a local file as a named table in DuckDB for repeated SQL access.

    Returns: dict with row_count and column schema.
    """
    conn = get_duckdb()
    _register_file(conn, file_path, table_name)

    # Get row count and schema
    row_count = conn.execute(f'SELECT COUNT(*) FROM "{table_name}"').fetchone()[0]
    schema_rows = conn.execute(f'DESCRIBE "{table_name}"').fetchdf()
    schema = [
        {"column": r["column_name"], "type": r["column_type"]}
        for _, r in schema_rows.iterrows()
    ]
    return {
        "table_name": table_name,
        "file_path": file_path,
        "row_count": row_count,
        "schema": schema,
    }


def list_tables() -> list[dict]:
    """List all tables/views in the local DuckDB instance."""
    conn = get_duckdb()
    tables_df = conn.execute(
        "SELECT table_name, table_type FROM information_schema.tables "
        "WHERE table_schema = 'main'"
    ).fetchdf()

    result = []
    for _, row in tables_df.iterrows():
        name = row["table_name"]
        try:
            count = conn.execute(f'SELECT COUNT(*) FROM "{name}"').fetchone()[0]
        except Exception:
            count = -1
        try:
            schema_rows = conn.execute(f'DESCRIBE "{name}"').fetchdf()
            cols = [
                {"column": r["column_name"], "type": r["column_type"]}
                for _, r in schema_rows.iterrows()
            ]
        except Exception:
            cols = []
        result.append({
            "table_name": name,
            "table_type": row["table_type"],
            "row_count": count,
            "columns": cols,
        })
    return result


def export_to_parquet(sql_or_table: str, output_path: str) -> dict:
    """Export a DuckDB query result or table to Parquet format.

    Args:
        sql_or_table: A SQL query string or a registered table name.
        output_path: Destination .parquet file path.
    """
    conn = get_duckdb()
    output_path = str(Path(output_path).resolve())

    # Determine if it's a table name or SQL query
    if " " not in sql_or_table.strip():
        # Likely a table name
        source = f'SELECT * FROM "{sql_or_table}"'
    else:
        source = sql_or_table

    conn.execute(f"COPY ({source}) TO '{output_path}' (FORMAT PARQUET)")

    # Get stats
    row_count = conn.execute(f"SELECT COUNT(*) FROM ({source})").fetchone()[0]
    file_size = os.path.getsize(output_path)

    return {
        "output_path": output_path,
        "row_count": row_count,
        "file_size_bytes": file_size,
        "file_size_mb": round(file_size / (1024 * 1024), 2),
    }


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _register_file(conn, file_path: str, table_name: str):
    """Register a file as a DuckDB view based on its extension."""
    ext = Path(file_path).suffix.lower()
    path_escaped = file_path.replace("'", "''")
    if ext == ".parquet":
        conn.execute(
            f"CREATE OR REPLACE VIEW \"{table_name}\" AS "
            f"SELECT * FROM read_parquet('{path_escaped}')"
        )
    elif ext == ".json" or ext == ".jsonl":
        conn.execute(
            f"CREATE OR REPLACE VIEW \"{table_name}\" AS "
            f"SELECT * FROM read_json_auto('{path_escaped}')"
        )
    else:
        # Default: CSV
        conn.execute(
            f"CREATE OR REPLACE VIEW \"{table_name}\" AS "
            f"SELECT * FROM read_csv_auto('{path_escaped}')"
        )
