package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/CalmoAI/mcp-server-signoz/pkg/mcp"
)

// SchemaValidator provides comprehensive validation for MCP tool schemas
type SchemaValidator struct {
	errors   []string
	warnings []string
	stats    ValidationStats
}

type ValidationStats struct {
	TotalTools        int
	ToolsWithArrays   int
	ArrayPropsChecked int
	ArrayPropsValid   int
	RequiredFields    int
	OptionalFields    int
}

func NewSchemaValidator() *SchemaValidator {
	return &SchemaValidator{
		errors:   make([]string, 0),
		warnings: make([]string, 0),
		stats:    ValidationStats{},
	}
}

func (sv *SchemaValidator) AddError(msg string) {
	sv.errors = append(sv.errors, msg)
}

func (sv *SchemaValidator) AddWarning(msg string) {
	sv.warnings = append(sv.warnings, msg)
}

func main() {
	// Parse command line arguments
	var mode string
	var export bool

	for _, arg := range os.Args[1:] {
		switch arg {
		case "--gemini-only":
			mode = "gemini"
		case "--export":
			export = true
		case "--help", "-h":
			printUsage()
			os.Exit(0)
		}
	}

	if mode == "gemini" {
		runGeminiValidation(export)
	} else {
		runComprehensiveValidation(export)
	}
}

func printUsage() {
	fmt.Println("🔍 SigNoz MCP Schema Validator")
	fmt.Println("=============================")
	fmt.Println("")
	fmt.Println("Usage:")
	fmt.Println("  go run scripts/comprehensive_schema_validator.go [options]")
	fmt.Println("")
	fmt.Println("Options:")
	fmt.Println("  --gemini-only    Run only Gemini-specific validation (faster)")
	fmt.Println("  --export         Export schemas and validation report to JSON files")
	fmt.Println("  --help, -h       Show this help message")
	fmt.Println("")
	fmt.Println("Examples:")
	fmt.Println("  go run scripts/comprehensive_schema_validator.go")
	fmt.Println("  go run scripts/comprehensive_schema_validator.go --gemini-only")
	fmt.Println("  go run scripts/comprehensive_schema_validator.go --export")
}

func runGeminiValidation(export bool) {
	fmt.Println("🔍 Focused Gemini API Schema Validation")
	fmt.Println("=======================================")

	// Test only the critical tool that was causing Gemini issues
	builderTool := getAllSigNozTools()["signoz_execute_builder_query"]

	validator := NewSchemaValidator()
	validator.stats.TotalTools = 1

	fmt.Printf("🔧 Testing: %s (Gemini critical validation)\n", builderTool.Name)
	validator.validateGeminiCritical(builderTool)

	// Test JSON serialization
	fmt.Println("\n🔄 Testing JSON serialization (what gets sent to Gemini)...")
	data, err := json.MarshalIndent(builderTool, "", "  ")
	if err != nil {
		validator.AddError(fmt.Sprintf("JSON marshaling failed: %v", err))
	} else {
		fmt.Printf("   ✅ JSON serialization successful (%d bytes)\n", len(data))

		jsonStr := string(data)
		if !strings.Contains(jsonStr, `"items":`) {
			validator.AddError("items definitions not found in JSON output")
		} else {
			fmt.Printf("   ✅ Found 'items' definitions in JSON\n")
		}
	}

	// Generate focused report
	fmt.Println("\n📊 GEMINI VALIDATION RESULTS")
	fmt.Println("============================")

	if len(validator.errors) == 0 {
		fmt.Println("🎉 SUCCESS: Gemini validation passed!")
		fmt.Println("✅ Array properties have proper 'items' definitions")
		fmt.Println("✅ Schema should be accepted by Gemini API")

		if export && len(data) > 0 {
			os.WriteFile("gemini_test_schema.json", data, 0644)
			fmt.Println("✅ Test schema exported to: gemini_test_schema.json")
		}
	} else {
		fmt.Printf("❌ FAILED: Found %d error(s):\n", len(validator.errors))
		for i, err := range validator.errors {
			fmt.Printf("   %d. %s\n", i+1, err)
		}
		os.Exit(1)
	}
}

