analysis/query_planner.py

"""
LLM-Powered Query Planner for Advanced Data Explorer.

Classifies user queries (FETCH / ANALYZE / EXPORT), decomposes complex ones
into structured ExecutionPlan JSON referencing the whitelisted operation registry.
"""
import json
from typing import Dict, List, Optional
from datetime import datetime
from core import config
from analysis.operations import get_registry_for_prompt


SYSTEM_PROMPT = """You are a query planner for a cryogenic hydrogen pump testing data system (CSH2 MURPHY). Your job is to decompose an engineer's natural language query into a structured execution plan.

You classify queries into three types:
- FETCH: Simple data retrieval (sensors + time range, no post-processing). Use for "show me", "plot", "display" type queries.
- ANALYZE: Data retrieval + one or more analysis operations. Use for "find", "identify", "detect", "count", "calculate" queries.
- EXPORT: Data retrieval + transformation + CSV export. Use for "export", "download", "save" queries with transformations.

CRITICAL RULES:
1. You MUST output valid JSON matching the ExecutionPlan schema below. No other text.
2. You MUST only use operations from the AVAILABLE_OPERATIONS list.
3. You MUST NOT generate arbitrary SQL, Python code, or operations not in the registry.
4. For FETCH queries, set operations to an empty array [].
5. Always include the correct sensors needed for each operation in data_requirements.sensors.
6. For fill/cycle analysis, always include motor speed tags (M130_Speed, MC130_VFD_Speed), PT130, FT140, and AvgPower.
7. Choose resolution wisely:
   - "raw": Only for sub-second analysis or frequency reduction from 10Hz. MAX 1 HOUR time range.
   - "1sec": For detailed analysis within a single day.
   - "15sec": For multi-day or full-range analysis.
   - "auto": Let the system decide based on time range duration.

ExecutionPlan JSON schema:
{
  "query_type": "FETCH | ANALYZE | EXPORT",
  "explanation": "Human-readable summary of what will happen",
  "data_requirements": {
    "sensors": ["PT130", "TT110"],
    "start_time": "YYYY-MM-DDTHH:MM:SS",
    "end_time": "YYYY-MM-DDTHH:MM:SS",
    "resolution": "raw | 1sec | 15sec | auto"
  },
  "operations": [
    {
      "op": "operation_name_from_registry",
      "params": { ... },
      "label": "Human-readable step description"
    }
  ],
  "export": {
    "enabled": true,
    "filename_hint": "descriptive_name"
  }
}
"""


