package diag

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

// idleTimeout is how long a language server may sit with no file to check
// before it is shut down. Five minutes is long enough that a burst of edits
// shares one server and short enough that a session which has moved on to
// another language is not carrying gopls's memory around for the rest of the
// afternoon.
const idleTimeout = 5 * time.Minute

// initTimeout is how long a server has to answer initialize. Every server in
// the table answers it in well under a second and does its real loading
// afterwards, so a server that has not answered in thirty seconds is one that
// never will — a broken shim, a binary for another architecture — and every
// check against it would otherwise pay the full checkTimeout for nothing.
const initTimeout = 30 * time.Second

// shutdownGrace is how long Close waits for a server to answer shutdown and
// then to exit before it is killed. A server that has not managed either in
// half a second is not going to.
const shutdownGrace = 500 * time.Millisecond

// Limits on what a server may make the client hold. A server that publishes
// for thousands of files or thousands of problems in one file is one whose
// output the model could not use anyway, and holding all of it would let a
// misbehaving server grow the harness without bound.
const (
	maxDiagnosticFiles   = 50
	maxDiagnosticsPerURI = 200
)

// errServerDead is what a check gets from a server that has exited or been
// closed. The Runner drops such a server from its table so the next check
// starts a fresh one.
var errServerDead = errors.New("language server is not running")

// session is one running language server and the client side of its
// connection. It lives across edits: starting gopls costs seconds, and paying
// that per edit would put every check past its deadline.
type session struct {
	name string
	spec serverSpec
	root string

	// stdin is written under writeMu so requests from Check and replies from
	// the reader never interleave.
	writeMu sync.Mutex
	stdin   io.WriteCloser
	// kill ends the process; wait reaps it. Both are no-ops for the in-process
	// fake the tests use.
	kill func()
	wait func()

	// ready is closed once initialize has been answered; initErr says whether
	// it went well. dead is closed when the connection is gone for any reason.
	ready   chan struct{}
	initErr error
	dead    chan struct{}
	dieOnce sync.Once
	onExit  func(*session)

	// busy is a one-slot semaphore: one check at a time per server, because
	// a collector is attached to the whole connection and two checks would
	// steal each other's batches. It is a channel rather than a mutex so a
	// waiting check can give up when its context does.
	busy chan struct{}

	mu        sync.Mutex
	nextID    int64
	pending   map[int64]chan rpcResponse
	versions  map[string]int
	collector *collector
	idle      *time.Timer
	idleD     time.Duration
	closed    bool

	// caps is what the server said it can do in its initialize result, keyed
	// by capability name, so a query the server never offered is refused here
	// rather than sent and answered with an error whose wording varies by
	// server. Written once before ready is closed and read only after.
	caps map[string]json.RawMessage
}

// rpcResponse is the part of a response the client uses.
type rpcResponse struct {
	Result json.RawMessage
	Error  *rpcError
}

// rpcError is a JSON-RPC error object.
type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func (e *rpcError) Error() string { return fmt.Sprintf("%s (%d)", e.Message, e.Code) }

// startServer spawns bin in root and begins initializing it. It returns nil if
// the process could not be started. Initialization continues in the
// background; Check waits for it under the caller's deadline, so a slow server
// costs one missed check rather than a hung one.
func startServer(bin string, spec serverSpec, root string, idle time.Duration, onExit func(*session)) *session {
	cmd := exec.Command(bin, spec.Args...)
	cmd.Dir = root
	// stderr is where servers log. It is discarded rather than captured
	// because gopls alone writes megabytes of it over a session, and nobody
	// reads a language server's log from inside a coding agent.
	cmd.Stderr = nil
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil
	}
	if err := cmd.Start(); err != nil {
		return nil
	}
	var waitOnce sync.Once
	s := newSession(spec, root, stdin, stdout,
		func() { _ = cmd.Process.Kill() },
		func() { waitOnce.Do(func() { _ = cmd.Wait() }) },
		idle, onExit)
	return s
}

