package lsp

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// Server represents a running LSP server process.
type Server struct {
	adapter *Adapter
	root    string
	cmd     *exec.Cmd
	stdin   io.WriteCloser
	stdout  io.ReadCloser

	mu      sync.Mutex
	reqID   atomic.Int64
	pending map[int64]chan json.RawMessage
	diags   map[string]json.RawMessage // uri → diagnostics
	open    map[string]bool

	done chan struct{}
}

// StartServer launches an LSP server and performs the initialize handshake.
func StartServer(adapter *Adapter, root string) (*Server, error) {
	cmd := exec.Command(adapter.Cmd[0], adapter.Cmd[1:]...)
	cmd.Dir = root
	cmd.Stderr = nil

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		stdin.Close()
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	s := &Server{
		adapter: adapter,
		root:    root,
		cmd:     cmd,
		stdin:   stdin,
		stdout:  stdout,
		pending: make(map[int64]chan json.RawMessage),
		diags:   make(map[string]json.RawMessage),
		open:    make(map[string]bool),
		done:    make(chan struct{}),
	}

	go s.readLoop()

	if err := s.initialize(); err != nil {
		s.Stop()
		return nil, err
	}
	return s, nil
}

func (s *Server) initialize() error {
	initOpts := s.adapter.InitOptions
	params := map[string]any{
		"processId": os.Getpid(),
		"rootUri":   FileURI(s.root),
		"rootPath":  s.root,
		"capabilities": map[string]any{
			"textDocument": map[string]any{
				"definition":         map[string]any{"linkSupport": true},
				"references":         map[string]any{},
				"hover":              map[string]any{"contentFormat": []string{"plaintext", "markdown"}},
				"rename":             map[string]any{"prepareSupport": true},
				"documentSymbol":     map[string]any{"hierarchicalDocumentSymbolSupport": true},
				"completion":         map[string]any{"completionItem": map[string]any{"snippetSupport": false}},
				"publishDiagnostics": map[string]any{},
			},
			"workspace": map[string]any{
				"workspaceFolders": true,
				"symbol":           map[string]any{},
			},
		},
		"workspaceFolders": []map[string]any{
			{"uri": FileURI(s.root), "name": filepath.Base(s.root)},
		},
	}
	if initOpts != nil {
		params["initializationOptions"] = initOpts
	}

	_, err := s.Request("initialize", params, 30*time.Second)
	if err != nil {
		return fmt.Errorf("initialize: %w", err)
	}
	s.Notify("initialized", map[string]any{})
	return nil
}

// Stop shuts down the LSP server gracefully.
func (s *Server) Stop() {
	s.Request("shutdown", map[string]any{}, 5*time.Second) //nolint:errcheck
	s.Notify("exit", nil)
	s.stdin.Close()
	timer := time.AfterFunc(5*time.Second, func() { s.cmd.Process.Kill() })
	s.cmd.Wait()
	timer.Stop()
	<-s.done
}

// EnsureOpen sends textDocument/didOpen if the file hasn't been opened yet.
func (s *Server) EnsureOpen(file string) {
	s.mu.Lock()
	if s.open[file] {
		s.mu.Unlock()
		return
	}
	s.open[file] = true
	s.mu.Unlock()

	text, _ := os.ReadFile(file)
	s.Notify("textDocument/didOpen", map[string]any{
		"textDocument": map[string]any{
			"uri":        FileURI(file),
			"languageId": LanguageID(file),
			"version":    1,
			"text":       string(text),
		},
	})
	// Brief pause to let server index
	time.Sleep(100 * time.Millisecond)
}

// Request sends a JSON-RPC request and waits for the response.
func (s *Server) Request(method string, params any, timeout time.Duration) (json.RawMessage, error) {
	id := s.reqID.Add(1)
	ch := make(chan json.RawMessage, 1)

	s.mu.Lock()
	s.pending[id] = ch
	s.mu.Unlock()

	msg := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"method":  method,
		"params":  params,
	}
	if err := s.send(msg); err != nil {
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
		return nil, err
	}

	select {
	case result := <-ch:
		return result, nil
	case <-time.After(timeout):
		s.mu.Lock()
		delete(s.pending, id)
		s.mu.Unlock()
		return nil, fmt.Errorf("%s: timeout after %v", method, timeout)
	}
}

// Notify sends a JSON-RPC notification (no response expected).
func (s *Server) Notify(method string, params any) {
	msg := map[string]any{
		"jsonrpc": "2.0",
		"method":  method,
	}
	if params != nil {
		msg["params"] = params
	}
	s.send(msg) //nolint:errcheck
}

// Diagnostics returns cached diagnostics for a file URI.
func (s *Server) Diagnostics(uri string) json.RawMessage {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.diags[uri]
}

func (s *Server) send(msg map[string]any) error {
	body, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	header := fmt.Sprintf("Content-Length: %d\r\n\r\n", len(body))
	s.mu.Lock()
	defer s.mu.Unlock()
	_, err = s.stdin.Write([]byte(header))
	if err != nil {
		return err
	}
	_, err = s.stdin.Write(body)
	return err
}

func (s *Server) readLoop() {
	defer close(s.done)
	reader := bufio.NewReader(s.stdout)
	for {
		// Read headers
		contentLength := 0
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				return
			}
			line = strings.TrimSpace(line)
			if line == "" {
				break
			}
			if strings.HasPrefix(line, "Content-Length:") {
				n, _ := strconv.Atoi(strings.TrimSpace(line[len("Content-Length:"):]))
				contentLength = n
			}
		}
		if contentLength == 0 {
			continue
		}

		body := make([]byte, contentLength)
		if _, err := io.ReadFull(reader, body); err != nil {
			return
		}

		var msg struct {
			ID     *json.RawMessage `json:"id"`
			Method string           `json:"method"`
			Result json.RawMessage  `json:"result"`
			Error  *struct {
				Message string `json:"message"`
			} `json:"error"`
			Params json.RawMessage `json:"params"`
		}
		if err := json.Unmarshal(body, &msg); err != nil {
			continue
		}

		if msg.ID != nil && msg.Method == "" {
			// Response
			var id int64
			json.Unmarshal(*msg.ID, &id)
			s.mu.Lock()
			ch := s.pending[id]
			delete(s.pending, id)
			s.mu.Unlock()
			if ch != nil {
				if msg.Error != nil {
					ch <- nil
				} else {
					ch <- msg.Result
				}
			}
		} else if msg.Method == "textDocument/publishDiagnostics" {
			// Cache diagnostics
			var params struct {
				URI         string          `json:"uri"`
				Diagnostics json.RawMessage `json:"diagnostics"`
			}
			if json.Unmarshal(msg.Params, &params) == nil {
				s.mu.Lock()
				s.diags[params.URI] = params.Diagnostics
				s.mu.Unlock()
			}
		}
	}
}