func runComprehensiveValidation(export bool) {
	fmt.Println("🔍 Comprehensive SigNoz MCP Schema Validation")
	fmt.Println("=============================================")

	// Define all expected SigNoz tools with their expected schemas
	tools := getAllSigNozTools()

	validator := NewSchemaValidator()
	validator.stats.TotalTools = len(tools)

	fmt.Printf("\n📋 Validating %d SigNoz MCP Tools:\n", len(tools))
	fmt.Println("===========================================")

	for toolName, tool := range tools {
		fmt.Printf("\n🔧 %s\n", toolName)
		validator.validateTool(toolName, tool)
	}

	// Generate comprehensive report
	validator.generateReport()

	// Export schemas if requested
	if export {
		validator.exportSchemas(tools)
	}

	// Exit with error code if validation failed
	if len(validator.errors) > 0 {
		os.Exit(1)
	}
}

func (sv *SchemaValidator) validateTool(toolName string, tool mcp.Tool) {
	// 1. Basic tool structure validation
	sv.validateBasicStructure(toolName, tool)

	// 2. Input schema validation
	sv.validateInputSchema(toolName, tool.InputSchema)

	// 3. Properties validation
	sv.validateProperties(toolName, tool.InputSchema.Properties)

	// 4. Required fields validation
	sv.validateRequiredFields(toolName, tool.InputSchema)

	// 5. SigNoz-specific validations
	sv.validateSigNozSpecific(toolName, tool)
}

func (sv *SchemaValidator) validateGeminiCritical(tool mcp.Tool) {
	// Focus only on the critical Gemini validation issues

	// Check the specific fields that were causing Gemini errors
	filtersProperty, exists := tool.InputSchema.Properties["filters"]
	if !exists {
		sv.AddError("filters property not found")
	} else {
		fmt.Printf("   📋 filters property: type=%s\n", filtersProperty.Type)

		if filtersProperty.Type != "array" {
			sv.AddError("filters should be array type")
		} else {
			if filtersProperty.Items == nil {
				sv.AddError("filters array missing Items definition (CRITICAL: Gemini rejects this!)")
			} else {
				fmt.Printf("   ✅ filters.items: type=%s\n", filtersProperty.Items.Type)
				if filtersProperty.Items.Type != "object" {
					sv.AddWarning("filters items should be object type")
				}
			}
		}
	}

	groupByProperty, exists := tool.InputSchema.Properties["group_by"]
	if !exists {
		sv.AddError("group_by property not found")
	} else {
		fmt.Printf("   📋 group_by property: type=%s\n", groupByProperty.Type)

		if groupByProperty.Type != "array" {
			sv.AddError("group_by should be array type")
		} else {
			if groupByProperty.Items == nil {
				sv.AddError("group_by array missing Items definition (CRITICAL: Gemini rejects this!)")
			} else {
				fmt.Printf("   ✅ group_by.items: type=%s\n", groupByProperty.Items.Type)
				if groupByProperty.Items.Type != "string" {
					sv.AddWarning("group_by items should be string type")
				}
			}
		}
	}

	// Update stats for the focused validation
	sv.stats.ArrayPropsChecked = 2
	if filtersProperty.Items != nil {
		sv.stats.ArrayPropsValid++
	}
	if groupByProperty.Items != nil {
		sv.stats.ArrayPropsValid++
	}

	fmt.Printf("   📊 Critical array validation: %d/%d valid\n", sv.stats.ArrayPropsValid, sv.stats.ArrayPropsChecked)
}

