// Package mcpclient manages MCP server configuration and short-lived discovery sessions.
package mcpclient

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/zalando/go-keyring"
)

type ServerConfig struct {
	ID               string   `json:"id"`
	Name             string   `json:"name"`
	Enabled          bool     `json:"enabled"`
	Transport        string   `json:"transport"`
	URL              string   `json:"url"`
	Command          string   `json:"command"`
	Args             []string `json:"args"`
	WorkingDirectory string   `json:"workingDirectory"`
	TimeoutSeconds   int      `json:"timeoutSeconds"`
}

// Secrets are write-only and never included in configuration files or views.
type Secrets struct {
	Env     map[string]string `json:"env"`
	Headers map[string]string `json:"headers"`
}
type ServerInput struct {
	Config         ServerConfig `json:"config"`
	Secrets        Secrets      `json:"secrets"`
	ReplaceSecrets bool         `json:"replaceSecrets"`
}
type ServerView struct {
	Config     ServerConfig `json:"config"`
	HasSecrets bool         `json:"hasSecrets"`
}
type ToolInfo struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	InputSchema string `json:"inputSchema"`
}
type TestResult struct {
	OK        bool       `json:"ok"`
	Message   string     `json:"message"`
	ElapsedMs int64      `json:"elapsedMs"`
	Tools     []ToolInfo `json:"tools"`
}

type secretStore interface {
	Get(string) (string, error)
	Set(string, string) error
	Delete(string) error
}
type systemSecrets struct{}

func (systemSecrets) Get(id string) (string, error) { return keyring.Get("com.ituoke.mcp", id) }
func (systemSecrets) Set(id, value string) error    { return keyring.Set("com.ituoke.mcp", id, value) }
func (systemSecrets) Delete(id string) error        { return keyring.Delete("com.ituoke.mcp", id) }

type Service struct {
	mu        sync.Mutex
	path      string
	secrets   secretStore
	servers   []ServerConfig
	loadErr   error
	agentGate chan struct{}
}

func NewService() *Service {
	dir, err := os.UserConfigDir()
	s := newService(filepath.Join(dir, "iTuoke", "mcp.json"), systemSecrets{})
	if err != nil {
		s.loadErr = fmt.Errorf("无法定位 MCP 配置目录")
	}
	if s.loadErr == nil {
		if err := s.ensureMobile(); err != nil {
			s.loadErr = err
		}
	}
	if s.loadErr == nil {
		if err := s.ensureFilesystem(); err != nil {
			s.loadErr = err
		}
	}
	return s
}
func newService(path string, secrets secretStore) *Service {
	s := &Service{path: path, secrets: secrets, servers: []ServerConfig{}}
	raw, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return s
	}
	if err == nil {
		err = json.Unmarshal(raw, &s.servers)
	}
	if err == nil {
		seen := map[string]bool{}
		for i, c := range s.servers {
			if c.ID == "" || seen[c.ID] {
				err = fmt.Errorf("duplicate ID")
				break
			}
			seen[c.ID] = true
			s.servers[i], err = normalize(c)
			if err != nil {
				break
			}
		}
	}
	if err != nil {
		s.loadErr = fmt.Errorf("MCP 配置读取失败，请检查 mcp.json 后重新启动应用")
	}
	return s
}