// newSession wires a session over an already-open connection and starts the
// reader and the initialize handshake. It is the seam the tests use to speak
// to an in-process fake server through pipes.
func newSession(spec serverSpec, root string, stdin io.WriteCloser, stdout io.Reader, kill, wait func(), idle time.Duration, onExit func(*session)) *session {
	if kill == nil {
		kill = func() {}
	}
	if wait == nil {
		wait = func() {}
	}
	if onExit == nil {
		onExit = func(*session) {}
	}
	s := &session{
		name:     filepath.Base(spec.Bin),
		spec:     spec,
		root:     root,
		stdin:    stdin,
		kill:     kill,
		wait:     wait,
		ready:    make(chan struct{}),
		dead:     make(chan struct{}),
		onExit:   onExit,
		busy:     make(chan struct{}, 1),
		pending:  map[int64]chan rpcResponse{},
		versions: map[string]int{},
		idleD:    idle,
	}
	if idle > 0 {
		s.idle = time.AfterFunc(idle, s.idleClose)
	}
	go s.readLoop(bufio.NewReaderSize(stdout, 64<<10))
	go s.initialize()
	return s
}

// Dead reports whether the connection is gone.
func (s *session) Dead() bool {
	select {
	case <-s.dead:
		return true
	default:
		return false
	}
}

// acquire takes the server's one request slot and then waits for the server
// to be ready, giving up when ctx does or the server dies. The slot is taken
// first so that a request arriving while the server is still initializing
// queues behind it rather than piling up several waiters that all wake at
// once. The returned release must be called exactly once.
func (s *session) acquire(ctx context.Context) (release func(), err error) {
	select {
	case s.busy <- struct{}{}:
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-s.dead:
		return nil, errServerDead
	}
	release = func() {
		s.touch()
		<-s.busy
	}
	s.touch()

	select {
	case <-s.ready:
	case <-ctx.Done():
		release()
		return nil, ctx.Err()
	case <-s.dead:
		release()
		return nil, errServerDead
	}
	if s.initErr != nil {
		release()
		return nil, s.initErr
	}
	return release, nil
}

// Check sends the file to the server and collects what it publishes about it.
func (s *session) Check(ctx context.Context, path string) ([]Diagnostic, error) {
	release, err := s.acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer release()

	content, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	uri := fileURI(path)
	c := newCollector(uri)
	s.setCollector(c)
	defer s.setCollector(nil)

	if err := s.sendDocument(uri, path, string(content)); err != nil {
		return nil, err
	}

	// First batch for the changed file, then a short window for the rest —
	// see settleWindow. If the window is cut short by the deadline what has
	// arrived so far is still returned; it is real.
	select {
	case <-c.got:
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-s.dead:
		return nil, errServerDead
	}
	settle := time.NewTimer(settleWindow)
	defer settle.Stop()
	select {
	case <-settle.C:
	case <-ctx.Done():
	case <-s.dead:
	}
	return c.result(s.name), nil
}

// Close asks the server to stop and makes sure it has. Safe to call more than
// once and concurrently with a check, which will see errServerDead.
func (s *session) Close() {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	if s.idle != nil {
		s.idle.Stop()
	}
	s.mu.Unlock()

	if !s.Dead() {
		// The polite sequence, each step bounded. Skipping straight to kill
		// would work, but gopls left uncleanly leaves stale state in its
		// cache directory that slows its next start.
		ctx, cancel := context.WithTimeout(context.Background(), shutdownGrace)
		_, _ = s.call(ctx, "shutdown", nil)
		cancel()
		_ = s.notify("exit", nil)
		_ = s.stdin.Close()
		select {
		case <-s.dead:
		case <-time.After(shutdownGrace):
		}
	}
	s.die(errServerDead)
	s.wait()
}

// idleClose is the idle timer's callback. A check that is in flight keeps
// the server: the timer is pushed back rather than the check interrupted. The
// check slot is held through Close so a check arriving during shutdown waits
// and then finds the server dead, rather than talking to it halfway down.
func (s *session) idleClose() {
	select {
	case s.busy <- struct{}{}:
	default:
		s.touch()
		return
	}
	defer func() { <-s.busy }()
	s.Close()
}

// failedToStart reports whether the server died or was closed before it had
// finished initializing. The Runner uses it to stop retrying a server that
// is installed but does not work.
func (s *session) failedToStart() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.initErr != nil
}

// touch pushes the idle deadline back to now plus the idle timeout.
func (s *session) touch() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.idle != nil && !s.closed {
		s.idle.Reset(s.idleD)
	}
}

