app/services/prompt_service.py

"""
Prompt construction service for AI-driven dataset analysis.

Provides canned system prompts and helpers to build context-aware
prompts for LLM-based CSV/dataframe analysis and visualization.
"""

from __future__ import annotations

import json
from typing import Any, Dict, List, Optional

# ---------------------------------------------------------------------------
# Canned System Prompts
# ---------------------------------------------------------------------------

CSV_SYSTEM_PROMPT: str = """\
You are a Senior Data Analyst AI and CSV analysis assistant. Your goal is to extract actionable insights, perform statistical analysis, answer complex questions, and generate professional visualizations using the provided dataset.

The pandas DataFrame is pre-loaded as 'df' - use this variable.

STRICT OPERATIONAL REQUIREMENTS:
1. NEVER guess, predict, or estimate values yourself. ALWAYS generate executable Python code to calculate precise answers.
2. USE THE EXISTING 'df' - Do not attempt to reload or recreate the dataframe.
3. VARIABLE ASSIGNMENT IS MANDATORY: Every result, calculation, filtered subset, or visualization must be assigned to a descriptive, snake_case variable name.
4. JSON FOR STRUCTURED DATA: For any data structures (Lists, Records, Tables, Dictionaries, etc.), return them as JSON with correct indentation so the UI can parse it.
5. CLEANLINESS: If the analysis requires handling missing values (NaNs) or data cleaning, perform it on a copy (e.g., 'cleaned_df') before analyzing.

ANALYSIS GUIDELINES:
- Descriptive Statistics: Use .describe(), .value_counts(), and .nunique().
- Relationships: Calculate correlations using .corr() or group data using .groupby().
- Filtering: Always store filtered results in a specific variable (e.g., 'high_value_customers = ...').
- Aggregation: When grouping, reset indices (.reset_index()) to keep results in a flat, readable format.
- Outliers: Use IQR or Z-score methods when asked to find anomalies.

VISUALIZATION STANDARDS:
- Use matplotlib/seaborn only.
- Professional quality: proper sizing, labels, titles.
- Figure size: (14, 8) for complex charts, (12, 6) for simple charts.
- Fonts: Clear titles (fontsize=16), labels (fontsize=14).
- Ticks: Rotate x-labels if needed (45°, fontsize=12).
- Aesthetics: Add annotations/gridlines where helpful; use colorblind-friendly palettes.
- Final Step: Always include plt.tight_layout() and plt.show().
- Variable Assignment: Assign figure/axis objects when needed (e.g., fig, ax = plt.subplots...).

VARIABLE ASSIGNMENT RULES:
1. Every operation must store its result in a variable.
2. Variable names should be descriptive and snake_case.
3. For DataFrame operations: result_df = df.operation()
4. For statistical results: summary_stats = df.describe(include='all')
5. For filtered data: filtered_data = df[df['column'] > value]
6. For grouped analysis: revenue_by_region = df.groupby('region')['revenue'].sum().reset_index()
7. For correlation matrices: correlation_matrix = df.corr(numeric_only=True)
8. For visualizations: fig, ax = plt.subplots(...)

Return complete, executable code that follows these rules.
Your response should be modular, precise, and favor variable assignment over direct printing."""


CSV_STRICT_OUTPUT_PROMPT: str = """\
### 2. STRICT OUTPUT FORMAT
Return your response ONLY as a JSON object.

- **If the user asks for analysis/charts:** Fill "analyze" and "visualization" arrays with Python code.
- **If the user greets you or asks a generic question:** Use the "message" field for your response and keep the arrays empty.

```json
{
  "analyze": [
    {
      "description": "Short explanation of the math",
      "python_code": "# Clean data first\\ndf['col'] = ...\\n\\n# Perform analysis\\nresult = df.groupby..."
    }
  ],
  "visualization": [
    {
      "description": "Short explanation of the chart",
      "python_code": "# Clean data first\\ndf['col'] = ...\\n\\n# Plot\\nplt.figure(figsize=(12,6))\\nsns.barplot(data=df, ...)"
    }
  ],
  "message": "Fill this ONLY if the user is greeting, asking non-data questions or asking for wrong information."
}
```"""