func normalize(c ServerConfig) (ServerConfig, error) {
	c.Name = strings.TrimSpace(c.Name)
	c.URL = strings.TrimSpace(c.URL)
	c.Command = strings.TrimSpace(c.Command)
	c.WorkingDirectory = strings.TrimSpace(c.WorkingDirectory)
	if c.Name == "" || len(c.Name) > 100 {
		return c, fmt.Errorf("服务名称应为 1–100 字节")
	}
	if c.TimeoutSeconds < 3 || c.TimeoutSeconds > 120 {
		return c, fmt.Errorf("连接超时应为 3–120 秒")
	}
	if len(c.ID) > 64 || strings.ContainsAny(c.ID, "/\\\x00") {
		return c, fmt.Errorf("服务标识无效")
	}
	if c.ID == filesystemServerID && c.Transport != "filesystem" {
		return c, fmt.Errorf("系统内置 Filesystem 的连接方式不可修改")
	}
	if c.ID == mobileServerID && c.Transport != "mobile" {
		return c, fmt.Errorf("系统内置 Mobile Next 的连接方式不可修改")
	}
	switch c.Transport {
	case "mobile":
		c.Command, c.WorkingDirectory, c.URL = "", "", ""
		c.Args = []string{}
	case "filesystem":
		if len(c.Args) == 0 || len(c.Args) > 20 {
			return c, fmt.Errorf("请指定 1–20 个允许访问的目录")
		}
		for _, dir := range c.Args {
			if !filepath.IsAbs(dir) || strings.ContainsAny(dir, "\x00\r\n") {
				return c, fmt.Errorf("允许目录必须为有效的绝对路径")
			}
		}
		c.Command, c.WorkingDirectory, c.URL = "", "", ""
	case "http":
		u, err := url.Parse(c.URL)
		if err != nil || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || (u.Scheme != "https" && u.Scheme != "http") {
			return c, fmt.Errorf("请填写有效的 MCP HTTP 地址，认证信息请放入请求头")
		}
		ip := net.ParseIP(u.Hostname())
		if u.Scheme == "http" && u.Hostname() != "localhost" && (ip == nil || (!ip.IsLoopback() && !ip.IsPrivate())) {
			return c, fmt.Errorf("公网 MCP 地址请使用 HTTPS")
		}
		c.Command, c.WorkingDirectory, c.Args = "", "", []string{}
	case "stdio":
		if c.Command == "" || len(c.Command) > 4096 || strings.ContainsAny(c.Command+c.WorkingDirectory, "\x00\r\n") {
			return c, fmt.Errorf("请填写有效的可执行文件路径或命令名称")
		}
		if c.WorkingDirectory != "" && !filepath.IsAbs(c.WorkingDirectory) {
			return c, fmt.Errorf("工作目录请使用绝对路径")
		}
		if len(c.Args) > 128 {
			return c, fmt.Errorf("启动参数不能超过 128 项")
		}
		for _, arg := range c.Args {
			if len(arg) > 8192 || strings.ContainsRune(arg, 0) {
				return c, fmt.Errorf("启动参数无效或过长")
			}
		}
		if c.Args == nil {
			c.Args = []string{}
		}
		c.URL = ""
	default:
		return c, fmt.Errorf("请选择 stdio 或 Streamable HTTP")
	}
	return c, nil
}

func validateSecrets(c ServerConfig, sec Secrets) error {
	if (c.Transport == "filesystem" || c.Transport == "mobile") && len(sec.Env)+len(sec.Headers) > 0 {
		return fmt.Errorf("内置 MCP 服务不需要凭据")
	}
	raw, _ := json.Marshal(sec)
	if len(raw) > 16000 || len(sec.Env)+len(sec.Headers) > 100 {
		return fmt.Errorf("环境变量或请求头配置过大")
	}
	if c.Transport == "http" && len(sec.Env) > 0 {
		return fmt.Errorf("HTTP 服务请配置请求头")
	}
	if c.Transport == "stdio" && len(sec.Headers) > 0 {
		return fmt.Errorf("stdio 服务请配置环境变量")
	}
	for k, v := range sec.Env {
		if k == "" || strings.ContainsAny(k, "=\x00\r\n") || strings.ContainsRune(v, 0) {
			return fmt.Errorf("环境变量名称或值无效")
		}
	}
	for k, v := range sec.Headers {
		if k == "" || strings.TrimSpace(k) != k || strings.ContainsAny(v, "\x00\r\n") {
			return fmt.Errorf("请求头名称或值无效")
		}
		for _, r := range k {
			if !strings.ContainsRune("!#$%&'*+-.^_`|~0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", r) {
				return fmt.Errorf("请求头名称无效")
			}
		}
		switch strings.ToLower(k) {
		case "host", "content-type", "accept", "content-length", "connection", "transfer-encoding":
			return fmt.Errorf("请勿覆盖 HTTP 协议请求头")
		}
		if strings.HasPrefix(strings.ToLower(k), "mcp-") {
			return fmt.Errorf("请勿覆盖 MCP 协议请求头")
		}
	}
	return nil
}

