package pkulaw

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner/mcp"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
)

const (
	connectionPreparationTimeout = 35 * time.Second
	connectionEndpointTimeout    = 15 * time.Second
	connectionCleanupTimeout     = 3 * time.Second
)

var fallbackProbeEndpoints = []EndpointID{
	EndpointLawSemantic,
	EndpointLawKeyword,
	EndpointCaseSemantic,
	EndpointCaseKeyword,
}

// PrepareConnection first probes the fixed aggregate service. If it is not
// usable, all four fixed fallback services must independently negotiate and
// clean up successfully. Only the closed topology enum is returned.
func PrepareConnection(
	ctx context.Context,
	credentials runner.Credentials,
) (runner.ConnectionConfig, error) {
	clients, err := newProbeClients(nil)
	if err != nil {
		return nil, runner.Errorf(
			runner.ErrorConnectionUnavailable,
			"PKULaw connection preparation is unavailable",
			nil,
		)
	}
	return prepareConnection(ctx, clients, credentials)
}

func newProbeClients(transport http.RoundTripper) (mcpClients, error) {
	if err := validateRouting(); err != nil {
		return nil, err
	}
	clients := make(mcpClients, len(endpointPolicies))
	for _, policy := range endpointPolicies {
		client, err := mcp.NewClient(mcp.Options{
			Endpoint:       policy.Endpoint,
			Auth:           AuthSpec(),
			ProbeOnly:      true,
			Transport:      transport,
			Timeout:        connectionEndpointTimeout,
			CleanupTimeout: connectionCleanupTimeout,
		})
		if err != nil {
			return nil, fmt.Errorf(
				"PKULaw: build %s MCP probe client: %w",
				policy.ID,
				err,
			)
		}
		clients[policy.ID] = client
	}
	return clients, nil
}

func prepareConnection(
	ctx context.Context,
	clients mcpClients,
	credentials runner.Credentials,
) (runner.ConnectionConfig, error) {
	if ctx == nil {
		return nil, runner.Errorf(
			runner.ErrorValidation,
			"PKULaw connection preparation context is unavailable",
			nil,
		)
	}
	if err := validateClientSet(clients); err != nil {
		return nil, runner.Errorf(
			runner.ErrorConnectionUnavailable,
			"PKULaw probe clients are incomplete",
			nil,
		)
	}
	preparationCtx, cancel := context.WithTimeout(
		ctx,
		connectionPreparationTimeout,
	)
	defer cancel()

	aggregateErr := clients[EndpointAggregate].Probe(
		preparationCtx,
		credentials.Clone(),
	)
	if aggregateErr == nil {
		config := aggregateConnectionConfig()
		if err := ConnectionConfigContract().Validate(config); err != nil {
			return nil, runner.Errorf(
				runner.ErrorConnectionUnavailable,
				"PKULaw aggregate topology is unavailable",
				nil,
			)
		}
		return config, nil
	}

	type probeResult struct {
		endpoint EndpointID
		err      error
	}
	results := make(chan probeResult, len(fallbackProbeEndpoints))
	for _, endpoint := range fallbackProbeEndpoints {
		endpoint := endpoint
		client := clients[endpoint]
		logging.Go(slog.Default(), func() {
			results <- probeResult{
				endpoint: endpoint,
				err: client.Probe(
					preparationCtx,
					credentials.Clone(),
				),
			}
		})
	}
	errorsByEndpoint := make(
		map[EndpointID]error,
		len(fallbackProbeEndpoints),
	)
	for range fallbackProbeEndpoints {
		result := <-results
		errorsByEndpoint[result.endpoint] = result.err
	}
	fallbackComplete := true
	for _, endpoint := range fallbackProbeEndpoints {
		if errorsByEndpoint[endpoint] != nil {
			fallbackComplete = false
			break
		}
	}
	if fallbackComplete {
		config := fallbackConnectionConfig()
		if err := ConnectionConfigContract().Validate(config); err != nil {
			return nil, runner.Errorf(
				runner.ErrorConnectionUnavailable,
				"PKULaw fallback topology is unavailable",
				nil,
			)
		}
		return config, nil
	}

	errorsByEndpoint[EndpointAggregate] = aggregateErr
	probePriority := make(
		[]EndpointID,
		0,
		len(fallbackProbeEndpoints)+1,
	)
	probePriority = append(probePriority, EndpointAggregate)
	probePriority = append(probePriority, fallbackProbeEndpoints...)
	for _, endpoint := range probePriority {
		err := errorsByEndpoint[endpoint]
		if runner.CodeOf(err) == runner.ErrorProviderAuthUnavailable {
			return nil, err
		}
	}
	for _, endpoint := range probePriority {
		if err := errorsByEndpoint[endpoint]; err != nil {
			return nil, err
		}
	}
	return nil, runner.Errorf(
		runner.ErrorConnectionUnavailable,
		"PKULaw connection topology is unavailable",
		nil,
	)
}