# ---------------------------------------------------------------------------
# Prompt Builder Helpers
# ---------------------------------------------------------------------------

def build_csv_context_prompt(
    metadata: Dict[str, Any],
    user_message: str,
    *,
    include_system_prompt: bool = True,
    include_strict_output: bool = True,
) -> str:
    """
    Build a complete prompt for LLM-based CSV analysis from extracted metadata.

    Parameters
    ----------
    metadata : dict
        The metadata dictionary returned by ``extract_metadata()``.
    user_message : str
        The user's natural-language query about the dataset.
    include_system_prompt : bool
        Whether to prepend the ``CSV_SYSTEM_PROMPT``.
    include_strict_output : bool
        Whether to append the ``CSV_STRICT_OUTPUT_PROMPT``.

    Returns
    -------
    str
        A fully assembled prompt string ready to send to an LLM.
    """
    parts: List[str] = []

    if include_system_prompt:
        parts.append(f"[SYSTEM PROMPT]\n{CSV_SYSTEM_PROMPT}")

    # ---- Dataset context section ----
    shape = metadata.get("shape", {})
    columns = metadata.get("columns", [])
    dtypes = metadata.get("dtypes", {})
    sample_data = metadata.get("sample_data", [])
    numeric_cols = metadata.get("numeric_columns", [])
    categorical_cols = metadata.get("categorical_columns", [])
    datetime_cols = metadata.get("datetime_columns", [])
    boolean_cols = metadata.get("boolean_columns", [])

    context_lines = ["\nCSV Info:"]
    context_lines.append(f"- Shape: {shape.get('rows', '?')} rows x {shape.get('columns', '?')} cols")
    context_lines.append(f"- Columns: {', '.join(columns)}")
    context_lines.append(f"- Data Types: {json.dumps(dtypes)}")

    if numeric_cols:
        context_lines.append(f"- Numeric Columns: {', '.join(numeric_cols)}")
    if categorical_cols:
        context_lines.append(f"- Categorical Columns: {', '.join(categorical_cols)}")
    if datetime_cols:
        context_lines.append(f"- Datetime Columns: {', '.join(datetime_cols)}")
    if boolean_cols:
        context_lines.append(f"- Boolean Columns: {', '.join(boolean_cols)}")

    if sample_data:
        context_lines.append(f"- Sample Data: {json.dumps(sample_data, indent=2)}")

    parts.append("\n".join(context_lines))

    # ---- User message ----
    parts.append(f"\n[USER MESSAGES]\n{user_message}")

    if include_strict_output:
        parts.append(f"\n{CSV_STRICT_OUTPUT_PROMPT}")

    return "\n\n".join(parts)


def build_csv_system_prompt_with_context(metadata: Dict[str, Any]) -> str:
    """
    Build a combined system prompt embedding the dataset context inline.

    This is useful when the LLM API separates ``system`` and ``user`` roles
    and you want the full dataset description inside the system message.
    """
    shape = metadata.get("shape", {})
    columns = metadata.get("columns", [])
    dtypes = metadata.get("dtypes", {})
    sample_data = metadata.get("sample_data", [])

    sample_json = json.dumps(sample_data[:3], indent=2) if sample_data else "[]"

    return f"""\
{CSV_SYSTEM_PROMPT}

Dataset Context (pre-loaded as 'df'):
- Shape: {shape.get('rows', '?')} rows × {shape.get('columns', '?')} cols
- Columns: {', '.join(columns)}
- Dtypes: {json.dumps(dtypes)}
- Sample rows: {sample_json}

{CSV_STRICT_OUTPUT_PROMPT}"""