func (s *Service) index(id string) int {
	for i, c := range s.servers {
		if c.ID == id {
			return i
		}
	}
	return -1
}
func (s *Service) readSecret(id string) (string, error) {
	if id == "" {
		return "", nil
	}
	raw, err := s.secrets.Get(id)
	if errors.Is(err, keyring.ErrNotFound) {
		return "", nil
	}
	if err != nil {
		return "", fmt.Errorf("无法访问系统凭据库")
	}
	return raw, nil
}
func sameTarget(a, b ServerConfig) bool {
	return a.Transport == b.Transport && a.URL == b.URL && a.Command == b.Command && a.WorkingDirectory == b.WorkingDirectory && reflect.DeepEqual(a.Args, b.Args)
}
func (s *Service) resolve(input ServerInput) (ServerConfig, Secrets, string, error) {
	c, err := normalize(input.Config)
	if err != nil {
		return c, Secrets{}, "", err
	}
	if i := s.index(c.ID); i >= 0 && systemTransport(s.servers[i].Transport) && c.Transport != s.servers[i].Transport {
		return c, Secrets{}, "", fmt.Errorf("系统内置 MCP 的连接方式不可修改")
	}
	if err := validateFilesystemDirs(c); err != nil {
		return c, Secrets{}, "", err
	}
	i := s.index(c.ID)
	if c.ID != "" && i < 0 {
		return c, Secrets{}, "", fmt.Errorf("MCP 服务不存在，请重新加载")
	}
	old, err := s.readSecret(c.ID)
	if err != nil {
		return c, Secrets{}, "", err
	}
	sec := input.Secrets
	if !input.ReplaceSecrets {
		if len(sec.Env)+len(sec.Headers) > 0 {
			return c, sec, old, fmt.Errorf("填写新凭据时请选择替换凭据")
		}
		if old != "" && !sameTarget(s.servers[i], c) {
			return c, sec, old, fmt.Errorf("连接目标已更改，请重新填写或清除凭据")
		}
		if old != "" && json.Unmarshal([]byte(old), &sec) != nil {
			return c, sec, old, fmt.Errorf("MCP 凭据读取失败，请重新填写")
		}
	}
	return c, sec, old, validateSecrets(c, sec)
}

func (s *Service) ListServers() ([]ServerView, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.loadErr != nil {
		return nil, s.loadErr
	}
	views := make([]ServerView, 0, len(s.servers))
	for _, c := range s.servers {
		raw, err := s.readSecret(c.ID)
		if err != nil {
			return nil, err
		}
		c.Args = append([]string{}, c.Args...)
		views = append(views, ServerView{Config: c, HasSecrets: raw != ""})
	}
	return views, nil
}

func (s *Service) persist(servers []ServerConfig) error {
	if err := os.MkdirAll(filepath.Dir(s.path), 0700); err != nil {
		return fmt.Errorf("无法创建 MCP 配置目录")
	}
	f, err := os.CreateTemp(filepath.Dir(s.path), "mcp-*.tmp")
	if err != nil {
		return fmt.Errorf("无法写入 MCP 配置")
	}
	defer os.Remove(f.Name())
	data, _ := json.MarshalIndent(servers, "", "  ")
	_, writeErr := f.Write(data)
	closeErr := f.Close()
	if writeErr != nil || closeErr != nil {
		return fmt.Errorf("MCP 配置写入失败")
	}
	if err := os.Rename(f.Name(), s.path); err != nil {
		return fmt.Errorf("MCP 配置保存失败")
	}
	return nil
}
func (s *Service) setSecret(id, raw string) error {
	var err error
	if raw == "" {
		err = s.secrets.Delete(id)
	} else {
		err = s.secrets.Set(id, raw)
	}
	if err != nil && !errors.Is(err, keyring.ErrNotFound) {
		return fmt.Errorf("MCP 凭据保存失败，请检查系统凭据库")
	}
	return nil
}

func (s *Service) SaveServer(input ServerInput) (ServerView, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.loadErr != nil {
		return ServerView{}, s.loadErr
	}
	c, sec, old, err := s.resolve(input)
	if err != nil {
		return ServerView{}, err
	}
	if c.ID == "" {
		if len(s.servers) >= 50 {
			return ServerView{}, fmt.Errorf("最多添加 50 个 MCP 服务")
		}
		var id [16]byte
		if _, err := rand.Read(id[:]); err != nil {
			return ServerView{}, fmt.Errorf("无法创建服务标识")
		}
		c.ID = hex.EncodeToString(id[:])
	}
	for _, other := range s.servers {
		if other.ID != c.ID && strings.EqualFold(other.Name, c.Name) {
			return ServerView{}, fmt.Errorf("服务名称已存在")
		}
	}
	raw := ""
	if len(sec.Env)+len(sec.Headers) > 0 {
		b, _ := json.Marshal(sec)
		raw = string(b)
	}
	if raw != old {
		if err := s.setSecret(c.ID, raw); err != nil {
			return ServerView{}, err
		}
	}
	next := append([]ServerConfig{}, s.servers...)
	if i := s.index(c.ID); i >= 0 {
		next[i] = c
	} else {
		next = append(next, c)
	}
	if err := s.persist(next); err != nil {
		if raw != old {
			if rollback := s.setSecret(c.ID, old); rollback != nil {
				return ServerView{}, fmt.Errorf("配置保存失败，凭据恢复失败，请重新填写凭据")
			}
		}
		return ServerView{}, err
	}
	s.servers = next
	return ServerView{Config: c, HasSecrets: raw != ""}, nil
}

