package mcpcheck

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
)

const defaultCheckTimeout = 10 * time.Second

type rpcRequest struct {
	JSONRPC string `json:"jsonrpc"`
	ID      any    `json:"id,omitempty"`
	Method  string `json:"method"`
	Params  any    `json:"params,omitempty"`
}

type rpcResponse struct {
	JSONRPC string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id"`
	Result  json.RawMessage  `json:"result,omitempty"`
	Error   *rpcError        `json:"error,omitempty"`
}

type rpcHTTPResponse struct {
	Response  rpcResponse
	SessionID string
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type toolsListResult struct {
	Tools []profile.MCPToolSummary `json:"tools"`
}

// Check validates and probes a configured external MCP server.
func Check(ctx context.Context, server profile.MCPServer) profile.MCPServerCheck {
	ctx, cancel := withDefaultTimeout(ctx)
	defer cancel()

	if err := server.Validate(); err != nil {
		return failure("validation", err.Error())
	}

	var result profile.MCPServerCheck
	switch server.Transport {
	case profile.MCPServerTransportHTTP:
		result = checkHTTP(ctx, server)
	case profile.MCPServerTransportStdio:
		result = checkStdio(ctx, server)
	default:
		result = failure("validation", fmt.Sprintf("unsupported MCP transport %q", server.Transport))
	}
	result.CheckedAt = time.Now().UTC()
	return result
}

func withDefaultTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
	if _, ok := ctx.Deadline(); ok {
		return context.WithCancel(ctx)
	}
	return context.WithTimeout(ctx, defaultCheckTimeout)
}

func checkHTTP(ctx context.Context, server profile.MCPServer) profile.MCPServerCheck {
	headers, err := resolveHTTPHeaders(server.HTTP.Headers)
	if err != nil {
		return failure("validation", err.Error())
	}
	client := &http.Client{}
	initResponse, err := postRPC(ctx, client, server.HTTP.URL, headers, "", rpcRequest{
		JSONRPC: "2.0",
		ID:      1,
		Method:  "initialize",
		Params: map[string]any{
			"protocolVersion": "2024-11-05",
			"capabilities":    map[string]any{},
			"clientInfo": map[string]string{
				"name":    "agentteam-desktop",
				"version": "0.0.0",
			},
		},
	})
	if err != nil {
		return classifyHTTPError(err)
	}
	sessionID := initResponse.SessionID
	if err := postRPCNotification(ctx, client, server.HTTP.URL, headers, sessionID, rpcRequest{
		JSONRPC: "2.0",
		Method:  "notifications/initialized",
	}); err != nil {
		return classifyHTTPError(err)
	}

	response, err := postRPC(ctx, client, server.HTTP.URL, headers, sessionID, rpcRequest{
		JSONRPC: "2.0",
		ID:      2,
		Method:  "tools/list",
		Params:  map[string]any{},
	})
	if err != nil {
		return classifyHTTPError(err)
	}
	tools, err := decodeToolsList(response.Response.Result)
	if err != nil {
		return failure("tools", err.Error())
	}
	return success(tools)
}

func resolveHTTPHeaders(headers []string) (map[string]string, error) {
	result := make(map[string]string, len(headers))
	for _, entry := range headers {
		name, value, ok := strings.Cut(entry, ":")
		if !ok {
			return nil, fmt.Errorf("header %q must contain ':'", entry)
		}
		name = strings.TrimSpace(name)
		value = strings.TrimSpace(profile.ExpandRuntimeEnv(value))
		if value == "" {
			return nil, fmt.Errorf("header %q value is not configured", name)
		}
		result[name] = value
	}
	return result, nil
}

type httpStatusError struct {
	status int
	body   string
}

func (e httpStatusError) Error() string {
	body := strings.TrimSpace(e.body)
	if body == "" {
		return fmt.Sprintf("HTTP %d", e.status)
	}
	return fmt.Sprintf("HTTP %d: %s", e.status, body)
}

