#!/usr/bin/env nu
# Schema Intelligence Tools for AuroraFrame MCP Server
#
# Provides AI-powered KCL schema assistance:
# - Generate KCL schemas from natural language
# - Validate and improve existing schemas
# - Migrate data between schema versions
# - Suggest best practices and optimizations

use mcp-server.nu call_openai_api
use mcp-server.nu validate_kcl_syntax

# Generate KCL schema from natural language description
export def generate_schema_tool [args: record, config: record, debug: bool] {
    let description = $args.description
    let examples = ($args.examples? | default [])

    debug_log $"Generating KCL schema from description: ($description)" $debug

    let system_prompt = (build_schema_generation_prompt)
    let user_prompt = (build_schema_user_prompt $description $examples)

    let messages = [
        { role: "system", content: $system_prompt }
        { role: "user", content: $user_prompt }
    ]

    let api_response = (call_openai_api $messages $config 0.3)

    if "error" in $api_response {
        return $api_response
    }

    # Validate the generated schema
    let validation = (validate_generated_schema $api_response.content)

    {
        content: [
            {
                type: "text"
                text: $"Generated KCL Schema:\n\n```kcl\n($api_response.content)\n```\n\nValidation Results:\n($validation)"
            }
        ]
    }
}

# Validate and suggest improvements for KCL schema
export def validate_schema_tool [args: record, config: record, debug: bool] {
    let schema = $args.schema
    let data = ($args.data? | default [])

    debug_log $"Validating KCL schema" $debug

    let system_prompt = $"You are an expert in KCL (KCL Configuration Language) and schema design. Analyze the provided KCL schema for:

1. **Syntax Correctness**: Check for proper KCL syntax and structure
2. **Type Safety**: Ensure proper type definitions and constraints
3. **Best Practices**: Follow KCL and schema design best practices
4. **Completeness**: Identify missing fields or validations
5. **Performance**: Suggest optimizations for better performance
6. **Maintainability**: Ensure schema is easy to understand and maintain

For each issue found, provide:
- **Issue type** (error, warning, suggestion)
- **Description** of the problem
- **Recommended fix** with example code
- **Rationale** explaining why the fix is important

If sample data is provided, validate it against the schema and identify any mismatches.

Return a detailed analysis with actionable recommendations."

    let user_prompt = $"KCL Schema to validate:
```kcl
($schema)
```

(if ($data | length) > 0 { $"\nSample data to validate:\n```json\n(($data | to json))\n```" } else { "" })

Please provide a comprehensive validation analysis."

    let messages = [
        { role: "system", content: $system_prompt }
        { role: "user", content: $user_prompt }
    ]

    let api_response = (call_openai_api $messages $config 0.2)

    if "error" in $api_response {
        return $api_response
    }

    # Add basic syntax validation
    let syntax_validation = (validate_kcl_syntax $schema)

    let combined_validation = if $syntax_validation.valid {
        $"✅ Basic syntax validation passed\n\nAI Analysis:\n($api_response.content)"
    } else {
        $"❌ Basic syntax issues found:\n($syntax_validation.issues | str join '\n- ')\n\nAI Analysis:\n($api_response.content)"
    }

    {
        content: [
            {
                type: "text"
                text: $"Schema Validation Results:\n\n($combined_validation)"
            }
        ]
    }
}

# Help migrate data between schema versions
export def migrate_schema_tool [args: record, config: record, debug: bool] {
    let old_schema = $args.old_schema
    let new_schema = $args.new_schema
    let data = ($args.data? | default [])

    debug_log "Analyzing schema migration" $debug

    let system_prompt = $"You are an expert in data migration and schema evolution. Analyze the differences between two KCL schema versions and provide:

1. **Change Analysis**: Identify all changes between schemas (added, removed, modified fields)
2. **Migration Strategy**: Provide step-by-step migration approach
3. **Data Transformation**: Show how to transform data from old to new format
4. **Risk Assessment**: Identify potential data loss or compatibility issues
5. **Rollback Plan**: Suggest how to rollback if needed
6. **Validation**: How to ensure migration was successful

If sample data is provided, show the actual transformation with examples.

Generate practical migration code and scripts where applicable."

    let user_prompt = $"Old Schema:
```kcl
($old_schema)
```

New Schema:
```kcl
($new_schema)
```

(if ($data | length) > 0 { $"\nSample data to migrate:\n```json\n(($data | to json))\n```" } else { "" })

Please provide a comprehensive migration plan."

    let messages = [
        { role: "system", content: $system_prompt }
        { role: "user", content: $user_prompt }
    ]

    let api_response = (call_openai_api $messages $config 0.3)

    if "error" in $api_response {
        return $api_response
    }

    {
        content: [
            {
                type: "text"
                text: $"Schema Migration Plan:\n\n($api_response.content)"
            }
        ]
    }
}