// die marks the connection gone, fails every request still waiting on it,
// and ends the process. Everything that can end a session funnels through
// here so it happens once.
func (s *session) die(err error) {
	s.dieOnce.Do(func() {
		s.mu.Lock()
		if s.initErr == nil && !s.isReady() {
			s.initErr = err
		}
		pending := s.pending
		s.pending = map[int64]chan rpcResponse{}
		if s.idle != nil {
			s.idle.Stop()
		}
		s.mu.Unlock()
		for _, ch := range pending {
			close(ch)
		}
		close(s.dead)
		s.readyOnce()
		s.kill()
		_ = s.stdin.Close()
		go func() {
			s.wait()
			s.onExit(s)
		}()
	})
}

func (s *session) isReady() bool {
	select {
	case <-s.ready:
		return true
	default:
		return false
	}
}

// readyOnce closes ready if it is not already closed.
func (s *session) readyOnce() {
	select {
	case <-s.ready:
	default:
		close(s.ready)
	}
}

// initialize runs the handshake. It has no deadline of its own: a server
// still loading a large module after the first check's two seconds is a
// server worth waiting for, and the checks are what carry the deadline. If
// the process dies the reader ends it.
func (s *session) initialize() {
	rootURI := fileURI(s.root)
	params := map[string]any{
		"processId": os.Getpid(),
		"clientInfo": map[string]any{
			"name": "codey",
		},
		"rootUri":  rootURI,
		"rootPath": s.root,
		"workspaceFolders": []map[string]any{{
			"uri":  rootURI,
			"name": filepath.Base(s.root),
		}},
		"capabilities": map[string]any{
			"textDocument": map[string]any{
				"publishDiagnostics": map[string]any{
					"relatedInformation": false,
				},
				"synchronization": map[string]any{
					"dynamicRegistration": false,
					"didSave":             true,
				},
				// Markdown first: it is what every server in the table
				// produces natively, and its code fences are how the
				// signature is told apart from the doc comment when the
				// hover is trimmed for the model.
				"hover": map[string]any{
					"contentFormat": []string{"markdown", "plaintext"},
				},
				// With hierarchy a struct's methods and fields come nested
				// under it, which is what lets the outline name the
				// container; without it gopls sends a flat list of top-level
				// names only.
				"documentSymbol": map[string]any{
					"hierarchicalDocumentSymbolSupport": true,
				},
			},
			"workspace": map[string]any{
				"workspaceFolders": true,
				"configuration":    false,
			},
			"window": map[string]any{
				"workDoneProgress": false,
			},
		},
		"initializationOptions": map[string]any{},
	}
	ctx, cancel := context.WithTimeout(context.Background(), initTimeout)
	defer cancel()
	result, err := s.call(ctx, "initialize", params)
	if err != nil {
		s.mu.Lock()
		s.initErr = fmt.Errorf("%s: initialize: %w", s.name, err)
		s.mu.Unlock()
		s.readyOnce()
		s.die(s.initErr)
		return
	}
	var init struct {
		Capabilities map[string]json.RawMessage `json:"capabilities"`
	}
	// A result that does not parse is not fatal: the server is up and can
	// still check files, which is the more important job. It just gets no
	// queries.
	_ = json.Unmarshal(result, &init)
	s.caps = init.Capabilities
	if err := s.notify("initialized", map[string]any{}); err != nil {
		s.mu.Lock()
		s.initErr = fmt.Errorf("%s: initialized: %w", s.name, err)
		s.mu.Unlock()
		s.readyOnce()
		s.die(s.initErr)
		return
	}
	s.readyOnce()
}

// sendDocument opens the file on the server or, if it is already open,
// replaces its content. Documents stay open across checks: a didClose after
// every check would have the server clear and recompute the file's
// diagnostics twice per edit, and some servers publish the clear late enough
// to be mistaken for the next check's answer. A didSave follows because the
// file genuinely is on disk, and rust-analyzer runs its cargo check only on
// save.
func (s *session) sendDocument(uri, path, content string) error {
	s.mu.Lock()
	version, open := s.versions[uri]
	version++
	s.versions[uri] = version
	s.mu.Unlock()

	var err error
	if open {
		err = s.notify("textDocument/didChange", map[string]any{
			"textDocument": map[string]any{"uri": uri, "version": version},
			"contentChanges": []map[string]any{
				{"text": content},
			},
		})
	} else {
		err = s.notify("textDocument/didOpen", map[string]any{
			"textDocument": map[string]any{
				"uri":        uri,
				"languageId": s.spec.languageID(path),
				"version":    version,
				"text":       content,
			},
		})
	}
	if err != nil {
		return err
	}
	return s.notify("textDocument/didSave", map[string]any{
		"textDocument": map[string]any{"uri": uri},
	})
}