EXAMPLES = """
WORKED EXAMPLES:

Example 1 — Query: "Find time blocks where PT130 exceeded 900 bar"
{
  "query_type": "ANALYZE",
  "explanation": "Find contiguous time blocks where discharge pressure PT130 exceeded 900 bar across the full data range",
  "data_requirements": {
    "sensors": ["PT130"],
    "start_time": "2025-03-14T00:00:00",
    "end_time": "2025-09-25T23:59:59",
    "resolution": "15sec"
  },
  "operations": [
    {
      "op": "find_peak_pressure_blocks",
      "params": {"pressure_tag": "PT130", "threshold_bar": 900.0},
      "label": "Find time blocks where PT130 > 900 bar"
    }
  ],
  "export": {"enabled": true, "filename_hint": "PT130_above_900bar"}
}

Example 2 — Query: "Find time blocks where both AOV140 and MC130 were constant for 5 minutes"
{
  "query_type": "ANALYZE",
  "explanation": "Find periods where both AOV140 (valve) and MC130_VFD_Speed (motor) were simultaneously constant for at least 5 minutes",
  "data_requirements": {
    "sensors": ["AOV140", "MC130_VFD_Speed"],
    "start_time": "2025-03-14T00:00:00",
    "end_time": "2025-09-25T23:59:59",
    "resolution": "15sec"
  },
  "operations": [
    {
      "op": "detect_constant_periods_multi",
      "params": {"tags": ["AOV140", "MC130_VFD_Speed"], "tolerance": 0.01, "min_duration_minutes": 5.0},
      "label": "Find overlapping constant periods for AOV140 and MC130_VFD_Speed"
    }
  ],
  "export": {"enabled": true, "filename_hint": "constant_AOV140_MC130"}
}

Example 3 — Query: "Export PT130, TT110, FT140 at 2Hz as CSV for Sep 25 10am-11am"
{
  "query_type": "EXPORT",
  "explanation": "Export PT130, TT110, FT140 at 2Hz (downsampled from 10Hz raw data) for Sep 25 10-11am",
  "data_requirements": {
    "sensors": ["PT130", "TT110", "FT140"],
    "start_time": "2025-09-25T10:00:00",
    "end_time": "2025-09-25T11:00:00",
    "resolution": "raw"
  },
  "operations": [
    {
      "op": "downsample_frequency",
      "params": {"target_hz": 2.0},
      "label": "Downsample from 10Hz to 2Hz"
    }
  ],
  "export": {"enabled": true, "filename_hint": "PT130_TT110_FT140_2Hz"}
}

Example 4 — Query: "Export the 3 testing windows May 5-6 where PT130 ramped from 100 to 800 bar"
{
  "query_type": "EXPORT",
  "explanation": "Find 3 testing windows in May 5-6 where PT130 ramped from 100 to 800 bar, export each window's data",
  "data_requirements": {
    "sensors": ["PT130", "TT110", "TT130", "FT140", "M130_Speed", "AOV140"],
    "start_time": "2025-05-05T00:00:00",
    "end_time": "2025-05-06T23:59:59",
    "resolution": "1sec"
  },
  "operations": [
    {
      "op": "detect_pressure_ramps",
      "params": {"pressure_tag": "PT130", "start_bar": 100.0, "end_bar": 800.0, "max_ramps": 3},
      "label": "Detect pressure ramps from 100 to 800 bar"
    },
    {
      "op": "extract_windows",
      "params": {"source": "previous_result"},
      "label": "Extract data for each ramp window"
    }
  ],
  "export": {"enabled": true, "filename_hint": "LN2_ramps_100_800bar"}
}

Example 5 — Query: "Identify fills where kWh per kg exceeded 0.6"
{
  "query_type": "ANALYZE",
  "explanation": "Detect all testing fills, compute kWh/kg for each, return those exceeding 0.6 kWh/kg",
  "data_requirements": {
    "sensors": ["M130_Speed", "MC130_VFD_Speed", "PT130", "FT140", "AvgPower", "TT110", "TT130"],
    "start_time": "2025-03-14T00:00:00",
    "end_time": "2025-09-25T23:59:59",
    "resolution": "15sec"
  },
  "operations": [
    {
      "op": "detect_and_analyze_fills",
      "params": {"kwh_per_kg_threshold": 0.6, "include_strokes": false},
      "label": "Detect fills and filter by kWh/kg > 0.6"
    }
  ],
  "export": {"enabled": true, "filename_hint": "fills_above_0.6_kwhkg"}
}

Example 6 — Query: "In Fill 3, what was the total number of pump strokes?"
{
  "query_type": "ANALYZE",
  "explanation": "Detect fills, find Fill 3, compute its total pump strokes from motor RPM",
  "data_requirements": {
    "sensors": ["M130_Speed", "MC130_VFD_Speed", "PT130", "FT140", "AvgPower"],
    "start_time": "2025-03-14T00:00:00",
    "end_time": "2025-09-25T23:59:59",
    "resolution": "15sec"
  },
  "operations": [
    {
      "op": "detect_and_analyze_fills",
      "params": {"fill_id": 3, "include_strokes": true},
      "label": "Find Fill 3 and compute pump strokes"
    }
  ],
  "export": {"enabled": false}
}

Example 7 — Query: "Between Sep 24 15:00 and 18:00, total pump strokes?"
{
  "query_type": "ANALYZE",
  "explanation": "Count total pump revolutions between Sep 24 15:00 and 18:00 UTC",
  "data_requirements": {
    "sensors": ["M130_Speed", "MC130_VFD_Speed"],
    "start_time": "2025-09-24T15:00:00",
    "end_time": "2025-09-24T18:00:00",
    "resolution": "1sec"
  },
  "operations": [
    {
      "op": "count_pump_strokes",
      "params": {},
      "label": "Count total pump revolutions in 3-hour window"
    }
  ],
  "export": {"enabled": false}
}
"""


