core/execution/mvcc_executor.py

"""
MVCC Executor.
Performs lock-free versioned operations on Parquet files using Polars LazyFrames.
"""
from __future__ import annotations

import os
import polars as pl

from config import DATA_DIR, LAZY_THRESHOLD_BYTES
from services.session_manager import session_manager
from core.column_registry import column_registry
from core.execution.polars_builder import build_polars_expression

def _build_predicate(column: str, condition: str, value) -> pl.Expr:
    """Build a Polars expression from a comparison condition string."""
    col = pl.col(column)
    if condition == ">":
        return col > value
    elif condition == "<":
        return col < value
    elif condition == ">=":
        return col >= value
    elif condition == "<=":
        return col <= value
    elif condition == "!=":
        return col != value
    elif condition == "==" or condition == "=":
        if isinstance(value, (int, float)):
            return col == value
        try:
            return col == float(value)
        except (ValueError, TypeError):
            return col.cast(pl.String).str.to_lowercase() == str(value).lower()
    elif condition in ("contains", "~"):
        return col.cast(pl.String).str.contains(str(value), literal=True)
    else:
        return col.cast(pl.String).str.to_lowercase() == str(value).lower()

def _format_number(val) -> str:
    """Pretty-print numbers with commas."""
    if val is None:
        return "N/A"
    if isinstance(val, float):
        if val == int(val):
            return f"{int(val):,}"
        return f"{val:,.2f}"
    if isinstance(val, int):
        return f"{val:,}"
    return str(val)