// call sends a request and waits for its response or for ctx.
func (s *session) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
	s.mu.Lock()
	s.nextID++
	id := s.nextID
	ch := make(chan rpcResponse, 1)
	if s.Dead() {
		s.mu.Unlock()
		return nil, errServerDead
	}
	s.pending[id] = ch
	s.mu.Unlock()

	msg := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
	if params != nil {
		msg["params"] = params
	}
	if err := s.write(msg); err != nil {
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
		return nil, err
	}
	select {
	case resp, ok := <-ch:
		if !ok {
			return nil, errServerDead
		}
		if resp.Error != nil {
			return nil, resp.Error
		}
		return resp.Result, nil
	case <-ctx.Done():
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
		return nil, ctx.Err()
	}
}

// notify sends a notification: a request with no id and no reply.
func (s *session) notify(method string, params any) error {
	msg := map[string]any{"jsonrpc": "2.0", "method": method}
	if params != nil {
		msg["params"] = params
	}
	return s.write(msg)
}

// reply answers a request the server made of us. It writes from its own
// goroutine rather than from the read loop that called it: a server that
// sends several requests in a row is blocked writing the second while we are
// blocked writing the answer to the first, and with both sides waiting on
// the other neither pipe drains. The OS pipe buffer hides this until a
// message is larger than it, which is exactly when it would be found in
// production and not before.
func (s *session) reply(id json.RawMessage, result any, err *rpcError) {
	msg := map[string]any{"jsonrpc": "2.0", "id": id}
	if err != nil {
		msg["error"] = err
	} else {
		msg["result"] = result
	}
	go func() { _ = s.write(msg) }()
}

func (s *session) write(msg any) error {
	if s.Dead() {
		return errServerDead
	}
	body, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	if err := writeMessage(s.stdin, body); err != nil {
		// A write that fails is a pipe with nobody on the other end.
		s.die(fmt.Errorf("%s: write: %w", s.name, err))
		return errServerDead
	}
	return nil
}

// incoming is the superset of the three JSON-RPC message shapes; which one
// it is depends on which fields are set.
type incoming struct {
	ID     json.RawMessage `json:"id"`
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
	Result json.RawMessage `json:"result"`
	Error  *rpcError       `json:"error"`
}

// readLoop reads messages until the connection fails, dispatching responses
// to their callers and diagnostics to the collector. Any framing or JSON
// error ends the session: a stream that has lost sync cannot be recovered
// and a server we cannot understand is a server to restart.
func (s *session) readLoop(r *bufio.Reader) {
	for {
		body, err := readMessage(r)
		if err != nil {
			if err == io.EOF {
				s.die(fmt.Errorf("%s exited", s.name))
			} else {
				s.die(fmt.Errorf("%s: %w", s.name, err))
			}
			return
		}
		var msg incoming
		if err := json.Unmarshal(body, &msg); err != nil {
			s.die(fmt.Errorf("%s sent invalid JSON: %w", s.name, err))
			return
		}
		switch {
		case msg.Method == "" && len(msg.ID) > 0:
			s.dispatchResponse(msg)
		case msg.Method != "" && len(msg.ID) > 0 && string(msg.ID) != "null":
			s.handleRequest(msg)
		case msg.Method != "":
			s.handleNotification(msg)
		}
	}
}

func (s *session) dispatchResponse(msg incoming) {
	var id int64
	if err := json.Unmarshal(msg.ID, &id); err != nil {
		return
	}
	s.mu.Lock()
	ch, ok := s.pending[id]
	delete(s.pending, id)
	s.mu.Unlock()
	if ok {
		ch <- rpcResponse{Result: msg.Result, Error: msg.Error}
	}
}

// handleRequest answers what servers ask of a client. Each gets the emptiest
// legal answer: this client has no UI to show progress in, no settings to
// hand over, and registers no capabilities dynamically. Not answering at all
// is not an option — gopls waits on workspace/configuration before it will
// type-check anything.
func (s *session) handleRequest(msg incoming) {
	switch msg.Method {
	case "workspace/configuration":
		var p struct {
			Items []json.RawMessage `json:"items"`
		}
		_ = json.Unmarshal(msg.Params, &p)
		s.reply(msg.ID, make([]any, len(p.Items)), nil)
	case "client/registerCapability", "client/unregisterCapability",
		"window/workDoneProgress/create", "workspace/applyEdit":
		s.reply(msg.ID, nil, nil)
	case "window/showMessageRequest":
		s.reply(msg.ID, nil, nil)
	case "workspace/workspaceFolders":
		s.reply(msg.ID, []map[string]any{{"uri": fileURI(s.root), "name": filepath.Base(s.root)}}, nil)
	default:
		s.reply(msg.ID, nil, &rpcError{Code: -32601, Message: "method not found"})
	}
}

