"""
AI Prompt Templates for KQL Generation

Enhanced prompts with KQL knowledge, schema context, and few-shot learning.
Uses special tokens and structured output for better AI accuracy.

Author: Arjun Trivedi
Email: arjuntrivedi42@yahoo.com
"""

from typing import Any, Dict, List, Optional


# System prompt with KQL expertise
KQL_SYSTEM_PROMPT = """You are an expert in Kusto Query Language (KQL). You generate accurate, efficient KQL queries.

CRITICAL RULES:
1. **Join Conditions**: ONLY use 'and' in join conditions, NEVER 'or'
   - ✅ CORRECT: Table1 | join Table2 on Col1 and Col2
   - ❌ WRONG: Table1 | join Table2 on Col1 or Col2

2. **Column Validation**: ONLY use columns that exist in the provided schema
   - Always check the schema before using a column name
   - Use exact column names (case-sensitive)

3. **Reserved Words**: Bracket reserved words and special characters
   - Use ['column-name'] for columns with hyphens or spaces
   - Use ['table name'] for tables with spaces

4. **Operator Best Practices**:
   - Use 'project' to select specific columns (avoid 'project *')
   - Use 'where' for filtering
   - Use 'summarize' for aggregations
   - Use 'extend' to add calculated columns
   - Use 'take' or 'limit' to limit results

5. **Data Types**: Use proper type conversions
   - toint(), tolong(), toreal() for numbers
   - tostring() for strings
   - todatetime() for dates
   - Handle nulls with isnull(), isnotnull(), iff()

OUTPUT FORMAT:
Return ONLY the KQL query, nothing else. No explanations, no markdown, just the query."""


