package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"
	"github.com/firebase/genkit/go/plugins/googlegenai"
	gkmcp "github.com/firebase/genkit/go/plugins/mcp"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func TestMCPClientTools(t *testing.T) {
	ctx := t.Context()
	mux := http.NewServeMux()
	mux.Handle("/mcpdemoservers/reverse", NewReverseListHandler())
	s := httptest.NewServer(mux)
	t.Cleanup(s.Close)

	// Initialize Genkit with Google AI
	g := genkit.Init(ctx, genkit.WithPlugins(&googlegenai.GoogleAI{}))

	// Create and connect to MCP server
	client, err := gkmcp.NewGenkitMCPClient(gkmcp.MCPClientOptions{
		Name:    "list-reverser",
		Version: "1.0.0",
		StreamableHTTP: &gkmcp.StreamableHTTPConfig{
			BaseURL: s.URL + "/mcpdemoservers/reverse",
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	defer func(client *gkmcp.GenkitMCPClient) {
		err := client.Disconnect()
		if err != nil {
			t.Logf("Failed to disconnect MCP client: %v", err)
		}
	}(client)

	// Get tools and generate response
	tools, _ := client.GetActiveTools(ctx, g)
	t.Logf("tools: %v", tools)

	var toolRefs []ai.ToolRef
	for _, tool := range tools {
		toolRefs = append(toolRefs, tool)
	}

	response, err := genkit.Generate(ctx, g,
		ai.WithModelName("googleai/gemini-2.5-flash"),
		ai.WithPrompt("Tell me about the tools you can access."),
		ai.WithTools(toolRefs...),
		ai.WithToolChoice(ai.ToolChoiceAuto),
	)
	if err != nil {
		t.Fatalf("Generation failed: %v", err)
	}
	t.Logf("Response: %s", response.Text())
}

func NewReverseListHandler() http.Handler {
	server := mcp.NewServer(&mcp.Implementation{
		Name:    "list-reverser",
		Version: "1.0.0",
	}, nil)

	// Only one tool: reverse a list of numbers
	mcp.AddTool(server, &mcp.Tool{
		Name:        "reverse_list",
		Description: "Reverses the order of a list of numbers.",
	}, ReverseList)

	return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { return server }, nil)
}

type ReverseListInput struct {
	Numbers []float64 `json:"numbers" jsonschema:"description:List of numbers to reverse"`
}

func ReverseList(ctx context.Context, req *mcp.CallToolRequest, input ReverseListInput) (*mcp.CallToolResult, any, error) {
	n := len(input.Numbers)
	reversed := make([]float64, n)
	for i := range input.Numbers {
		reversed[i] = input.Numbers[n-1-i]
	}

	return &mcp.CallToolResult{
		Content: []mcp.Content{&mcp.TextContent{Text: prettyJSON(reversed)}},
	}, nil, nil
}

func prettyJSON(v any) string {
	b, _ := json.MarshalIndent(v, "", "  ")
	return string(b)
}

func TestToGeminiSchema(t *testing.T) {
	tests := []struct {
		name         string
		genkitSchema map[string]any
		want         *genai.Schema
		wantErr      bool
	}{
		{
			name: "string type",
			genkitSchema: map[string]any{
				"type": "string",
			},
			want: &genai.Schema{
				Type: genai.TypeString,
			},
		},
		{
			name: "array type [string, null]",
			genkitSchema: map[string]any{
				"type": []any{"string", "null"},
			},
			want: &genai.Schema{
				Type:     genai.TypeString,
				Nullable: genai.Ptr(true),
			},
		},
		{
			name: "array type [null, integer]",
			genkitSchema: map[string]any{
				"type": []any{"null", "integer"},
			},
			want: &genai.Schema{
				Type:     genai.TypeInteger,
				Nullable: genai.Ptr(true),
			},
		},
		{
			name: "array type [null]",
			genkitSchema: map[string]any{
				"type": []any{"null"},
			},
			want: &genai.Schema{
				Nullable: genai.Ptr(true),
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := toGeminiSchema(tt.genkitSchema, tt.genkitSchema)
			if (err != nil) != tt.wantErr {
				t.Errorf("toGeminiSchema() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if diff := cmp.Diff(tt.want, got); diff != "" {
				t.Errorf("toGeminiSchema() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func toGeminiSchema(originalSchema map[string]any, genkitSchema map[string]any) (*genai.Schema, error) {
	// this covers genkitSchema == nil and {}
	// genkitSchema will be {} if it's any
	if len(genkitSchema) == 0 {
		return nil, nil
	}
	if v, ok := genkitSchema["$ref"]; ok {
		ref, ok := v.(string)
		if !ok {
			return nil, fmt.Errorf("invalid $ref value: not a string")
		}
		s, err := resolveRef(originalSchema, ref)
		if err != nil {
			return nil, err
		}
		return toGeminiSchema(originalSchema, s)
	}

	// Handle "anyOf" subschemas by finding the first valid schema definition
	if v, ok := genkitSchema["anyOf"]; ok {
		if anyOfList, isList := v.([]map[string]any); isList {
			for _, subSchema := range anyOfList {
				if subSchemaType, hasType := subSchema["type"]; hasType {
					if typeStr, isString := subSchemaType.(string); isString && typeStr != "null" {
						if title, ok := genkitSchema["title"]; ok {
							subSchema["title"] = title
						}
						if description, ok := genkitSchema["description"]; ok {
							subSchema["description"] = description
						}
						// Found a schema like: {"type": "string"}
						return toGeminiSchema(originalSchema, subSchema)
					}
				}
			}
		}
	}

	schema := &genai.Schema{}
	typeVal, ok := genkitSchema["type"]
	if !ok {
		return nil, fmt.Errorf("schema is missing the 'type' field: %#v", genkitSchema)
	}

	typeStr, ok := typeVal.(string)
	if !ok {
		return nil, fmt.Errorf("schema 'type' field is not a string, but %T", typeVal)
	}

	switch typeStr {
	case "string":
		schema.Type = genai.TypeString
	case "float64", "number":
		schema.Type = genai.TypeNumber
	case "integer":
		schema.Type = genai.TypeInteger
	case "boolean":
		schema.Type = genai.TypeBoolean
	case "object":
		schema.Type = genai.TypeObject
	case "array":
		schema.Type = genai.TypeArray
	default:
		return nil, fmt.Errorf("schema type %q not allowed", genkitSchema["type"])
	}
	if v, ok := genkitSchema["required"]; ok {
		schema.Required = castToStringArray(v)
	}
	if v, ok := genkitSchema["propertyOrdering"]; ok {
		schema.PropertyOrdering = castToStringArray(v)
	}
	if v, ok := genkitSchema["description"]; ok {
		schema.Description = v.(string)
	}
	if v, ok := genkitSchema["format"]; ok {
		schema.Format = v.(string)
	}
	if v, ok := genkitSchema["title"]; ok {
		schema.Title = v.(string)
	}
	if v, ok := genkitSchema["minItems"]; ok {
		if i64, ok := castToInt64(v); ok {
			schema.MinItems = genai.Ptr(i64)
		}
	}
	if v, ok := genkitSchema["maxItems"]; ok {
		if i64, ok := castToInt64(v); ok {
			schema.MaxItems = genai.Ptr(i64)
		}
	}
	if v, ok := genkitSchema["maximum"]; ok {
		if f64, ok := castToFloat64(v); ok {
			schema.Maximum = genai.Ptr(f64)
		}
	}
	if v, ok := genkitSchema["minimum"]; ok {
		if f64, ok := castToFloat64(v); ok {
			schema.Minimum = genai.Ptr(f64)
		}
	}
	if v, ok := genkitSchema["enum"]; ok {
		schema.Enum = castToStringArray(v)
	}
	if v, ok := genkitSchema["items"]; ok {
		items, err := toGeminiSchema(originalSchema, v.(map[string]any))
		if err != nil {
			return nil, err
		}
		schema.Items = items
	}
	if val, ok := genkitSchema["properties"]; ok {
		props := map[string]*genai.Schema{}
		for k, v := range val.(map[string]any) {
			p, err := toGeminiSchema(originalSchema, v.(map[string]any))
			if err != nil {
				return nil, err
			}
			props[k] = p
		}
		schema.Properties = props
	}
	// Nullable -- not supported in jsonschema.Schema

	return schema, nil
}