package mcp

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime/debug"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"codey/internal/config"
	"codey/internal/tools"
)

// clientVersion is what initialize reports as clientInfo.version. It is
// read from the build rather than copied from cmd/codey, which keeps this
// package from holding a second copy of the version string; "(devel)" is
// what a go build outside a module release gets, and is the truth.
var clientVersion = func() string {
	if bi, ok := debug.ReadBuildInfo(); ok && bi.Main.Version != "" {
		return bi.Main.Version
	}
	return "unknown"
}()

// stderrTail is how much of a server's stderr is kept for the error a
// failed start reports. A server that will not initialize usually says why
// on stderr — a missing API key, a module it cannot find — and without it
// the user is left with "timed out" and nowhere to look. Four kilobytes is
// the end of a stack trace, which is the part that names the cause.
const stderrTail = 4 << 10

// Connect starts every server in cfg.MCPServers, discovers its tools, and
// returns them ready to register. Servers are started together rather than
// one after another, so five broken servers cost one handshake timeout and
// not five; the tools come back sorted by name so the registry's order, and
// with it the prompt cache, does not depend on which server answered first.
//
// Errors are per server and per tool, never fatal: a server that is not
// installed, will not start, or speaks a protocol this client does not is
// reported and its tools are simply not there. The closer is always
// returned, even when nothing started, so the caller has one thing to defer.
func Connect(ctx context.Context, cfg config.Config) ([]tools.Tool, io.Closer, []error) {
	names := make([]string, 0, len(cfg.MCPServers))
	for name := range cfg.MCPServers {
		names = append(names, name)
	}
	sort.Strings(names)

	results := make([]started, len(names))
	var wg sync.WaitGroup
	for i, name := range names {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c, infos, err := start(ctx, name, cfg.MCPServers[name])
			results[i] = started{name: name, c: c, tools: infos, err: err}
		}()
	}
	wg.Wait()
	return assemble(results)
}

// started is what starting one server came to: its client, what it listed,
// or why it did neither. The client is set even on failure so it can be
// closed; nil means the process never started.
type started struct {
	name  string
	c     *client
	tools []toolInfo
	err   error
}

// assemble turns started servers into registrable tools. It is separate from
// Connect so the tests can feed it in-process fakes; Connect's only other
// work is starting processes.
func assemble(results []started) ([]tools.Tool, *servers, []error) {
	all := &servers{}
	var out []tools.Tool
	var errs []error
	// Registration panics on a duplicate, and two servers whose sanitized
	// names collide — "my db" and "my_db" — would produce one. The first
	// keeps the name; the second's tool is dropped and reported.
	seen := map[string]bool{}
	for _, r := range results {
		if r.err != nil {
			errs = append(errs, fmt.Errorf("mcp: %s: %w", r.name, r.err))
			if r.c != nil {
				_ = r.c.Close()
			}
			continue
		}
		kept := 0
		for _, info := range r.tools {
			t, err := newTool(r.c, r.name, info)
			if err != nil {
				errs = append(errs, fmt.Errorf("mcp: %s: %w", r.name, err))
				continue
			}
			if seen[t.Name()] {
				errs = append(errs, fmt.Errorf("mcp: %s: tool %q would register as %q, which another tool already has", r.name, info.Name, t.Name()))
				continue
			}
			seen[t.Name()] = true
			out = append(out, t)
			kept++
		}
		if kept == 0 {
			// A server with nothing to offer is a process with nothing to do;
			// it is stopped now rather than carried to the end of the session.
			errs = append(errs, fmt.Errorf("mcp: %s: server offers no usable tools", r.name))
			_ = r.c.Close()
			continue
		}
		all.clients = append(all.clients, r.c)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name() < out[j].Name() })
	return out, all, errs
}

// start runs one server through to a listed set of tools. The client is
// returned even on failure so the caller can close it; a nil client means
// the process never started.
func start(ctx context.Context, name string, spec config.MCPServer) (*client, []toolInfo, error) {
	command := strings.TrimSpace(spec.Command)
	if command == "" {
		return nil, nil, fmt.Errorf("no command configured")
	}
	cmd := exec.Command(command, spec.Args...)
	cmd.Env = append(os.Environ(), envList(spec.Env)...)
	// A new process group, so a server started through npx or uvx can be
	// killed whole. Killing only the wrapper leaves the real server holding
	// the pipes, which is how the bash tool learned the same lesson.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	// WaitDelay bounds Wait when a grandchild keeps stderr open after the
	// server has been killed; without it Close could hang on a pipe.
	cmd.WaitDelay = 2 * time.Second
	stderr := &tail{max: stderrTail}
	cmd.Stderr = stderr
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, nil, fmt.Errorf("starting %s: %w", command, err)
	}
	var waitOnce sync.Once
	c := newClient(name, stdin, stdout,
		func() {
			if cmd.Process != nil {
				_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			}
		},
		func() { waitOnce.Do(func() { _ = cmd.Wait() }) })

	ctx, cancel := context.WithTimeout(ctx, initTimeout)
	defer cancel()
	infos, err := handshake(ctx, c)
	if err != nil {
		if s := strings.TrimSpace(stderr.String()); s != "" {
			err = fmt.Errorf("%w\nserver stderr:\n%s", err, s)
		}
		return c, nil, err
	}
	return c, infos, nil
}

// handshake is initialize, initialized and tools/list under one deadline. It
// is separate from start so the tests can drive it over pipes.
func handshake(ctx context.Context, c *client) ([]toolInfo, error) {
	if err := c.initialize(ctx); err != nil {
		return nil, err
	}
	return c.listTools(ctx)
}

// envList is env as KEY=VALUE pairs, sorted so two runs with the same config
// give the server the same environment in the same order.
func envList(env map[string]string) []string {
	out := make([]string, 0, len(env))
	for k, v := range env {
		out = append(out, k+"="+v)
	}
	sort.Strings(out)
	return out
}

// servers is the closer Connect returns: every client that made it through
// the handshake, closed together.
type servers struct {
	mu      sync.Mutex
	clients []*client
}

// Close kills every server and its process group. Safe to call more than
// once; the second call finds nothing to do.
func (s *servers) Close() error {
	s.mu.Lock()
	clients := s.clients
	s.clients = nil
	s.mu.Unlock()
	var wg sync.WaitGroup
	for _, c := range clients {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_ = c.Close()
		}()
	}
	wg.Wait()
	return nil
}

// tail keeps the last max bytes written to it. It is what a server's stderr
// goes into: enough to say why it failed, bounded so a chatty server that
// runs all session cannot grow the harness.
type tail struct {
	mu  sync.Mutex
	max int
	buf []byte
}

func (t *tail) Write(p []byte) (int, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.buf = append(t.buf, p...)
	if len(t.buf) > t.max {
		t.buf = t.buf[len(t.buf)-t.max:]
	}
	return len(p), nil
}

func (t *tail) String() string {
	t.mu.Lock()
	defer t.mu.Unlock()
	return string(t.buf)
}