# Specialized Mermaid Visualization Prompt
MERMAID_VISUALIZATION_PROMPT = """
You are an expert in Data Visualization using Mermaid.js.
Your goal is to create STUNNING, MODERN, and HIGH-CONTRAST diagrams using a CYBERPUNK/NEON aesthetic.

🎨 CYBERPUNK COLOR PALETTE:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Background & Fills:
  - Primary Dark:   #0a0e27 (deep space blue)
  - Secondary Dark: #1a1a2e (midnight blue)
  - Medium Dark:    #16213e (dark slate)
  - Medium:         #1a1a40 (navy blue)
  - Accent Dark:    #0f3460 (deep ocean)

Neon Accents & Strokes:
  - Cyan Neon:      #00d9ff (electric cyan)
  - Bright Cyan:    #00ffff (pure cyan)
  - Hot Pink:       #ff0080 (magenta neon)
  - Purple Neon:    #9d4edd (violet)
  - Light Purple:   #c77dff (lavender neon)
  - Gold Neon:      #ffaa00 (amber gold)
  - Orange Neon:    #ff6600 (fire orange)
  - Success Green:  #00ff88 (neon green)
  - Error Red:      #ff3366 (neon red)

Typography:
  - Font Family: 'Inter', 'Segoe UI', 'Roboto', sans-serif
  - Base Size: 16px-18px
  - Labels: 14px-16px

⚠️ STRICT COLOR COMPLIANCE:
Use ONLY the colors listed above. Do NOT use:
  - #00ff88 (neon green) - NOT in palette
  - #ff3366 (neon red) - NOT in palette
  - Any custom colors not explicitly listed
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

📋 CRITICAL INSTRUCTIONS:
1. **ALWAYS** use `%%{init: {...}}%%` directive for theme configuration (MANDATORY)
2. **COLOR PALETTE**: Use ONLY colors from the defined cyberpunk palette above
   - Fill Colors: Choose from #0a0e27 (darkest), #1a1a2e (dark), #16213e (medium-dark), #1a1a40 (medium), #0f3460 (medium-light)
   - Stroke Colors: Choose from #00d9ff (cyan-blue), #ff6600 (orange), #00ffff (cyan), #ff0080 (pink), #9d4edd (purple), #c77dff (light-purple), #ffaa00 (gold)
   - Example: `style NodeName fill:#0a0e27,stroke:#00d9ff,stroke-width:3px,color:#00d9ff`
3. **WATERMARK**: ALWAYS add report attribution in markdown after the mermaid diagram block:
   - Format: "**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)"
   - Place this text OUTSIDE and AFTER the closing ```
4. **High Contrast**: Ensure text is readable against dark backgrounds (use bright neon colors)
5. **Consistent Styling**: Apply node-specific styles for visual hierarchy
6. **Proper Spacing**: Use adequate node spacing (40-60px) for clarity

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

📊 DIAGRAM TYPE TEMPLATES:

═══════════════════════════════════════════════════════
1. FLOWCHART / PROCESS DIAGRAM (graph TB/LR)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#1a1a2e',
  'primaryTextColor':'#00d9ff',
  'primaryBorderColor':'#00d9ff',
  'secondaryColor':'#16213e',
  'secondaryTextColor':'#c77dff',
  'secondaryBorderColor':'#c77dff',
  'tertiaryColor':'#0f3460',
  'tertiaryTextColor':'#ffaa00',
  'tertiaryBorderColor':'#ffaa00',
  'lineColor':'#00d9ff',
  'textColor':'#ffffff',
  'mainBkg':'#0a0e27',
  'nodeBorder':'#00d9ff',
  'clusterBkg':'#16213e',
  'clusterBorder':'#9d4edd',
  'titleColor':'#00ffff',
  'edgeLabelBackground':'#1a1a2e',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'16px',
  'flowchart':{'nodeSpacing':50, 'rankSpacing':60, 'curve':'basis', 'padding':20}
}}}%%
graph TB
  Start([🚀 Start Process]) --> Validate[🔍 Validate Input]
  Validate --> Decision{✅ Valid?}
  Decision -->|Yes| Process[⚙️ Process Data]
  Decision -->|No| Error([❌ Error State])
  Process --> Success([✅ Success])

  %% Node Styles (using only palette colors)
  style Start fill:#1a1a2e,stroke:#00d9ff,stroke-width:3px,color:#00ffff
  style Validate fill:#16213e,stroke:#c77dff,stroke-width:2px,color:#c77dff
  style Decision fill:#1a1a40,stroke:#ffaa00,stroke-width:3px,color:#ffaa00
  style Process fill:#16213e,stroke:#9d4edd,stroke-width:2px,color:#9d4edd
  style Success fill:#0f3460,stroke:#00d9ff,stroke-width:3px,color:#00d9ff
  style Error fill:#1a1a2e,stroke:#ff0080,stroke-width:3px,color:#ff0080
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
2. SEQUENCE DIAGRAM (Interactions/Flow)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'actorBkg':'#1a1a2e',
  'actorBorder':'#00d9ff',
  'actorTextColor':'#00ffff',
  'actorLineColor':'#00d9ff',
  'signalColor':'#c77dff',
  'signalTextColor':'#ffffff',
  'labelBoxBkgColor':'#16213e',
  'labelBoxBorderColor':'#9d4edd',
  'labelTextColor':'#ffffff',
  'loopTextColor':'#ffaa00',
  'noteBkgColor':'#0f3460',
  'noteBorderColor':'#ff6600',
  'noteTextColor':'#ffffff',
  'activationBkgColor':'#9d4edd',
  'activationBorderColor':'#c77dff',
  'sequenceNumberColor':'#ffaa00',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'16px'
}}}%%
sequenceDiagram
  autonumber
  participant U as 👤 User
  participant S as 🖥️ System
  participant D as 💾 Database

  U->>+S: 📤 Send Request
  S->>+D: 🔍 Query Data
  D-->>-S: 📊 Return Results
  S-->>-U: ✅ Response
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
3. ER DIAGRAM (Database Schema)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#1a1a2e',
  'primaryTextColor':'#00d9ff',
  'primaryBorderColor':'#00d9ff',
  'lineColor':'#c77dff',
  'secondaryColor':'#16213e',
  'tertiaryColor':'#0f3460',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'16px',
  'entityBkg':'#1a1a2e',
  'entityBorder':'#00d9ff',
  'entityTextColor':'#00ffff',
  'relationshipLabelColor':'#ffaa00',
  'relationshipLineColor':'#c77dff'
}}}%%
erDiagram
  USERS ||--o{ ORDERS : "places"
  USERS {
    string userId PK
    string username
    string email
    datetime createdAt
  }
  ORDERS {
    string orderId PK
    string userId FK
    decimal amount
    datetime orderDate
  }
  ORDERS ||--|{ ORDER_ITEMS : "contains"
  ORDER_ITEMS {
    string itemId PK
    string orderId FK
    string productId FK
    int quantity
  }
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
4. STATE DIAGRAM (Workflow States)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#1a1a2e',
  'primaryTextColor':'#00d9ff',
  'primaryBorderColor':'#00d9ff',
  'lineColor':'#c77dff',
  'secondaryColor':'#16213e',
  'tertiaryColor':'#0f3460',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'16px'
}}}%%
stateDiagram-v2
  [*] --> Idle: 🚀 Initialize
  Idle --> Processing: ▶️ Start
  Processing --> Success: ✅ Complete
  Processing --> Failed: ❌ Error
  Failed --> Retry: 🔄 Retry
  Retry --> Processing
  Success --> [*]
  Failed --> [*]: 🛑 Abort
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
5. GANTT CHART (Timeline/Schedule)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#1a1a2e',
  'primaryTextColor':'#00d9ff',
  'primaryBorderColor':'#00d9ff',
  'lineColor':'#c77dff',
  'secondaryColor':'#16213e',
  'tertiaryColor':'#0f3460',
  'gridColor':'#00d9ff',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'14px'
}}}%%
gantt
  title 📊 Project Timeline
  dateFormat YYYY-MM-DD
  section Phase 1
  Task 1 :a1, 2024-01-01, 30d
  Task 2 :a2, after a1, 20d
  section Phase 2
  Task 3 :a3, after a2, 25d
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
6. PIE CHART (Distribution)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#00d9ff',
  'primaryTextColor':'#ffffff',
  'pie1':'#00d9ff',
  'pie2':'#ff0080',
  'pie3':'#9d4edd',
  'pie4':'#ffaa00',
  'pie5':'#00ff88',
  'pie6':'#ff6600',
  'pieTitleTextColor':'#00ffff',
  'pieSectionTextColor':'#ffffff',
  'pieLegendTextColor':'#ffffff',
  'pieStrokeColor':'#0a0e27',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'16px'
}}}%%
pie title 📊 Data Distribution
  "Category A" : 45
  "Category B" : 30
  "Category C" : 15
  "Category D" : 10
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

═══════════════════════════════════════════════════════
7. CLASS DIAGRAM (Object Structure)
═══════════════════════════════════════════════════════
```mermaid
%%{init: {'theme':'dark', 'themeVariables': {
  'primaryColor':'#1a1a2e',
  'primaryTextColor':'#00d9ff',
  'primaryBorderColor':'#00d9ff',
  'lineColor':'#c77dff',
  'secondaryColor':'#16213e',
  'tertiaryColor':'#0f3460',
  'fontFamily':'Inter, Segoe UI, sans-serif',
  'fontSize':'14px',
  'classText':'#00ffff'
}}}%%
classDiagram
  class KQLQuery {
    +string query
    +string database
    +execute()
    +validate()
  }
  class Result {
    +array data
    +int rowCount
    +format()
  }
  KQLQuery --> Result : returns
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

🎯 SELECTION GUIDE:
- **Process/Workflow**: Use Flowchart (graph TB/LR)
- **Interactions/API Calls**: Use Sequence Diagram
- **Data Relationships**: Use ER Diagram
- **Status/Lifecycle**: Use State Diagram
- **Timeline/Schedule**: Use Gantt Chart
- **Percentages/Parts**: Use Pie Chart
- **Code/Classes**: Use Class Diagram

⚡ BEST PRACTICES:
1. **MANDATORY**: Add markdown attribution AFTER closing ``` : "**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)"
2. **COLOR RESTRICTION**: Use ONLY the defined cyberpunk palette colors (no custom colors)
   - Fills: #0a0e27, #1a1a2e, #16213e, #1a1a40, #0f3460
   - Strokes: #00d9ff, #ff6600, #00ffff, #ff0080, #9d4edd, #c77dff, #ffaa00
3. Always include emojis for visual appeal and clarity
4. Keep node labels concise (max 3-4 words)
5. Test that text colors have sufficient contrast (AA standard minimum)
6. Use line thickness to indicate importance (3px=critical, 2px=normal, 1px=minor)
7. Add meaningful arrow labels when needed
8. Group related items in subgraphs when applicable

🚫 COMMON MISTAKES TO AVOID:
- ❌ Don't use colors outside the defined cyberpunk palette (#00ff88, #ff3366, etc.)
- ❌ Don't put watermarks INSIDE the mermaid diagram
- ❌ Don't forget the markdown attribution after the closing ```
- ❌ Don't forget the init directive
- ❌ Don't use low-contrast color combinations
- ❌ Don't overcrowd diagrams (max 10-15 nodes per graph)

📝 DIAGRAM GENERATION WORKFLOW:
1. Analyze the data/query results structure
2. Choose the MOST APPROPRIATE diagram type from templates above
3. Apply the `%%{init: {...}}%%` directive with cyberpunk theme
4. Build the diagram using ONLY palette colors for fills and strokes
5. Add clear labels with emojis for visual interest
6. Close the mermaid code block with ```
7. **CRITICAL**: Add markdown attribution line: "**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)"
8. Verify color palette compliance and readability

✨ OUTPUT FORMAT EXAMPLE:
```mermaid
%%{init: {...}}%%
[your diagram here]
```

**Report Generated by MCP-KQL-Server** | [⭐ Star this repo on GitHub](https://github.com/4R9UN/mcp-kql-server)

Remember: The goal is STUNNING, PREMIUM, HIGH-CONTRAST visualizations using the strict cyberpunk color palette! 🎨✨
"""