func (s *Service) SetEnabled(id string, enabled bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.loadErr != nil {
		return s.loadErr
	}
	i := s.index(id)
	if i < 0 {
		return fmt.Errorf("MCP 服务不存在")
	}
	next := append([]ServerConfig{}, s.servers...)
	next[i].Enabled = enabled
	if err := s.persist(next); err != nil {
		return err
	}
	s.servers = next
	return nil
}

func (s *Service) DeleteServer(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.loadErr != nil {
		return s.loadErr
	}
	if id == mobileServerID {
		return fmt.Errorf("系统内置 Mobile Next 可停用，无需删除")
	}
	i := s.index(id)
	if i < 0 {
		return fmt.Errorf("MCP 服务不存在")
	}
	if systemTransport(s.servers[i].Transport) {
		return fmt.Errorf("系统内置 MCP 可停用或编辑，不允许删除")
	}
	old, err := s.readSecret(id)
	if err != nil {
		return err
	}
	if old != "" {
		if err := s.setSecret(id, ""); err != nil {
			return err
		}
	}
	next := append([]ServerConfig{}, s.servers[:i]...)
	next = append(next, s.servers[i+1:]...)
	if err := s.persist(next); err != nil {
		if old != "" {
			if rollback := s.setSecret(id, old); rollback != nil {
				return fmt.Errorf("删除失败，凭据恢复失败，请重新填写凭据")
			}
		}
		return err
	}
	s.servers = next
	return nil
}

// TestServer discovers tools using the draft without saving or executing tools.
// Sessions are closed after discovery; enabled is a saved availability preference.
func (s *Service) TestServer(ctx context.Context, input ServerInput) TestResult {
	start := time.Now()
	result := TestResult{Tools: []ToolInfo{}}
	s.mu.Lock()
	c, sec, _, err := s.resolve(input)
	if s.loadErr != nil {
		err = s.loadErr
	}
	s.mu.Unlock()
	if err == nil {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Duration(c.TimeoutSeconds)*time.Second)
		defer cancel()
		result.Tools, err = discover(ctx, c, sec)
		if err != nil {
			switch {
			case errors.Is(ctx.Err(), context.Canceled):
				err = fmt.Errorf("连接测试已取消")
			case errors.Is(ctx.Err(), context.DeadlineExceeded):
				err = fmt.Errorf("连接超时，请检查服务是否启动以及连接地址")
			default:
				err = fmt.Errorf("连接或工具发现失败，请检查启动命令、运行环境、MCP 地址和认证配置")
			}
		}
	}
	result.ElapsedMs = time.Since(start).Milliseconds()
	result.OK = err == nil
	if err != nil {
		result.Message = err.Error()
		result.Tools = []ToolInfo{}
	} else {
		result.Message = fmt.Sprintf("连接成功，发现 %d 个工具", len(result.Tools))
	}
	return result
}

// Request headers are attached only to this service's own HTTP client.
type headerTransport struct {
	base    http.RoundTripper
	headers map[string]string
	ctx     context.Context
}

func (t headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	ctx, cancel := context.WithCancel(req.Context())
	stop := context.AfterFunc(t.ctx, cancel)
	cleanup := func() { stop(); cancel() }
	clone := req.Clone(ctx)
	for k, v := range t.headers {
		clone.Header.Set(k, v)
	}
	res, err := t.base.RoundTrip(clone)
	if err != nil {
		cleanup()
		return nil, err
	}
	res.Body = &cleanupBody{ReadCloser: res.Body, cleanup: cleanup}
	return res, nil
}

type cleanupBody struct {
	io.ReadCloser
	cleanup func()
}

func (b *cleanupBody) Close() error { defer b.cleanup(); return b.ReadCloser.Close() }