# Suggest schema improvements based on usage patterns
export def suggest_schema_improvements [schema: string, usage_data: list, config: record] {
    let system_prompt = $"You are a schema optimization expert. Analyze the provided KCL schema and usage data to suggest improvements for:

1. **Performance Optimization**: Reduce validation time and memory usage
2. **Type Safety**: Strengthen type constraints based on actual usage
3. **Maintainability**: Improve schema structure and documentation
4. **Extensibility**: Make schema easier to extend in the future
5. **Best Practices**: Apply KCL and schema design best practices
6. **Error Prevention**: Add constraints to prevent common errors

Provide specific, actionable recommendations with example code."

    let user_prompt = $"Current Schema:
```kcl
($schema)
```

Usage Data Analysis:
```json
($usage_data | to json)
```

Please suggest specific improvements."

    let messages = [
        { role: "system", content: $system_prompt }
        { role: "user", content: $user_prompt }
    ]

    let api_response = (call_openai_api $messages $config 0.4)

    if "error" in $api_response {
        return $api_response.error
    }

    $api_response.content
}

# Build schema generation system prompt
def build_schema_generation_prompt [] {
    $"You are an expert KCL (KCL Configuration Language) schema designer for AuroraFrame, a type-safe static site generator.

Your task is to generate well-structured KCL schemas based on natural language descriptions and optional example data.

KCL Schema Best Practices:
1. **Use descriptive names** for schemas and fields
2. **Apply appropriate type constraints** (str, int, bool, [str], etc.)
3. **Add default values** where sensible
4. **Use optional fields** (field?) when appropriate
5. **Include validation constraints** for ranges, patterns, etc.
6. **Structure nested schemas** for complex objects
7. **Add comments** for documentation

Common AuroraFrame Schema Patterns:
- **Content schemas**: title, author, date, tags, content
- **Site configuration**: navigation, metadata, features
- **Page schemas**: layout, template, seo data
- **Component schemas**: reusable UI components

Example KCL Schema:
```kcl
schema BlogPost:
    title: str
    slug: str
    author: str
    published_date: str
    modified_date?: str
    tags: [str] = []
    featured: bool = False
    content: str
    seo: SEOData

schema SEOData:
    description: str
    keywords: [str]
    og_image?: str
```

Generate clean, type-safe, and well-documented KCL schemas."
}

# Build user prompt for schema generation
def build_schema_user_prompt [description: string, examples: list] {
    let base_prompt = $"Generate a KCL schema based on this description:\n\n($description)"

    let examples_section = if ($examples | length) > 0 {
        $"\n\nExample data objects to consider:\n```json\n($examples | to json)\n```"
    } else {
        ""
    }

    $"($base_prompt)($examples_section)\n\nPlease generate a comprehensive KCL schema with appropriate types, constraints, and documentation."
}

# Validate generated schema
def validate_generated_schema [schema: string] {
    let syntax_validation = (validate_kcl_syntax $schema)

    if $syntax_validation.valid {
        "✅ Schema appears to be well-formed\n✅ Proper KCL syntax detected\n✅ Type annotations present"
    } else {
        $"Validation Issues:\n($syntax_validation.issues | each { |issue| $"⚠️  ($issue)" } | str join '\n')"
    }
}