# Few-shot learning examples
FEW_SHOT_EXAMPLES = [
    {
        "description": "Show recent failed login attempts",
        "schema": {
            "table": "SigninLogs",
            "columns": {
                "TimeGenerated": {"data_type": "datetime"},
                "UserPrincipalName": {"data_type": "string"},
                "ResultType": {"data_type": "string"},
                "ResultDescription": {"data_type": "string"},
                "IPAddress": {"data_type": "string"},
                "Location": {"data_type": "string"}
            }
        },
        "query": "SigninLogs | where ResultType != '0' | where TimeGenerated > ago(1h) | project TimeGenerated, UserPrincipalName, ResultDescription, IPAddress, Location | take 100"
    },
    {
        "description": "Count events by severity in the last 24 hours",
        "schema": {
            "table": "SecurityEvent",
            "columns": {
                "TimeGenerated": {"data_type": "datetime"},
                "EventID": {"data_type": "int"},
                "Level": {"data_type": "string"},
                "Computer": {"data_type": "string"},
                "Account": {"data_type": "string"}
            }
        },
        "query": "SecurityEvent | where TimeGenerated > ago(24h) | summarize Count=count() by Level | order by Count desc"
    },
    {
        "description": "Find top 10 users by activity",
        "schema": {
            "table": "AuditLogs",
            "columns": {
                "TimeGenerated": {"data_type": "datetime"},
                "OperationName": {"data_type": "string"},
                "InitiatedBy": {"data_type": "string"},
                "TargetResources": {"data_type": "dynamic"},
                "Result": {"data_type": "string"}
            }
        },
        "query": "AuditLogs | summarize ActivityCount=count() by InitiatedBy | top 10 by ActivityCount desc"
    },
    {
        "description": "Join two tables to correlate data",
        "schema": {
            "table1": "Alerts",
            "columns1": {
                "AlertId": {"data_type": "string"},
                "Severity": {"data_type": "string"},
                "DeviceId": {"data_type": "string"},
                "TimeGenerated": {"data_type": "datetime"}
            },
            "table2": "Devices",
            "columns2": {
                "DeviceId": {"data_type": "string"},
                "DeviceName": {"data_type": "string"},
                "OSPlatform": {"data_type": "string"}
            }
        },
        "query": "Alerts | join kind=inner Devices on DeviceId | project TimeGenerated, AlertId, Severity, DeviceName, OSPlatform | take 50"
    }
]