func (sv *SchemaValidator) validateBasicStructure(toolName string, tool mcp.Tool) {
	if tool.Name == "" {
		sv.AddError(fmt.Sprintf("%s: tool name is empty", toolName))
	}

	if tool.Name != toolName {
		sv.AddError(fmt.Sprintf("%s: tool name mismatch (expected %s, got %s)", toolName, toolName, tool.Name))
	}

	if tool.Description == "" {
		sv.AddWarning(fmt.Sprintf("%s: missing description", toolName))
	}

	if len(tool.Description) < 10 {
		sv.AddWarning(fmt.Sprintf("%s: description too short (%d chars)", toolName, len(tool.Description)))
	}

	fmt.Printf("   ✅ Basic structure: name=%s, description=%d chars\n", tool.Name, len(tool.Description))
}

func (sv *SchemaValidator) validateInputSchema(toolName string, schema mcp.ToolInputSchema) {
	if schema.Type != "object" {
		sv.AddError(fmt.Sprintf("%s: input schema type must be 'object', got '%s'", toolName, schema.Type))
		return
	}

	if schema.Properties == nil {
		sv.AddError(fmt.Sprintf("%s: input schema missing properties", toolName))
		return
	}

	fmt.Printf("   ✅ Input schema: type=%s, properties=%d\n", schema.Type, len(schema.Properties))
}

func (sv *SchemaValidator) validateProperties(toolName string, properties map[string]mcp.Property) {
	arrayCount := 0
	validArrayCount := 0

	for propName, prop := range properties {
		// Basic property validation
		if prop.Type == "" {
			sv.AddError(fmt.Sprintf("%s.%s: missing type field", toolName, propName))
			continue
		}

		// Array property validation (the critical Gemini issue)
		if prop.Type == "array" {
			arrayCount++
			sv.stats.ArrayPropsChecked++

			fmt.Printf("   🔍 Array property: %s\n", propName)

			if prop.Items == nil {
				sv.AddError(fmt.Sprintf("%s.%s: array property missing 'items' definition (CRITICAL: Gemini will reject this!)", toolName, propName))
			} else {
				validArrayCount++
				sv.stats.ArrayPropsValid++

				if prop.Items.Type == "" {
					sv.AddError(fmt.Sprintf("%s.%s: array items missing type", toolName, propName))
				} else {
					fmt.Printf("      ✅ Items type: %s\n", prop.Items.Type)
				}

				// Validate expected SigNoz array types
				sv.validateSigNozArrayTypes(toolName, propName, prop)
			}
		}

		// Date-time format validation
		if prop.Format == "date-time" && prop.Type != "string" {
			sv.AddError(fmt.Sprintf("%s.%s: date-time format requires string type, got %s", toolName, propName, prop.Type))
		}

		// Enum validation
		if len(prop.Enum) > 0 && prop.Type != "string" {
			sv.AddWarning(fmt.Sprintf("%s.%s: enum values with non-string type may cause issues", toolName, propName))
		}

		// Number constraints validation
		if prop.Type == "number" {
			if prop.Minimum != nil && prop.Maximum != nil && *prop.Minimum > *prop.Maximum {
				sv.AddError(fmt.Sprintf("%s.%s: minimum (%f) greater than maximum (%f)", toolName, propName, *prop.Minimum, *prop.Maximum))
			}
		}
	}

	if arrayCount > 0 {
		sv.stats.ToolsWithArrays++
		fmt.Printf("   📊 Array properties: %d/%d valid\n", validArrayCount, arrayCount)
	}
}

func (sv *SchemaValidator) validateSigNozArrayTypes(toolName, propName string, prop mcp.Property) {
	switch propName {
	case "filters":
		if prop.Items.Type != "object" {
			sv.AddWarning(fmt.Sprintf("%s.%s: expected items type 'object' for filters, got '%s'", toolName, propName, prop.Items.Type))
		}
	case "group_by":
		if prop.Items.Type != "string" {
			sv.AddWarning(fmt.Sprintf("%s.%s: expected items type 'string' for group_by, got '%s'", toolName, propName, prop.Items.Type))
		}
	}
}

