package tools

import (
	"context"
	"fmt"

	"github.com/ClareAI/astra-insight-mcp/internal/gateway"
	"github.com/ClareAI/astra-insight-mcp/internal/mcp"
)

func registerStartTestPipeline(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "start_test_pipeline",
		Description: "Start the full test pipeline: Generate test cases → Execute → Evaluate → Analyze. Returns a job_id for polling. This is the easiest way to run a complete test cycle.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":       {Type: "string", Description: "Agent ID"},
				"user_id":        {Type: "string", Description: "User ID"},
				"instruction":    {Type: "string", Description: "Optional instruction override for the agent"},
				"run_analyze":    {Type: "boolean", Description: "Whether to run analysis after evaluation (default true)"},
			},
			Required: []string{"agent_id", "user_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		userID, _ := args["user_id"].(string)
		if apiKey == "" || agentID == "" || userID == "" {
			return nil, fmt.Errorf("api_key, agent_id, and user_id are required")
		}

		req := gateway.PipelineRequest{
			AgentID: agentID,
			UserID:  userID,
		}
		if v, ok := args["instruction"].(string); ok {
			req.Instruction = v
		}
		if v, ok := args["run_analyze"].(bool); ok {
			req.RunAnalyze = &v
		}

		return gw.StartTestPipeline(ctx, apiKey, req)
	})
}

func registerListPipelineJobs(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "list_pipeline_jobs",
		Description: "List pipeline jobs for an agent — see all past and running test pipelines.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":  {Type: "string", Description: "Agent ID"},
				"limit":     {Type: "integer", Description: "Max number of jobs to return (default 50)"},
			},
			Required: []string{"agent_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		if apiKey == "" || agentID == "" {
			return nil, fmt.Errorf("api_key and agent_id are required")
		}

		limit := 50
		if v, ok := args["limit"].(float64); ok && v > 0 {
			limit = int(v)
		}

		return gw.ListPipelineJobs(ctx, apiKey, agentID, limit)
	})
}

func registerGetPipelineJob(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_pipeline_job",
		Description: "Get pipeline job status and results. Poll this after starting a pipeline to check progress and see results when complete.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"job_id":  {Type: "string", Description: "Pipeline job ID"},
			},
			Required: []string{"job_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		jobID, _ := args["job_id"].(string)
		if apiKey == "" || jobID == "" {
			return nil, fmt.Errorf("api_key and job_id are required")
		}

		return gw.GetPipelineJob(ctx, apiKey, jobID)
	})
}

func registerCreateTestSuite(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "create_test_suite",
		Description: "Generate test cases for an agent asynchronously. Creates a test suite with auto-generated test cases across 6 categories. Poll get_test_suite_progress for status.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":     {Type: "string", Description: "Agent ID"},
				"instruction":  {Type: "string", Description: "Optional instruction override"},
				"name":         {Type: "string", Description: "Optional name for the test suite"},
			},
			Required: []string{"agent_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		if apiKey == "" || agentID == "" {
			return nil, fmt.Errorf("api_key and agent_id are required")
		}

		req := gateway.CreateTestSuiteRequest{
			AgentID: agentID,
		}
		if v, ok := args["instruction"].(string); ok {
			req.Instruction = v
		}
		if v, ok := args["name"].(string); ok {
			req.Name = v
		}

		return gw.CreateTestSuite(ctx, apiKey, req)
	})
}

func registerListTestSuites(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "list_test_suites",
		Description: "List test suites for an agent. Filter by status (generating/ready/failed).",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":  {Type: "string", Description: "Agent ID (optional)"},
				"status":    {Type: "string", Description: "Filter by status", Enum: []string{"generating", "ready", "failed"}},
				"limit":     {Type: "integer", Description: "Max number of suites to return"},
			},
			Required: []string{},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		if apiKey == "" {
			return nil, fmt.Errorf("api_key is required")
		}

		opts := gateway.TestSuiteListOpts{}
		if v, ok := args["agent_id"].(string); ok {
			opts.AgentID = v
		}
		if v, ok := args["status"].(string); ok {
			opts.Status = v
		}
		if v, ok := args["limit"].(float64); ok && v > 0 {
			opts.Limit = int(v)
		}

		return gw.ListTestSuites(ctx, apiKey, opts)
	})
}