def build_generation_prompt(
    nl_query: str,
    schema: Dict[str, Any],
    table_name: Optional[str] = None,
    include_examples: bool = True,
    include_visualization: bool = False,
    similar_queries: Optional[List[Dict[str, Any]]] = None
) -> str:
    """
    Build optimized prompt for KQL generation with schema context.

    Args:
        nl_query: Natural language query from user
        schema: Table schema with columns and types
        table_name: Target table name
        include_examples: Whether to include few-shot examples

    Returns:
        Formatted prompt string
    """
    prompt_parts = []

    # Add schema context
    if schema and schema.get("columns"):
        prompt_parts.append("SCHEMA CONTEXT:")
        prompt_parts.append(f"Table: {table_name or schema.get('table', 'Unknown')}")

        # Extract column names for critical rules section
        column_names = list(schema["columns"].keys())

        prompt_parts.append("Columns:")

        for col_name, col_info in list(schema["columns"].items())[:20]:  # Limit to 20 columns
            data_type = col_info.get("data_type", "unknown")
            description = col_info.get("description", "")

            col_line = f"  - {col_name} ({data_type})"

            # Add sample values if available
            sample_values = col_info.get("sample_values", [])
            if sample_values:
                # Format samples nicely
                samples_str = ", ".join(str(v) for v in sample_values[:2])
                col_line += f" [Samples: {samples_str}]"

            if description:
                col_line += f": {description[:50]}"
            prompt_parts.append(col_line)

        if len(schema["columns"]) > 20:
            prompt_parts.append(f"  ... and {len(schema['columns']) - 20} more columns")

        prompt_parts.append("")

        # Add CRITICAL RULES section
        prompt_parts.append("⚠️ CRITICAL RULES - COLUMN NAMES:")
        prompt_parts.append("1. Use ONLY these EXACT column names (case-sensitive):")
        # Show first 15 column names explicitly
        for i, col_name in enumerate(column_names[:15], 1):
            prompt_parts.append(f"   {i}. {col_name}")
        if len(column_names) > 15:
            prompt_parts.append(f"   ... and {len(column_names) - 15} more columns listed above")
        prompt_parts.append("2. DO NOT modify, abbreviate, or invent column names")
        prompt_parts.append("3. Column names are CASE-SENSITIVE - use exact capitalization")
        prompt_parts.append("4. If you need a column not in this list, the query cannot be fulfilled")
        prompt_parts.append("")

    # Add few-shot examples (optional)
    if include_examples:
        prompt_parts.append("EXAMPLES:")

        # Add dynamic similar queries if available
        if similar_queries:
            prompt_parts.append("--- RELEVANT PAST QUERIES (High Confidence) ---")
            for i, q in enumerate(similar_queries, 1):
                desc = q.get('description', 'Similar query')
                query = q.get('query', '').strip()
                prompt_parts.append(f"// {desc}")
                prompt_parts.append(query)
                prompt_parts.append("")
            prompt_parts.append("--- GENERAL EXAMPLES ---")

        for i, example in enumerate(FEW_SHOT_EXAMPLES[:2], 1):  # Include 2 examples
            prompt_parts.append(f"\nExample {i}:")
            prompt_parts.append(f"User: {example['description']}")
            prompt_parts.append(f"KQL: {example['query']}")
        prompt_parts.append("")

    # Add user query
    prompt_parts.append("USER REQUEST:")
    prompt_parts.append(nl_query)
    prompt_parts.append("")
    prompt_parts.append("Generate a KQL query that:")
    prompt_parts.append("1. Uses ONLY columns from the schema above")
    prompt_parts.append("2. Follows all KQL syntax rules")
    prompt_parts.append("3. Is efficient, optimize and accurate")
    prompt_parts.append("")
    prompt_parts.append("KQL Query:")

    if include_visualization:
        prompt_parts.append("\n" + MERMAID_VISUALIZATION_PROMPT)

    return "\n".join(prompt_parts)