func (sv *SchemaValidator) validateRequiredFields(toolName string, schema mcp.ToolInputSchema) {
	requiredCount := len(schema.Required)
	sv.stats.RequiredFields += requiredCount
	sv.stats.OptionalFields += len(schema.Properties) - requiredCount

	// Check that all required fields exist in properties
	for _, requiredField := range schema.Required {
		if _, exists := schema.Properties[requiredField]; !exists {
			sv.AddError(fmt.Sprintf("%s: required field '%s' not found in properties", toolName, requiredField))
		}
	}

	fmt.Printf("   📝 Required fields: %d/%d (%v)\n", requiredCount, len(schema.Properties), schema.Required)
}

func (sv *SchemaValidator) validateSigNozSpecific(toolName string, tool mcp.Tool) {
	// Validate SigNoz-specific naming conventions
	if !strings.HasPrefix(toolName, "signoz_") {
		sv.AddWarning(fmt.Sprintf("%s: tool name doesn't follow SigNoz convention (signoz_*)", toolName))
	}

	// Validate common SigNoz parameters
	commonTimeParams := []string{"start_time", "end_time"}
	for _, param := range commonTimeParams {
		if prop, exists := tool.InputSchema.Properties[param]; exists {
			if prop.Type != "string" || prop.Format != "date-time" {
				sv.AddError(fmt.Sprintf("%s.%s: time parameters should be string with date-time format", toolName, param))
			}
		}
	}

	// Check for SigNoz service_name parameter consistency
	if prop, exists := tool.InputSchema.Properties["service_name"]; exists {
		if prop.Type != "string" {
			sv.AddError(fmt.Sprintf("%s.service_name: should be string type, got %s", toolName, prop.Type))
		}
		if prop.Description == "" {
			sv.AddWarning(fmt.Sprintf("%s.service_name: missing description", toolName))
		}
	}
}

func (sv *SchemaValidator) generateReport() {
	fmt.Println("\n📊 COMPREHENSIVE VALIDATION REPORT")
	fmt.Println("==========================================")

	// Statistics
	fmt.Printf("📈 Statistics:\n")
	fmt.Printf("   • Total tools validated: %d\n", sv.stats.TotalTools)
	fmt.Printf("   • Tools with array properties: %d\n", sv.stats.ToolsWithArrays)
	fmt.Printf("   • Array properties checked: %d\n", sv.stats.ArrayPropsChecked)
	fmt.Printf("   • Array properties valid: %d\n", sv.stats.ArrayPropsValid)
	fmt.Printf("   • Total required fields: %d\n", sv.stats.RequiredFields)
	fmt.Printf("   • Total optional fields: %d\n", sv.stats.OptionalFields)

	// Error report
	if len(sv.errors) == 0 {
		fmt.Println("\n🎉 SUCCESS: All schemas are valid!")
		fmt.Println("✅ No critical validation errors found")
		fmt.Println("✅ All array properties have proper 'items' definitions")
		fmt.Println("✅ Schemas should be accepted by all AI providers (including Gemini)")
	} else {
		fmt.Printf("\n❌ VALIDATION FAILED: Found %d error(s):\n", len(sv.errors))
		for i, err := range sv.errors {
			fmt.Printf("   %d. %s\n", i+1, err)
		}
	}

	// Warning report
	if len(sv.warnings) > 0 {
		fmt.Printf("\n⚠️  Found %d warning(s):\n", len(sv.warnings))
		for i, warning := range sv.warnings {
			fmt.Printf("   %d. %s\n", i+1, warning)
		}
	}

	// Schema compliance summary
	fmt.Println("\n📋 Schema Compliance Summary:")
	fmt.Printf("   • JSON Schema compliant: %s\n", sv.getComplianceStatus(len(sv.errors) == 0))
	fmt.Printf("   • Gemini API compatible: %s\n", sv.getComplianceStatus(sv.stats.ArrayPropsValid == sv.stats.ArrayPropsChecked))
	fmt.Printf("   • SigNoz conventions: %s\n", sv.getComplianceStatus(sv.countWarningsOfType("convention") == 0))
}