def execute_mvcc_command(session_id: str, intent: dict) -> dict:
    """Execute the command and return a result dict.

    Updates metadata and creates a new parquet version on write.
    """
    meta = session_manager.get(session_id)
    if meta is None:
        return {
            "status": "error",
            "message": "Session nahi mila. Pehle file upload karo.",
            "diff": {}
        }

    version = getattr(meta, "current_version", 0)

    # Resolve read path
    # If version is 0, we can fall back to the initial upload path (data/sessions/{session_id}.parquet).
    v_path = os.path.join(DATA_DIR, f"{session_id}_v{version}.parquet")
    if not os.path.exists(v_path):
        if version == 0:
            v_path = os.path.join(DATA_DIR, f"{session_id}.parquet")
            if not os.path.exists(v_path):
                raise FileNotFoundError(f"Initial Parquet file not found for session {session_id}")
        else:
            raise FileNotFoundError(f"Parquet file not found for version {version} of session {session_id}")

    op = intent.get("operation")
    if not op:
        return {
            "message": "Operation specifies nahi kiya gaya",
            "diff": {}
        }

    # ── Read-only operations (Aggregates) ──
    if op in ("sum", "average", "count", "min", "max"):
        lf = pl.scan_parquet(v_path)
        if op == "sum":
            col = intent["column"]
            result = lf.select(pl.col(col).sum()).collect().item()
            return {
                "message": f"{col} ka sum = {_format_number(result)}",
                "diff": {"operation": "sum", "column": col, "result": result},
            }
        elif op == "average":
            col = intent["column"]
            result = lf.select(pl.col(col).mean()).collect().item()
            return {
                "message": f"{col} ka average = {_format_number(result)}",
                "diff": {"operation": "average", "column": col, "result": result},
            }
        elif op == "count":
            col = intent.get("column")
            if col:
                result = lf.select(pl.col(col).count()).collect().item()
                return {
                    "message": f"{col} mein {_format_number(result)} non-null values hain",
                    "diff": {"operation": "count", "column": col, "result": result},
                }
            else:
                result = lf.select(pl.len()).collect().item()
                return {
                    "message": f"Total {_format_number(result)} rows hain",
                    "diff": {"operation": "count", "column": None, "result": result},
                }
        elif op == "min":
            col = intent["column"]
            result = lf.select(pl.col(col).min()).collect().item()
            return {
                "message": f"{col} ka minimum = {_format_number(result)}",
                "diff": {"operation": "min", "column": col, "result": result},
            }
        elif op == "max":
            col = intent["column"]
            result = lf.select(pl.col(col).max()).collect().item()
            return {
                "message": f"{col} ka maximum = {_format_number(result)}",
                "diff": {"operation": "max", "column": col, "result": result},
            }

    # ── Write operations ──
    lf = pl.scan_parquet(v_path)
    new_version = version + 1
    new_path = os.path.join(DATA_DIR, f"{session_id}_v{new_version}.parquet")

    before_count = lf.select(pl.len()).collect().item()

    if op in ("increase", "decrease"):
        col = intent["column"]
        value = intent.get("value")
        if value is None:
            value = 10.0
        is_percent = intent.get("is_percent", True)
        
        expr = build_polars_expression(op, col, {"value": value, "is_percent": is_percent})
        lf = lf.with_columns(expr)

    elif op == "cast_type":
        col = intent["column"]
        target = intent["target_dtype"]
        expr = build_polars_expression("cast_type", col, {"target_dtype": target})
        lf = lf.with_columns(expr)

    elif op == "find_replace":
        col = intent["column"]
        old_val = intent["old_value"]
        new_val = intent["new_value"]
        
        schema = lf.collect_schema()
        dtype = schema[col]
        is_numeric = dtype in (pl.Float64, pl.Int64, pl.Int32, pl.Float32)
        
        if is_numeric:
            try:
                float(old_val)
                float(new_val)
            except (ValueError, TypeError):
                return {
                    "message": f"'{col}' numeric hai, replace values bhi numbers hone chahiye",
                    "diff": {},
                }
        
        expr = build_polars_expression("find_replace", col, {
            "old_value": old_val,
            "new_value": new_val,
            "is_numeric": is_numeric
        })
        lf = lf.with_columns(expr)

    elif op == "delete_column":
        col = intent["column"]
        lf = lf.drop(col)

    elif op == "rename_column":
        old = intent["column"]
        new = intent["new_name"]
        lf = lf.rename({old: new})

    elif op == "add_column":
        new_col = intent.get("new_column_name") or intent.get("column", "new_col")
        default_val = intent.get("value", 0)
        lf = lf.with_columns(pl.lit(default_val).alias(new_col))

    elif op == "remove_duplicates":
        col = intent.get("column")
        if col:
            lf = lf.unique(subset=[col], keep="first")
        else:
            lf = lf.unique(keep="first")

    elif op == "filter":
        col = intent["column"]
        condition = intent["condition"]
        filter_value = intent["filter_value"]
        predicate = _build_predicate(col, condition, filter_value)
        lf = lf.filter(predicate)

    elif op in ("sort_asc", "sort_desc"):
        col = intent["column"]
        descending = (op == "sort_desc")
        lf = lf.sort(col, descending=descending)

    else:
        return {
            "message": f"'{op}' operation supported nahi hai",
            "diff": {},
        }

    # Lock-free versioned write
    prev_size = os.path.getsize(v_path)
    try:
        if prev_size > LAZY_THRESHOLD_BYTES:
            lf.sink_parquet(new_path)
        else:
            lf.collect().write_parquet(new_path)
    except Exception:
        # Fallback to eager write if streaming sink is not supported for this query plan (e.g. sort/unique)
        lf.collect().write_parquet(new_path)

    # Read updated metadata
    lf_new = pl.scan_parquet(new_path)
    new_schema = lf_new.collect_schema()
    after_count = lf_new.select(pl.len()).collect().item()
    columns_meta = [{"name": name, "dtype": str(dtype)} for name, dtype in new_schema.items()]

    # Format result message and diff
    message = ""
    diff = {}

    if op in ("increase", "decrease"):
        col = intent["column"]
        value = intent.get("value")
        if value is None:
            value = 10.0
        is_percent = intent.get("is_percent", True)
        op_label = f"{value}% badha diya" if op == "increase" else f"{value}% ghata diya"
        if not is_percent:
            op_label = f"{value} joda" if op == "increase" else f"{value} ghata diya"
        message = f"{col} ko {op_label} ({after_count:,} rows updated)"
        diff = {
            "operation": op,
            "column": col,
            "value": value,
            "is_percent": is_percent,
            "affected_rows": after_count,
        }

    elif op == "cast_type":
        col = intent["column"]
        target = intent["target_dtype"]
        message = f"'{col}' ka type {target} mein change ho gaya"
        diff = {"operation": "cast_type", "column": col, "target_dtype": target}

    elif op == "find_replace":
        col = intent["column"]
        old_val = intent["old_value"]
        new_val = intent["new_value"]
        
        dtype = new_schema[col]
        is_numeric = dtype in (pl.Float64, pl.Int64, pl.Int32, pl.Float32)
        if is_numeric:
            try:
                count = lf_new.filter(pl.col(col) == float(new_val)).select(pl.len()).collect().item()
            except (ValueError, TypeError):
                count = "N/A"
        else:
            count = lf_new.filter(pl.col(col).cast(pl.String) == str(new_val)).select(pl.len()).collect().item()
            
        message = f"'{col}' mein '{old_val}' ko '{new_val}' se replace kiya ({count} rows changed)"
        diff = {
            "operation": "find_replace",
            "column": col,
            "old_value": old_val,
            "new_value": new_val,
            "affected_rows": count,
        }

    elif op == "delete_column":
        col = intent["column"]
        message = f"'{col}' column delete ho gaya"
        diff = {"operation": "delete_column", "column": col}

    elif op == "rename_column":
        old = intent["column"]
        new = intent["new_name"]
        message = f"'{old}' ka naam badal ke '{new}' ho gaya"
        diff = {"operation": "rename_column", "old_name": old, "new_name": new}

    elif op == "add_column":
        new_col = intent.get("new_column_name") or intent.get("column", "new_col")
        default_val = intent.get("value", 0)
        message = f"Naya column '{new_col}' add ho gaya (default: {default_val})"
        diff = {"operation": "add_column", "column": new_col}

    elif op == "remove_duplicates":
        col = intent.get("column")
        removed = before_count - after_count
        if col:
            message = f"'{col}' ke duplicate rows hata diye ({removed} rows removed, {after_count:,} remaining)"
        else:
            message = f"Duplicate rows hata diye ({removed} rows removed, {after_count:,} remaining)"
        diff = {
            "operation": "remove_duplicates",
            "column": col,
            "rows_removed": removed,
            "rows_remaining": after_count,
        }

    elif op == "filter":
        col = intent["column"]
        condition = intent["condition"]
        filter_value = intent["filter_value"]
        message = f"{col} {condition} {filter_value} → {after_count:,} rows bach gaye"
        diff = {
            "operation": "filter",
            "column": col,
            "condition": condition,
            "filter_value": filter_value,
            "rows_after": after_count,
        }

    elif op in ("sort_asc", "sort_desc"):
        col = intent["column"]
        descending = (op == "sort_desc")
        order = "descending (bada se chhota)" if descending else "ascending (chhota se bada)"
        message = f"{col} ko {order} sort kiya ({after_count:,} rows)"
        diff = {
            "operation": "sort",
            "column": col,
            "descending": descending,
            "affected_rows": after_count,
        }

    # Update session metadata in session_manager
    meta.current_version = new_version
    meta.row_count = after_count
    meta.columns = columns_meta
    meta.touch()

    # Update column registry
    column_registry.register(session_id, [c["name"] for c in columns_meta])

    # Clean up the previous intermediate versioned file
    if version > 0:
        prev_version_path = os.path.join(DATA_DIR, f"{session_id}_v{version}.parquet")
        if os.path.exists(prev_version_path):
            try:
                os.remove(prev_version_path)
            except Exception as e:
                print(f"[MVCC Executor] Warning cleaning up version {version}: {e}")

    return {
        "message": message,
        "diff": diff,
    }