def build_schema_description_prompt(
    table_name: str,
    columns: Dict[str, Any]
) -> str:
    """
    Build prompt for LLM to generate table description.

    Args:
        table_name: Name of the table
        columns: Dictionary of columns with metadata

    Returns:
        Formatted prompt for description generation
    """
    col_names = ", ".join(list(columns.keys())[:15])
    if len(columns) > 15:
        col_names += f", and {len(columns) - 15} more"

    prompt = f"""Generate a concise, natural language description for a database table.

Table Name: {table_name}
Columns: {col_names}

Description should:
1. Be 1-2 sentences
2. Explain what data this table likely contains
3. Mention key columns if obvious (e.g., timestamps, IDs, names)
4. Be helpful for accurate query generation

Description:"""

    return prompt


def build_error_feedback_prompt(
    original_query: str,
    error_message: str,
    schema: Dict[str, Any]
) -> str:
    """
    Build prompt for LLM to fix a failed query.

    Args:
        original_query: The query that failed
        error_message: Error message from Kusto
        schema: Table schema

    Returns:
        Formatted prompt for query correction
    """
    prompt_parts = []

    prompt_parts.append("QUERY ERROR - PLEASE FIX:")
    prompt_parts.append("")
    prompt_parts.append("Original Query:")
    prompt_parts.append(original_query)
    prompt_parts.append("")
    prompt_parts.append("Error Message:")
    prompt_parts.append(error_message)
    prompt_parts.append("")

    if schema and schema.get("columns"):
        prompt_parts.append("Available Columns:")
        for col_name in list(schema["columns"].keys())[:20]:
            prompt_parts.append(f"  - {col_name}")
        prompt_parts.append("")

    prompt_parts.append("Generate a CORRECTED KQL query that:")
    prompt_parts.append("1. Fixes the error")
    prompt_parts.append("2. Uses only valid columns from the schema")
    prompt_parts.append("3. Maintains the original intent")
    prompt_parts.append("")
    prompt_parts.append("Corrected KQL Query:")

    return "\n".join(prompt_parts)


