package mcpclient

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"github.com/google/jsonschema-go/jsonschema"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// AgentTool keeps provider function names unique across MCP servers.
type AgentTool struct {
	Name        string
	Server      string
	Tool        string
	Description string
	Schema      json.RawMessage
	validator   *jsonschema.Resolved
}

type liveServer struct {
	config  ServerConfig
	secrets Secrets
	session *mcp.ClientSession
	cancel  context.CancelFunc
	done    chan struct{}
}

type Toolset struct {
	service *Service
	servers []*liveServer
	tools   []AgentTool
	targets map[string]*liveServer
	gate    chan struct{}
}

// OpenTools gives each chat its own sessions. It discovers only enabled servers;
// configuration/credentials are snapshotted under lock, never during network IO.
func OpenTools(ctx context.Context, service *Service, warning func(string), filters ...func(ServerConfig) bool) (*Toolset, error) {
	set := &Toolset{service: service, targets: map[string]*liveServer{}}
	if service == nil {
		return set, nil
	}
	service.mu.Lock()
	if service.loadErr != nil {
		service.mu.Unlock()
		return nil, fmt.Errorf("MCP 配置读取失败，请检查设置")
	}
	var candidates []*liveServer
	for _, config := range service.servers {
		if !config.Enabled {
			continue
		}
		allowed := true
		for _, filter := range filters {
			if !filter(config) {
				allowed = false
				break
			}
		}
		if !allowed {
			continue
		}
		_, secret, _, err := service.resolve(ServerInput{Config: config})
		if err != nil {
			if warning != nil {
				warning("MCP「" + config.Name + "」配置不可用")
			}
			continue
		}
		candidates = append(candidates, &liveServer{config: config, secrets: secret})
	}
	service.mu.Unlock()
	for _, live := range candidates {
		if ctx.Err() != nil {
			set.Close()
			return nil, ctx.Err()
		}
		child, cancel := context.WithCancel(ctx)
		live.cancel, live.done = cancel, make(chan struct{})
		ready := make(chan error, 1)
		var found []AgentTool
		go func() {
			defer close(live.done)
			err := withSession(child, live.config, live.secrets, func(session *mcp.ClientSession) error {
				live.session = session
				for tool, err := range session.Tools(child, nil) {
					if err != nil {
						return err
					}
					if !supportedTool(live.config, tool.Name) {
						continue
					}
					if len(found) >= 128 {
						return fmt.Errorf("服务工具超过 128 个")
					}
					schema, err := json.Marshal(tool.InputSchema)
					if err != nil || len(schema) > 64*1024 {
						return fmt.Errorf("工具定义过大或无效")
					}
					var definition jsonschema.Schema
					if err := json.Unmarshal(schema, &definition); err != nil {
						return fmt.Errorf("工具参数定义无效")
					}
					validator, err := definition.Resolve(nil)
					if err != nil {
						return fmt.Errorf("工具参数定义无法校验")
					}
					hash := fmt.Sprintf("%x", sha256.Sum256([]byte(live.config.ID+":"+tool.Name)))[:12]
					name := strings.Map(func(r rune) rune {
						if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' {
							return r
						}
						return '_'
					}, tool.Name)
					if len(name) > 45 {
						name = name[:45]
					}
					found = append(found, AgentTool{Name: "mcp_" + hash + "_" + name, Server: live.config.Name, Tool: tool.Name, Description: toolDescription(live.config, tool.Name, tool.Description), Schema: schema, validator: validator})
				}
				ready <- nil
				<-child.Done()
				return nil
			})
			if err != nil {
				select {
				case ready <- err:
				default:
				}
			}
		}()
		timer := time.NewTimer(time.Duration(live.config.TimeoutSeconds) * time.Second)
		var err error
		select {
		case err = <-ready:
		case <-ctx.Done():
			err = ctx.Err()
		case <-timer.C:
			err = context.DeadlineExceeded
		}
		timer.Stop()
		if err != nil || len(set.tools)+len(found) > 128 {
			cancel()
			<-live.done
			if ctx.Err() != nil {
				set.Close()
				return nil, ctx.Err()
			}
			if warning != nil {
				warning("MCP「" + live.config.Name + "」连接或工具发现失败（最多支持 128 个工具）")
			}
			continue
		}
		set.servers = append(set.servers, live)
		for _, tool := range found {
			set.tools = append(set.tools, tool)
			set.targets[tool.Name] = live
		}
	}
	return set, nil
}

func (s *Toolset) Close() {
	for _, server := range s.servers {
		server.cancel()
	}
	for _, server := range s.servers {
		<-server.done
	}
	if s.gate != nil {
		<-s.gate
		s.gate = nil
	}
}
func (s *Toolset) Tools() []AgentTool { return s.tools }

func (s *Toolset) Call(ctx context.Context, name, arguments string) (*mcp.CallToolResult, error) {
	live := s.targets[name]
	if live == nil {
		return nil, fmt.Errorf("未知的 MCP 工具")
	}
	// Serialize action sequences across chats. One lease for all MCP servers
	// avoids deadlocks when two chats need the same servers in opposite order.
	// Text-only generation remains concurrent; waiting for the lease can cancel.
	if s.gate == nil {
		s.service.mu.Lock()
		gate := s.service.agentGate
		if gate == nil {
			gate = make(chan struct{}, 1)
			s.service.agentGate = gate
		}
		s.service.mu.Unlock()
		select {
		case gate <- struct{}{}:
			s.gate = gate
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	var args map[string]any
	if len(arguments) > 64*1024 || json.Unmarshal([]byte(arguments), &args) != nil || args == nil {
		return nil, fmt.Errorf("工具参数必须是 JSON 对象（最多 64 KB）")
	}
	// An already-open session must not bypass a later disable/delete/target edit.
	s.service.mu.Lock()
	i := s.service.index(live.config.ID)
	valid := i >= 0 && s.service.servers[i].Enabled && sameTarget(s.service.servers[i], live.config)
	s.service.mu.Unlock()
	if !valid {
		return nil, fmt.Errorf("MCP 服务已停用、删除或修改，请重新发起对话")
	}
	var original string
	for _, tool := range s.tools {
		if tool.Name == name {
			if err := tool.validator.Validate(args); err != nil {
				return nil, fmt.Errorf("工具参数不符合定义: %.1000s", err.Error())
			}
			original = tool.Tool
			break
		}
	}
	callCtx, cancel := context.WithTimeout(ctx, time.Duration(live.config.TimeoutSeconds)*time.Second)
	defer cancel()
	result, err := live.session.CallTool(callCtx, &mcp.CallToolParams{Name: original, Arguments: args})
	if err != nil {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		message := err.Error()
		for _, values := range []map[string]string{live.secrets.Env, live.secrets.Headers} {
			for _, value := range values {
				if value != "" {
					message = strings.ReplaceAll(message, value, "[已隐藏凭据]")
				}
			}
		}
		return nil, fmt.Errorf("MCP 工具执行失败: %.1000s", message)
	}
	return result, nil
}