analysis/execution_engine.py

"""
Execution Engine for Advanced Data Explorer.

Validates an ExecutionPlan, fetches the required data, dispatches
operations sequentially via OperationDispatcher, and packages results.
"""
import time
import pandas as pd
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field
from datetime import datetime

from core import config
from core.db_connector import get_db_connector
from core.cached_queries import cached_get_pivot_data
from analysis.operations import OperationDispatcher, OperationOutput


# ── Result Dataclasses ───────────────────────────────────────────

@dataclass
class ExecutionResult:
    """Complete result of executing a plan."""
    success: bool
    query_type: str
    explanation: str
    source_data: Optional[pd.DataFrame] = None
    operation_results: List[OperationOutput] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    execution_time_seconds: float = 0.0
    table_used: str = ""
    row_count: int = 0


# ── Table label mapping (for display) ───────────────────────────

TABLE_LABELS = {
    config.TABLES['raw']: 'Raw (10Hz)',
    config.TABLES['agg_1sec']: '1-second aggregates',
    config.TABLES['agg_15sec']: '15-second aggregates',
}


# ── Execution Engine ─────────────────────────────────────────────

class ExecutionEngine:
    """
    Validate an ExecutionPlan, fetch data, run operations, package results.

    Usage:
        engine = ExecutionEngine()
        result = engine.execute(plan)
    """

    def __init__(self):
        self.dispatcher = OperationDispatcher()

    def execute(self, plan: Dict) -> ExecutionResult:
        """Execute a validated plan end-to-end."""
        t0 = time.time()
        errors = []
        warnings = []

        # ── 1. Validate plan structure ──
        qt = plan.get("query_type", "FETCH")
        explanation = plan.get("explanation", "")
        dr = plan.get("data_requirements", {})
        operations = plan.get("operations", [])

        sensors = dr.get("sensors", [])
        start_time = dr.get("start_time")
        end_time = dr.get("end_time")
        resolution = dr.get("resolution", "auto")

        if not sensors:
            return self._fail("No sensors specified in plan", qt, explanation, t0)

        if not start_time or not end_time:
            return self._fail("Missing start_time or end_time in plan", qt, explanation, t0)

        # Ensure datetimes are datetime objects
        if isinstance(start_time, str):
            start_time = datetime.fromisoformat(start_time)
        if isinstance(end_time, str):
            end_time = datetime.fromisoformat(end_time)

        # ── 2. Resolve table from resolution ──
        table_override = None
        if resolution in config.RESOLUTION_MAP:
            table_override = config.RESOLUTION_MAP[resolution]
        elif resolution == "auto":
            # Let the DB connector decide via smart routing
            table_override = None
        else:
            warnings.append(f"Unknown resolution '{resolution}', using auto routing.")

        # ── 3. Safety: raw table time-range cap ──
        if table_override == config.TABLES['raw']:
            duration_h = (end_time - start_time).total_seconds() / 3600
            if duration_h > 1.0:
                return self._fail(
                    f"Raw resolution limited to 1 hour. Requested {duration_h:.1f} hours. "
                    f"Use '1sec' or '15sec' for longer ranges.",
                    qt, explanation, t0,
                )

        # ── 4. Fetch data ──
        try:
            db = get_db_connector()
            df = cached_get_pivot_data(
                tag_names=tuple(sensors),
                start_time=start_time,
                end_time=end_time,
                table_override=table_override,
            )
        except Exception as e:
            return self._fail(f"Data fetch failed: {e}", qt, explanation, t0)

        if df is None or df.empty:
            return self._fail(
                "No data returned for the specified sensors and time range. "
                "Check that the sensors exist and the time range is within March 14 – September 25, 2025.",
                qt, explanation, t0,
            )

        # Determine which table was actually used (for display)
        if table_override:
            table_used = table_override
        else:
            table_used = db._select_table(start_time, end_time)

        table_label = TABLE_LABELS.get(table_used, table_used)

        # Check which sensors actually came back
        returned_sensors = [s for s in sensors if s in df.columns]
        missing = [s for s in sensors if s not in df.columns and s != "timestamp"]
        if missing:
            warnings.append(f"Sensors not found in data: {', '.join(missing)}")

        # ── 5. Run operations sequentially ──
        op_results = []
        prev_result = None

        for i, op_spec in enumerate(operations):
            op_name = op_spec.get("op", "")
            op_params = op_spec.get("params", {})
            op_label = op_spec.get("label", op_name)

            try:
                output = self.dispatcher.run(
                    op_name=op_name,
                    df=df,
                    params=op_params,
                    prev_result=prev_result,
                    db_connector=get_db_connector(),
                )
                op_results.append(output)
                prev_result = output

                # Check for operation-level errors
                if output.metadata.get("error"):
                    warnings.append(f"Operation '{op_label}': {output.metadata['error']}")

            except Exception as e:
                error_msg = f"Operation '{op_label}' failed: {e}"
                errors.append(error_msg)
                # Create a placeholder output so downstream chaining can detect the failure
                op_results.append(OperationOutput(
                    op_name=op_name,
                    label=f"{op_label} (FAILED)",
                    data=None,
                    metadata={"error": str(e)},
                ))
                prev_result = op_results[-1]

        # ── 6. Package result ──
        elapsed = time.time() - t0
        success = len(errors) == 0

        return ExecutionResult(
            success=success,
            query_type=qt,
            explanation=explanation,
            source_data=df,
            operation_results=op_results,
            errors=errors,
            warnings=warnings,
            execution_time_seconds=round(elapsed, 2),
            table_used=table_label,
            row_count=len(df),
        )

    @staticmethod
    def _fail(msg: str, qt: str, explanation: str, t0: float) -> ExecutionResult:
        """Return a failed ExecutionResult."""
        return ExecutionResult(
            success=False,
            query_type=qt,
            explanation=explanation,
            errors=[msg],
            execution_time_seconds=round(time.time() - t0, 2),
        )