def extract_kql_from_response(response: str) -> str:
    """
    Extract KQL query from LLM response.

    Handles various response formats:
    - Plain query
    - Query in code blocks
    - Query with explanations

    Args:
        response: LLM response text

    Returns:
        Extracted KQL query
    """
    # Remove markdown code blocks
    if "```" in response:
        # Extract content between ```kql or ``` blocks
        import re
        pattern = r'```(?:kql)?\s*\n(.*?)\n```'
        matches = re.findall(pattern, response, re.DOTALL)
        if matches:
            response = matches[0]

    # Remove common prefixes
    prefixes = [
        "KQL Query:",
        "Query:",
        "Here's the query:",
        "Here is the query:",
        "The query is:",
    ]

    for prefix in prefixes:
        if response.strip().startswith(prefix):
            response = response.strip()[len(prefix):].strip()

    # Take first line if multi-line with explanations
    lines = response.strip().split('\n')
    if len(lines) > 1:
        # Check if first line looks like a query
        first_line = lines[0].strip()
        if any(kw in first_line.lower() for kw in ['|', 'where', 'project', 'summarize', 'take']):
            # Multi-line query, join all lines that look like query parts
            query_lines = []
            for line in lines:
                line = line.strip()
                if line and not line.startswith('#') and not line.startswith('//'):
                    query_lines.append(line)
                elif query_lines:  # Stop at first non-query line after query started
                    break
            response = ' '.join(query_lines)

    return response.strip()


# Special tokens for structured output
SPECIAL_TOKENS = {
    "query_start": "<KQL>",
    "query_end": "</KQL>",
    "error_start": "<ERROR>",
    "error_end": "</ERROR>",
    "suggestion_start": "<SUGGESTION>",
    "suggestion_end": "</SUGGESTION>"
}


def build_structured_prompt(
    nl_query: str,
    schema: Dict[str, Any],
    use_special_tokens: bool = False
) -> str:
    """
    Build prompt with optional special tokens for structured output.

    Args:
        nl_query: Natural language query
        schema: Table schema
        use_special_tokens: Whether to use special tokens

    Returns:
        Formatted prompt
    """
    base_prompt = build_generation_prompt(nl_query, schema, include_examples=True)

    if use_special_tokens:
        base_prompt += f"\n\nWrap your KQL query in {SPECIAL_TOKENS['query_start']} and {SPECIAL_TOKENS['query_end']} tags."

    return base_prompt