func (s *session) handleNotification(msg incoming) {
	if msg.Method != "textDocument/publishDiagnostics" {
		return
	}
	var p struct {
		URI         string `json:"uri"`
		Diagnostics []struct {
			Range struct {
				Start struct {
					Line      int `json:"line"`
					Character int `json:"character"`
				} `json:"start"`
			} `json:"range"`
			Severity int    `json:"severity"`
			Message  string `json:"message"`
		} `json:"diagnostics"`
	}
	if err := json.Unmarshal(msg.Params, &p); err != nil {
		return
	}
	path := uriPath(p.URI)
	if path == "" {
		return
	}
	ds := make([]Diagnostic, 0, min(len(p.Diagnostics), maxDiagnosticsPerURI))
	for _, d := range p.Diagnostics {
		if len(ds) == maxDiagnosticsPerURI {
			break
		}
		sev := Severity(d.Severity)
		if sev == 0 {
			// The spec lets a server omit severity, and says the client
			// should then interpret it as it likes. An unqualified
			// complaint from a compiler is an error.
			sev = SeverityError
		}
		ds = append(ds, Diagnostic{
			File:     path,
			Line:     d.Range.Start.Line + 1,
			Col:      d.Range.Start.Character + 1,
			Severity: sev,
			Message:  d.Message,
		})
	}
	s.mu.Lock()
	c := s.collector
	s.mu.Unlock()
	if c != nil {
		c.add(p.URI, path, ds)
	}
}

func (s *session) setCollector(c *collector) {
	s.mu.Lock()
	s.collector = c
	s.mu.Unlock()
}

// collector gathers publishDiagnostics batches during one check. The latest
// batch per file wins, because a server that publishes twice for a file is
// replacing, not adding.
type collector struct {
	target     string
	targetPath string
	got        chan struct{}
	mu         sync.Mutex
	byPath     map[string][]Diagnostic
}

func newCollector(target string) *collector {
	return &collector{
		target:     target,
		targetPath: uriPath(target),
		got:        make(chan struct{}, 1),
		byPath:     map[string][]Diagnostic{},
	}
}

// add records a batch. The file being checked is always kept; neighbours are
// kept until the file cap is reached, so a server that publishes for the
// whole workspace cannot push out the one file the model asked about.
func (c *collector) add(uri, path string, ds []Diagnostic) {
	c.mu.Lock()
	_, seen := c.byPath[path]
	if seen || path == c.targetPath || len(c.byPath) < maxDiagnosticFiles {
		c.byPath[path] = ds
	}
	c.mu.Unlock()
	if sameURI(uri, c.target) {
		select {
		case c.got <- struct{}{}:
		default:
		}
	}
}

// result flattens what was collected and stamps it with the server's name.
func (c *collector) result(source string) []Diagnostic {
	c.mu.Lock()
	defer c.mu.Unlock()
	var out []Diagnostic
	for _, ds := range c.byPath {
		for _, d := range ds {
			d.Source = source
			out = append(out, d)
		}
	}
	return out
}

// fileURI is path as a file: URI, with the characters a URI cannot carry
// escaped. url.URL does the escaping; building the string by hand would break
// on the first path with a space in it.
func fileURI(path string) string {
	path = filepath.ToSlash(path)
	if runtime.GOOS == "windows" && !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	u := url.URL{Scheme: "file", Path: path}
	return u.String()
}

// uriPath is the inverse of fileURI, or "" for a URI that is not a file.
func uriPath(uri string) string {
	u, err := url.Parse(uri)
	if err != nil || u.Scheme != "file" {
		return ""
	}
	p := u.Path
	if runtime.GOOS == "windows" && len(p) > 2 && p[0] == '/' && p[2] == ':' {
		p = p[1:]
	}
	return filepath.Clean(filepath.FromSlash(p))
}

// sameURI compares two file URIs by the path they name, so a server that
// percent-encodes differently from us still matches.
func sameURI(a, b string) bool {
	if a == b {
		return true
	}
	pa, pb := uriPath(a), uriPath(b)
	return pa != "" && pa == pb
}