func (sv *SchemaValidator) getComplianceStatus(compliant bool) string {
	if compliant {
		return "✅ PASS"
	}
	return "❌ FAIL"
}

func (sv *SchemaValidator) countWarningsOfType(warningType string) int {
	count := 0
	for _, warning := range sv.warnings {
		if strings.Contains(strings.ToLower(warning), warningType) {
			count++
		}
	}
	return count
}

func (sv *SchemaValidator) exportSchemas(tools map[string]mcp.Tool) {
	fmt.Println("\n💾 Exporting schemas for external validation...")

	// Export as JSON
	data, err := json.MarshalIndent(tools, "", "  ")
	if err != nil {
		fmt.Printf("❌ Failed to marshal schemas: %v\n", err)
		return
	}

	if err := os.WriteFile("signoz_schemas_complete.json", data, 0644); err != nil {
		fmt.Printf("❌ Failed to write schemas file: %v\n", err)
		return
	}

	// Export validation report
	report := map[string]interface{}{
		"validation_timestamp": "2025-08-04T18:00:00Z",
		"total_tools":          sv.stats.TotalTools,
		"errors":               sv.errors,
		"warnings":             sv.warnings,
		"statistics":           sv.stats,
		"gemini_compatible":    sv.stats.ArrayPropsValid == sv.stats.ArrayPropsChecked,
		"validation_passed":    len(sv.errors) == 0,
	}

	reportData, _ := json.MarshalIndent(report, "", "  ")
	os.WriteFile("signoz_validation_report.json", reportData, 0644)

	fmt.Printf("✅ Schemas exported to: signoz_schemas_complete.json\n")
	fmt.Printf("✅ Validation report exported to: signoz_validation_report.json\n")
	fmt.Printf("📊 Total exported: %d tools, %d bytes\n", len(tools), len(data))
}