class QueryPlanner:
    """LLM-powered query classification and decomposition into ExecutionPlan JSON."""

    def __init__(self):
        self.api_key = config.ANTHROPIC_API_KEY
        self.api_available = bool(self.api_key and self.api_key != 'your_api_key_here')
        self.model = "claude-sonnet-4-20250514"
        self._client = None

    @property
    def client(self):
        if self._client is None and self.api_available:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self.api_key)
        return self._client

    def plan(self, user_query: str, available_sensors: List[str]) -> Optional[Dict]:
        """
        Decompose a natural language query into a structured execution plan.

        Uses Anthropic prompt caching: the system prompt and static context
        (sensor catalog, operation registry, worked examples) are cached for
        5 minutes, so repeated queries only pay for the dynamic user query.

        Returns:
            Dict (ExecutionPlan) or dict with 'error' key on failure.
        """
        if not self.api_available:
            return {"error": "Claude API key not configured."}

        static_context = self._build_static_context(available_sensors)
        dynamic_query = f'User query: "{user_query}"\n\nReturn ONLY the JSON execution plan. No other text.'

        try:
            message = self.client.messages.create(
                model=self.model,
                max_tokens=2000,
                system=[
                    {
                        "type": "text",
                        "text": SYSTEM_PROMPT,
                        "cache_control": {"type": "ephemeral"},
                    }
                ],
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": static_context,
                                "cache_control": {"type": "ephemeral"},
                            },
                            {
                                "type": "text",
                                "text": dynamic_query,
                            },
                        ],
                    }
                ],
            )
            response_text = message.content[0].text
            plan = self._extract_json(response_text)

            if plan:
                validation = self._validate_plan(plan, available_sensors)
                if validation:
                    return {"error": validation, "raw_plan": plan}
                return self._process_plan(plan, available_sensors)

            # Fallback: try V1 parser for simple queries
            return self._fallback_parse(user_query, available_sensors)

        except Exception as e:
            return {"error": f"Planning failed: {e}"}

    def _build_static_context(self, available_sensors: List[str]) -> str:
        """Build the cacheable static context block.

        Contains sensor catalog, groups, aliases, operation registry, and
        worked examples. This is ~25-30KB of stable text that changes only
        when sensors or operations are added — ideal for Anthropic's prompt
        caching (5-minute TTL, 90% cost reduction on cache hit).
        """
        # Sensor reference with descriptions
        sensor_lines = []
        for tag_name, info in config.SENSOR_TAGS.items():
            desc = info.get("description", "")
            units = info.get("units", "")
            sensor_lines.append(f"  {tag_name}: {desc} ({units})")
        sensor_reference = "\n".join(sensor_lines)

        # Sensor group reference
        group_lines = []
        for group_name, tags in config.SENSOR_GROUPS.items():
            group_lines.append(f"  {group_name}: {', '.join(tags)}")
        group_reference = "\n".join(group_lines)

        # Operation registry
        registry_text = get_registry_for_prompt()

        current_date = datetime.now()

        return f"""Current date/time: {current_date.strftime('%Y-%m-%d %H:%M:%S')}
Data is available from March 14, 2025 to September 25, 2025.

AVAILABLE SENSORS (tag name: description (units)):
{sensor_reference}

SENSOR GROUPS:
{group_reference}

COMMON ALIASES:
- "pressure" or "pressures" -> PT130 (Discharge), PT01T (Cryotank), PT110 (Inlet)
- "temperature" or "temperatures" -> TT110 (Pump Feed Line), TT130 (Discharge)
- "flow" -> FT140 (Flow Transmitter)
- "motor" or "speed" or "rpm" -> M130_Speed (Motor Speed)
- "power" -> AvgPower, PeakPower, PowerConsumption
- "discharge pressure" -> PT130
- "inlet pressure" or "supply pressure" -> PT110
- "tank pressure" or "cryotank pressure" -> PT01T
- "discharge temperature" -> TT130
- "feed temperature" or "pump temperature" -> TT110
- "valve" -> AOV140
- "fill" or "fills" or "test cycle" -> requires detect_and_analyze_fills operation
- "strokes" or "pump strokes" or "revolutions" -> requires count_pump_strokes or detect_and_analyze_fills
- "efficiency" or "kWh/kg" or "specific energy" -> requires detect_and_analyze_fills with AvgPower + FT140

{registry_text}

{EXAMPLES}"""

    def _build_prompt(self, user_query: str, available_sensors: List[str]) -> str:
        """Build the full user prompt (used by fallback path).

        The main plan() method uses _build_static_context() + dynamic query
        separately for prompt caching. This method is kept for backward
        compatibility with the fallback parser.
        """
        static = self._build_static_context(available_sensors)
        return f"""{static}

User query: "{user_query}"

Return ONLY the JSON execution plan. No other text."""

    def _extract_json(self, text: str) -> Optional[Dict]:
        """Extract JSON object from LLM response."""
        try:
            start = text.find("{")
            end = text.rfind("}") + 1
            if start != -1 and end > start:
                json_str = text[start:end]
                return json.loads(json_str)
            return None
        except json.JSONDecodeError:
            return None

    def _validate_plan(self, plan: Dict, available_sensors: List[str]) -> Optional[str]:
        """Validate plan structure. Returns error string or None if valid."""
        # Check required top-level fields
        if "query_type" not in plan:
            return "Missing 'query_type' field"
        if plan["query_type"] not in ("FETCH", "ANALYZE", "EXPORT"):
            return f"Invalid query_type: {plan['query_type']}"
        if "data_requirements" not in plan:
            return "Missing 'data_requirements' field"

        dr = plan["data_requirements"]
        if "sensors" not in dr or not dr["sensors"]:
            return "No sensors specified"
        if "start_time" not in dr or "end_time" not in dr:
            return "Missing start_time or end_time"

        # Validate resolution
        resolution = dr.get("resolution", "auto")
        if resolution not in ("raw", "1sec", "15sec", "auto"):
            return f"Invalid resolution: {resolution}. Use: raw, 1sec, 15sec, auto"

        # Safety: raw table time range limit
        if resolution == "raw":
            try:
                start = datetime.fromisoformat(dr["start_time"])
                end = datetime.fromisoformat(dr["end_time"])
                duration_h = (end - start).total_seconds() / 3600
                if duration_h > 1.0:
                    return f"Raw resolution limited to 1 hour. Requested {duration_h:.1f} hours. Use '1sec' or '15sec' for longer ranges."
            except (ValueError, TypeError):
                pass

        # Validate operations exist in registry
        from analysis.operations import OPERATION_REGISTRY
        for op in plan.get("operations", []):
            if op.get("op") not in OPERATION_REGISTRY:
                known = ", ".join(OPERATION_REGISTRY.keys())
                return f"Unknown operation: '{op['op']}'. Available: {known}"

        return None

    def _process_plan(self, plan: Dict, available_sensors: List[str]) -> Dict:
        """Process and clean up the plan: resolve sensors, parse datetimes."""
        dr = plan.get("data_requirements", {})

        # Validate sensors against available list (case-insensitive)
        sensor_map = {s.upper(): s for s in available_sensors}
        valid_sensors = []
        for s in dr.get("sensors", []):
            matched = sensor_map.get(s.upper())
            if matched:
                valid_sensors.append(matched)

        if not valid_sensors:
            return {"error": f"No matching sensors found. Requested: {dr.get('sensors', [])}"}

        # Parse datetimes
        try:
            start_time = datetime.fromisoformat(dr["start_time"])
            end_time = datetime.fromisoformat(dr["end_time"])
        except (ValueError, TypeError, KeyError) as e:
            return {"error": f"Invalid time format: {e}"}

        # Clean up and return
        plan["data_requirements"]["sensors"] = valid_sensors
        plan["data_requirements"]["start_time"] = start_time
        plan["data_requirements"]["end_time"] = end_time
        plan["data_requirements"]["resolution"] = dr.get("resolution", "auto")

        return plan

    def _fallback_parse(self, user_query: str, available_sensors: List[str]) -> Optional[Dict]:
        """Fallback: attempt to parse as simple FETCH query via V1 NLQueryParser."""
        try:
            from analysis.nl2sql import NLQueryParser
            parser = NLQueryParser()
            result = parser.parse(user_query, available_sensors)
            if result and "error" not in result:
                return {
                    "query_type": "FETCH",
                    "explanation": result.get("explanation", ""),
                    "data_requirements": {
                        "sensors": result["sensors"],
                        "start_time": result["start_time"],
                        "end_time": result["end_time"],
                        "resolution": "auto",
                    },
                    "operations": [],
                    "export": {"enabled": False},
                }
            return result  # Pass through error
        except Exception as e:
            return {"error": f"Fallback parse failed: {e}"}