func postRPC(ctx context.Context, client *http.Client, endpoint string, headers map[string]string, sessionID string, request rpcRequest) (rpcHTTPResponse, error) {
	payload, err := json.Marshal(request)
	if err != nil {
		return rpcHTTPResponse{}, err
	}
	req, err := newRPCHTTPRequest(ctx, endpoint, headers, sessionID, payload)
	if err != nil {
		return rpcHTTPResponse{}, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return rpcHTTPResponse{}, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return rpcHTTPResponse{}, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return rpcHTTPResponse{}, httpStatusError{status: resp.StatusCode, body: string(body)}
	}
	var decoded rpcResponse
	if err := json.Unmarshal(body, &decoded); err != nil {
		return rpcHTTPResponse{}, fmt.Errorf("decode JSON-RPC response: %w", err)
	}
	if decoded.Error != nil {
		return rpcHTTPResponse{}, fmt.Errorf("JSON-RPC error %d: %s", decoded.Error.Code, decoded.Error.Message)
	}
	if len(decoded.Result) == 0 {
		return rpcHTTPResponse{}, fmt.Errorf("JSON-RPC response missing result")
	}
	return rpcHTTPResponse{
		Response:  decoded,
		SessionID: resp.Header.Get("Mcp-Session-Id"),
	}, nil
}

func postRPCNotification(ctx context.Context, client *http.Client, endpoint string, headers map[string]string, sessionID string, request rpcRequest) error {
	payload, err := json.Marshal(request)
	if err != nil {
		return err
	}
	req, err := newRPCHTTPRequest(ctx, endpoint, headers, sessionID, payload)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return httpStatusError{status: resp.StatusCode, body: string(body)}
	}
	return nil
}

func newRPCHTTPRequest(ctx context.Context, endpoint string, headers map[string]string, sessionID string, payload []byte) (*http.Request, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	if strings.TrimSpace(sessionID) != "" {
		req.Header.Set("Mcp-Session-Id", strings.TrimSpace(sessionID))
	}
	for name, value := range headers {
		req.Header.Set(name, value)
	}
	return req, nil
}

func classifyHTTPError(err error) profile.MCPServerCheck {
	var statusErr httpStatusError
	if asHTTPStatusError(err, &statusErr) {
		stage := "http_status"
		if statusErr.status == http.StatusUnauthorized || statusErr.status == http.StatusForbidden {
			stage = "auth"
		}
		return failure(stage, statusErr.Error())
	}
	message := err.Error()
	if strings.Contains(message, "JSON-RPC") || strings.Contains(message, "decode") || strings.Contains(message, "missing result") {
		return failure("protocol", message)
	}
	return failure("network", message)
}

func asHTTPStatusError(err error, target *httpStatusError) bool {
	if typed, ok := err.(httpStatusError); ok {
		*target = typed
		return true
	}
	return false
}

func checkStdio(ctx context.Context, server profile.MCPServer) profile.MCPServerCheck {
	command, err := resolveCommand(server.Stdio.Command, server.Stdio.CWD)
	if err != nil {
		return failure("validation", err.Error())
	}
	cmd := exec.CommandContext(ctx, command, server.Stdio.Args...)
	if server.Stdio.CWD != "" {
		cmd.Dir = server.Stdio.CWD
	}
	cmd.Env = append(runtimeconfig.RuntimeProcessEnv(), server.Stdio.ResolvedEnv()...)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return failure("stdio", fmt.Sprintf("open stdin: %v", err))
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return failure("stdio", fmt.Sprintf("open stdout: %v", err))
	}
	var stderr limitedBuffer
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		return failure("stdio", fmt.Sprintf("start command: %v", err))
	}
	defer func() {
		_ = stdin.Close()
		if cmd.Process != nil {
			_ = cmd.Process.Kill()
		}
		_ = cmd.Wait()
	}()

	reader := bufio.NewReader(stdout)
	if err := writeRPCLine(stdin, rpcRequest{
		JSONRPC: "2.0",
		ID:      1,
		Method:  "initialize",
		Params: map[string]any{
			"protocolVersion": "2024-11-05",
			"capabilities":    map[string]any{},
			"clientInfo": map[string]string{
				"name":    "agentteam-desktop",
				"version": "0.0.0",
			},
		},
	}); err != nil {
		return failure("stdio", fmt.Sprintf("write initialize: %v", err))
	}
	if _, err := readRPCResponse(ctx, reader, "1"); err != nil {
		return stdioReadFailure(err, stderr.String())
	}
	_ = writeRPCLine(stdin, rpcRequest{JSONRPC: "2.0", Method: "notifications/initialized"})
	if err := writeRPCLine(stdin, rpcRequest{JSONRPC: "2.0", ID: 2, Method: "tools/list", Params: map[string]any{}}); err != nil {
		return failure("stdio", fmt.Sprintf("write tools/list: %v", err))
	}
	response, err := readRPCResponse(ctx, reader, "2")
	if err != nil {
		return stdioReadFailure(err, stderr.String())
	}
	tools, err := decodeToolsList(response.Result)
	if err != nil {
		return failure("tools", err.Error())
	}
	return success(tools)
}