func registerGetTestSuite(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_test_suite",
		Description: "Get a test suite with all its test cases. Can be called during generation to see partial results.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"generator_id":  {Type: "string", Description: "Test suite generator ID"},
			},
			Required: []string{"generator_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		generatorID, _ := args["generator_id"].(string)
		if apiKey == "" || generatorID == "" {
			return nil, fmt.Errorf("api_key and generator_id are required")
		}

		return gw.GetTestSuite(ctx, apiKey, generatorID)
	})
}

func registerGetTestSuiteProgress(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_test_suite_progress",
		Description: "Get lightweight generation progress for a test suite. Optimized for frequent polling (1-5 second intervals).",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"generator_id": {Type: "string", Description: "Test suite generator ID"},
			},
			Required: []string{"generator_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		generatorID, _ := args["generator_id"].(string)
		if apiKey == "" || generatorID == "" {
			return nil, fmt.Errorf("api_key and generator_id are required")
		}

		return gw.GetTestSuiteProgress(ctx, apiKey, generatorID)
	})
}

func registerAddTestCase(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "add_test_case",
		Description: "Manually add a test case to an existing test suite. Optionally use LLM to auto-classify category and extract answer assertions.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"generator_id":     {Type: "string", Description: "Test suite generator ID"},
				"user_query":       {Type: "string", Description: "User query for the test case"},
				"expected_answer":  {Type: "string", Description: "Optional expected answer"},
				"use_llm":          {Type: "boolean", Description: "Use LLM to classify and extract assertions (default true)"},
			},
			Required: []string{"generator_id", "user_query"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		generatorID, _ := args["generator_id"].(string)
		userQuery, _ := args["user_query"].(string)
		if apiKey == "" || generatorID == "" || userQuery == "" {
			return nil, fmt.Errorf("api_key, generator_id, and user_query are required")
		}

		req := gateway.AddTestCaseRequest{
			UserQuery: userQuery,
		}
		if v, ok := args["expected_answer"].(string); ok {
			req.ExpectedAnswer = v
		}
		if v, ok := args["use_llm"].(bool); ok {
			req.UseLLM = &v
		}

		return gw.AddTestCase(ctx, apiKey, generatorID, req)
	})
}

func registerCreateTestRun(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "create_test_run",
		Description: "Execute and evaluate all test cases in a suite against the agent. Optionally override the instruction for A/B testing.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"generator_id":     {Type: "string", Description: "Test suite generator ID"},
				"user_id":          {Type: "string", Description: "User ID"},
				"new_instruction":  {Type: "string", Description: "Optional instruction override for A/B testing"},
				"run_analyze":     {Type: "boolean", Description: "Whether to run analysis after evaluation"},
			},
			Required: []string{"generator_id", "user_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		generatorID, _ := args["generator_id"].(string)
		userID, _ := args["user_id"].(string)
		if apiKey == "" || generatorID == "" || userID == "" {
			return nil, fmt.Errorf("api_key, generator_id, and user_id are required")
		}

		req := gateway.CreateRunRequest{
			UserID: userID,
		}
		if v, ok := args["new_instruction"].(string); ok {
			req.NewInstruction = v
		}
		if v, ok := args["run_analyze"].(bool); ok {
			req.RunAnalyze = &v
		}

		return gw.CreateTestRun(ctx, apiKey, generatorID, req)
	})
}

func registerListTestRuns(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "list_test_runs",
		Description: "List test runs for an agent — see execution history and pass/fail rates.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":  {Type: "string", Description: "Agent ID"},
				"limit":     {Type: "integer", Description: "Max number of runs to return (default 50)"},
			},
			Required: []string{"agent_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		if apiKey == "" || agentID == "" {
			return nil, fmt.Errorf("api_key and agent_id are required")
		}

		limit := 50
		if v, ok := args["limit"].(float64); ok && v > 0 {
			limit = int(v)
		}

		return gw.ListTestRuns(ctx, apiKey, agentID, limit)
	})
}