# Extract schema metadata for analysis
export def extract_schema_metadata [schema: string] {
    mut metadata = {
        schemas: []
        fields: []
        types: []
        constraints: []
    }

    let lines = ($schema | lines)

    # Extract schema definitions
    for line in $lines {
        if ($line | str contains "schema ") and ($line | str contains ":") {
            let schema_name = ($line | str replace "schema " "" | str replace ":" "" | str trim)
            $metadata.schemas = ($metadata.schemas | append $schema_name)
        }

        # Extract field definitions
        if ($line | str trim | str contains ": ") and not ($line | str contains "schema") {
            let trimmed = ($line | str trim)
            let parts = ($trimmed | split column ": ")
            if ($parts | length) >= 2 {
                let field_name = ($parts | get column1 | str trim)
                let field_type = ($parts | get column2 | str trim)
                $metadata.fields = ($metadata.fields | append $field_name)
                $metadata.types = ($metadata.types | append $field_type)
            }
        }
    }

    $metadata
}

# Generate schema documentation
export def generate_schema_documentation [schema: string] {
    let metadata = (extract_schema_metadata $schema)

    $"# Schema Documentation

## Schemas Defined
($metadata.schemas | each { |s| $"- ($s)" } | str join '\n')

## Field Types Used
(($metadata.types | uniq) | each { |t| $"- ($t)" } | str join '\n')

## Total Fields
($metadata.fields | length) fields defined across all schemas

## Usage Example
```kcl
# Import and use the schema
import \"schema.k\"

# Create an instance
instance: SchemaName = {
    // Fill required fields
}
```"
}

# Analyze schema complexity
export def analyze_schema_complexity [schema: string] {
    let metadata = (extract_schema_metadata $schema)
    let total_schemas = ($metadata.schemas | length)
    let total_fields = ($metadata.fields | length)
    let unique_types = ($metadata.types | uniq | length)

    let complexity_score = ($total_fields * 1) + ($total_schemas * 2) + ($unique_types * 0.5) | math round

    {
        total_schemas: $total_schemas
        total_fields: $total_fields
        unique_types: $unique_types
        complexity_score: $complexity_score
        complexity_level: (if $complexity_score <= 10 { "Simple" } else if $complexity_score <= 25 { "Moderate" } else { "Complex" })
    }
}

# Suggest schema optimizations
export def suggest_schema_optimizations [schema: string] {
    let complexity = (analyze_schema_complexity $schema)
    let metadata = (extract_schema_metadata $schema)

    mut suggestions = []

    # Check for overly complex schemas
    if $complexity.complexity_level == "Complex" {
        $suggestions = ($suggestions | append "Consider breaking down large schemas into smaller, focused schemas")
    }

    # Check for missing optional fields
    let optional_fields = ($metadata.fields | where { |field| ($field | str contains "?") } | length)
    if $optional_fields == 0 and ($metadata.fields | length) > 3 {
        $suggestions = ($suggestions | append "Consider making some fields optional with the '?' syntax")
    }

    # Check for missing default values
    if not ($schema | str contains " = ") {
        $suggestions = ($suggestions | append "Consider adding default values for common fields")
    }

    # Check for documentation
    if not ($schema | str contains "#") {
        $suggestions = ($suggestions | append "Add documentation comments using # syntax")
    }

    $suggestions
}

# Generate schema from existing data
export def infer_schema_from_data [data: list, schema_name: string] {
    if ($data | length) == 0 {
        return "# No data provided to infer schema from"
    }

    let first_item = ($data | first)
    let all_keys = ($data | each { |item| $item | columns } | flatten | uniq)

    mut schema_fields = []

    for key in $all_keys {
        # Analyze field across all data items
        let values = ($data | each { |item|
            if $key in ($item | columns) {
                $item | get $key
            } else {
                null
            }
        })

        let non_null_values = ($values | where $it != null)
        let field_required = (($non_null_values | length) == ($data | length))

        # Infer type from values
        let sample_value = ($non_null_values | first)
        let inferred_type = match ($sample_value | describe) {
            "string" => "str"
            "int" => "int"
            "bool" => "bool"
            "list" => "[str]"  # Simplified assumption
            "record" => "record"  # Could be nested schema
            _ => "str"  # Default fallback
        }

        let field_def = if $field_required {
            $"    ($key): ($inferred_type)"
        } else {
            $"    ($key)?: ($inferred_type)"
        }

        $schema_fields = ($schema_fields | append $field_def)
    }

    $"schema ($schema_name):
($schema_fields | str join '\n')"
}

# Debug helper
def debug_log [message: string, debug: bool] {
    if $debug {
        print $"🐛 SCHEMA-INT: ($message)"
    }
}