func resolveCommand(command, cwd string) (string, error) {
	if strings.ContainsRune(command, os.PathSeparator) || filepath.IsAbs(command) {
		path := command
		if !filepath.IsAbs(path) && cwd != "" {
			path = filepath.Join(cwd, path)
		}
		info, err := os.Stat(path)
		if err != nil {
			return "", fmt.Errorf("stdio command is not accessible: %w", err)
		}
		if info.IsDir() {
			return "", fmt.Errorf("stdio command must be a file")
		}
		return command, nil
	}
	resolver := runtimeconfig.DefaultRuntimeResolver()
	path, err := resolver.ResolveBinary(command, nil)
	if err != nil {
		return "", fmt.Errorf("stdio command %q not found on PATH", command)
	}
	return path, nil
}

func writeRPCLine(w io.Writer, request rpcRequest) error {
	payload, err := json.Marshal(request)
	if err != nil {
		return err
	}
	payload = append(payload, '\n')
	_, err = w.Write(payload)
	return err
}

func readRPCResponse(ctx context.Context, reader *bufio.Reader, expectedID string) (rpcResponse, error) {
	for {
		lineCh := make(chan struct {
			line []byte
			err  error
		}, 1)
		go func() {
			line, err := reader.ReadBytes('\n')
			lineCh <- struct {
				line []byte
				err  error
			}{line: line, err: err}
		}()
		select {
		case <-ctx.Done():
			return rpcResponse{}, ctx.Err()
		case item := <-lineCh:
			if item.err != nil {
				return rpcResponse{}, item.err
			}
			var decoded rpcResponse
			if err := json.Unmarshal(bytes.TrimSpace(item.line), &decoded); err != nil {
				return rpcResponse{}, fmt.Errorf("decode JSON-RPC response: %w", err)
			}
			if decoded.ID == nil || rpcIDString(*decoded.ID) != expectedID {
				continue
			}
			if decoded.Error != nil {
				return rpcResponse{}, fmt.Errorf("JSON-RPC error %d: %s", decoded.Error.Code, decoded.Error.Message)
			}
			if len(decoded.Result) == 0 {
				return rpcResponse{}, fmt.Errorf("JSON-RPC response missing result")
			}
			return decoded, nil
		}
	}
}

func rpcIDString(raw json.RawMessage) string {
	return strings.Trim(strings.TrimSpace(string(raw)), "\"")
}

func stdioReadFailure(err error, stderr string) profile.MCPServerCheck {
	message := err.Error()
	if trimmed := strings.TrimSpace(stderr); trimmed != "" {
		message += ": " + trimmed
	}
	if err == context.DeadlineExceeded || err == context.Canceled {
		return failure("timeout", message)
	}
	if strings.Contains(message, "JSON-RPC") || strings.Contains(message, "decode") || strings.Contains(message, "missing result") {
		return failure("protocol", message)
	}
	return failure("stdio", message)
}

func decodeToolsList(raw json.RawMessage) ([]profile.MCPToolSummary, error) {
	var result toolsListResult
	if err := json.Unmarshal(raw, &result); err != nil {
		return nil, fmt.Errorf("decode tools/list result: %w", err)
	}
	if result.Tools == nil {
		return nil, fmt.Errorf("tools/list result missing tools array")
	}
	return result.Tools, nil
}

func success(tools []profile.MCPToolSummary) profile.MCPServerCheck {
	return profile.MCPServerCheck{
		OK:      true,
		Stage:   "tools",
		Message: "MCP server check passed",
		Tools:   tools,
	}
}

func failure(stage string, message string) profile.MCPServerCheck {
	return profile.MCPServerCheck{
		OK:      false,
		Stage:   stage,
		Message: message,
	}
}

type limitedBuffer struct {
	buf bytes.Buffer
}

func (b *limitedBuffer) Write(p []byte) (int, error) {
	const limit = 4096
	remaining := limit - b.buf.Len()
	if remaining > 0 {
		if len(p) > remaining {
			_, _ = b.buf.Write(p[:remaining])
		} else {
			_, _ = b.buf.Write(p)
		}
	}
	return len(p), nil
}

func (b *limitedBuffer) String() string {
	return b.buf.String()
}