// getAllSigNozTools returns all SigNoz MCP tools with their expected schemas
func getAllSigNozTools() map[string]mcp.Tool {
	return map[string]mcp.Tool{
		"signoz_test_connection": {
			Name:        "signoz_test_connection",
			Description: "Verify connectivity to your Signoz instance and configuration",
			InputSchema: mcp.ToolInputSchema{
				Type:       "object",
				Properties: map[string]mcp.Property{},
				Required:   []string{},
			},
		},
		"signoz_fetch_dashboards": {
			Name:        "signoz_fetch_dashboards",
			Description: "List all available dashboards from Signoz",
			InputSchema: mcp.ToolInputSchema{
				Type:       "object",
				Properties: map[string]mcp.Property{},
				Required:   []string{},
			},
		},
		"signoz_fetch_dashboard_details": {
			Name:        "signoz_fetch_dashboard_details",
			Description: "Retrieve detailed information about a specific dashboard by its ID",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"dashboard_id": {
						Type:        "string",
						Description: "The ID of the dashboard to fetch details for",
					},
				},
				Required: []string{"dashboard_id"},
			},
		},
		"signoz_fetch_dashboard_data": {
			Name:        "signoz_fetch_dashboard_data",
			Description: "Fetch all panel data for a given dashboard by name and time range",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"dashboard_name": {
						Type:        "string",
						Description: "The name of the dashboard to fetch data for",
					},
					"start_time": {
						Type:        "string",
						Description: "Start time in RFC3339 format (e.g., 2023-01-01T00:00:00Z)",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "End time in RFC3339 format (e.g., 2023-01-01T01:00:00Z)",
						Format:      "date-time",
					},
				},
				Required: []string{"dashboard_name", "start_time", "end_time"},
			},
		},
		"signoz_fetch_apm_metrics": {
			Name:        "signoz_fetch_apm_metrics",
			Description: "Retrieve standard APM metrics (request rate, error rate, latency, apdex) for a given service and time range",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"service_name": {
						Type:        "string",
						Description: "The name of the service to fetch metrics for",
					},
					"start_time": {
						Type:        "string",
						Description: "Start time in RFC3339 format",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "End time in RFC3339 format",
						Format:      "date-time",
					},
				},
				Required: []string{"service_name", "start_time", "end_time"},
			},
		},
		"signoz_fetch_services": {
			Name:        "signoz_fetch_services",
			Description: "Fetch all instrumented services from Signoz with optional time range filtering",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"start_time": {
						Type:        "string",
						Description: "Optional start time in RFC3339 format for filtering",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "Optional end time in RFC3339 format for filtering",
						Format:      "date-time",
					},
				},
				Required: []string{},
			},
		},
		"signoz_execute_clickhouse_query": {
			Name:        "signoz_execute_clickhouse_query",
			Description: "Execute custom Clickhouse SQL queries via the Signoz API with time range support",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"sql": {
						Type:        "string",
						Description: "The SQL query to execute",
					},
					"start_time": {
						Type:        "string",
						Description: "Start time in RFC3339 format",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "End time in RFC3339 format",
						Format:      "date-time",
					},
					"params": {
						Type:        "object",
						Description: "Optional query parameters",
					},
				},
				Required: []string{"sql", "start_time", "end_time"},
			},
		},
		"signoz_execute_builder_query": {
			Name:        "signoz_execute_builder_query",
			Description: "Execute Signoz builder queries for custom metrics and aggregations with time range support",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"query_type": {
						Type:        "string",
						Description: "The type of query (e.g., 'metrics', 'traces')",
					},
					"metric_name": {
						Type:        "string",
						Description: "The name of the metric to query",
					},
					"aggregation": {
						Type:        "string",
						Description: "The aggregation function (e.g., 'sum', 'avg', 'max')",
					},
					"start_time": {
						Type:        "string",
						Description: "Start time in RFC3339 format",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "End time in RFC3339 format",
						Format:      "date-time",
					},
					"filters": {
						Type:        "array",
						Description: "Optional filters for the query",
						Items: &mcp.Property{
							Type: "object",
						},
					},
					"group_by": {
						Type:        "array",
						Description: "Optional fields to group by",
						Items: &mcp.Property{
							Type: "string",
						},
					},
					"step": {
						Type:        "string",
						Description: "Optional step duration (e.g., '1m', '5m')",
					},
					"options": {
						Type:        "object",
						Description: "Optional query options",
					},
				},
				Required: []string{"query_type", "metric_name", "aggregation", "start_time", "end_time"},
			},
		},
		"signoz_fetch_traces_or_logs": {
			Name:        "signoz_fetch_traces_or_logs",
			Description: "Fetch traces or logs from SigNoz using ClickHouse SQL. Specify data_type ('traces' or 'logs'), time range, service name, and limit",
			InputSchema: mcp.ToolInputSchema{
				Type: "object",
				Properties: map[string]mcp.Property{
					"data_type": {
						Type:        "string",
						Description: "The type of data to fetch ('traces' or 'logs')",
						Enum:        []string{"traces", "logs"},
					},
					"service_name": {
						Type:        "string",
						Description: "The name of the service to fetch data for",
					},
					"start_time": {
						Type:        "string",
						Description: "Start time in RFC3339 format",
						Format:      "date-time",
					},
					"end_time": {
						Type:        "string",
						Description: "End time in RFC3339 format",
						Format:      "date-time",
					},
					"limit": {
						Type:        "number",
						Description: "Maximum number of records to return (1-10000)",
						Minimum:     func() *float64 { v := 1.0; return &v }(),
						Maximum:     func() *float64 { v := 10000.0; return &v }(),
					},
				},
				Required: []string{"data_type", "service_name", "start_time", "end_time", "limit"},
			},
		},
	}
}