func registerGetTestRun(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_test_run",
		Description: "Get full test run results — individual test case results, pass/fail status, scores, and agent responses.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"run_id":   {Type: "string", Description: "Test run ID"},
			},
			Required: []string{"run_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		runID, _ := args["run_id"].(string)
		if apiKey == "" || runID == "" {
			return nil, fmt.Errorf("api_key and run_id are required")
		}

		return gw.GetTestRun(ctx, apiKey, runID)
	})
}

func registerCreateRunAnalysis(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "create_run_analysis",
		Description: "Analyze a completed test run — generates improvement suggestions and a recommended instruction based on failures.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"run_id":  {Type: "string", Description: "Test run ID"},
			},
			Required: []string{"run_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		runID, _ := args["run_id"].(string)
		if apiKey == "" || runID == "" {
			return nil, fmt.Errorf("api_key and run_id are required")
		}

		return gw.CreateRunAnalysis(ctx, apiKey, runID)
	})
}

func registerGetRunAnalysis(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_run_analysis",
		Description: "Get the analysis results for a test run — improvement suggestions, failure patterns, and recommended instruction.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"run_id":  {Type: "string", Description: "Test run ID"},
			},
			Required: []string{"run_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		runID, _ := args["run_id"].(string)
		if apiKey == "" || runID == "" {
			return nil, fmt.Errorf("api_key and run_id are required")
		}

		return gw.GetRunAnalysis(ctx, apiKey, runID)
	})
}

func registerAnalyzeFailures(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "analyze_failures",
		Description: "Analyze test failures without generating a new instruction (Phase 1). Returns analysis_json that can be passed to generate_instruction (Phase 2).",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"run_id":  {Type: "string", Description: "Test run ID"},
			},
			Required: []string{"run_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		runID, _ := args["run_id"].(string)
		if apiKey == "" || runID == "" {
			return nil, fmt.Errorf("api_key and run_id are required")
		}

		return gw.AnalyzeFailures(ctx, apiKey, runID)
	})
}

func registerGenerateInstruction(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "generate_instruction",
		Description: "Generate an improved instruction from analysis results (Phase 2). Pass analysis_json from analyze_failures or a list of improvement suggestions.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"instruction":             {Type: "string", Description: "Original agent instruction"},
				"analysis_json":           {Type: "object", Description: "Analysis JSON from analyze_failures"},
				"improvement_suggestions":  {Type: "array", Description: "Array of {priority, suggestion} objects from analyze_failures"},
				"agent_description":        {Type: "string", Description: "Short description of the agent's purpose"},
			},
			Required: []string{"instruction"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		instruction, _ := args["instruction"].(string)
		if apiKey == "" || instruction == "" {
			return nil, fmt.Errorf("api_key and instruction are required")
		}

		req := gateway.GenerateInstructionRequest{
			Instruction: instruction,
		}
		if v, ok := args["analysis_json"]; ok && v != nil {
			req.AnalysisJSON = v
		}
		if v, ok := args["agent_description"].(string); ok && v != "" {
			req.AgentDescription = v
		}
		if v, ok := args["improvement_suggestions"]; ok {
			if arr, ok := v.([]interface{}); ok {
				for _, item := range arr {
					switch val := item.(type) {
					case map[string]interface{}:
						s := gateway.ImprovementSuggestion{}
						if p, ok := val["priority"].(string); ok {
							s.Priority = p
						}
						if sg, ok := val["suggestion"].(string); ok {
							s.Suggestion = sg
						}
						req.ImprovementSuggestions = append(req.ImprovementSuggestions, s)
					case string:
						req.ImprovementSuggestions = append(req.ImprovementSuggestions, gateway.ImprovementSuggestion{Suggestion: val})
					}
				}
			}
		}

		return gw.GenerateInstruction(ctx, apiKey, req)
	})
}

func registerGetAgentTestOverview(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "get_agent_test_overview",
		Description: "Get a testing summary for an agent — total test suites, total runs, latest run results, and overall pass rate.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id": {Type: "string", Description: "Agent ID"},
			},
			Required: []string{"agent_id"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		if apiKey == "" || agentID == "" {
			return nil, fmt.Errorf("api_key and agent_id are required")
		}

		return gw.GetAgentTestOverview(ctx, apiKey, agentID)
	})
}
