package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"sync"
	"sync/atomic"
)

// StdioConfig configures a subprocess-backed stdio MCP transport.
type StdioConfig struct {
	Command string
	Args    []string
	Env     []string // additional environment (KEY=VALUE), appended to os env
	Dir     string
	// Stderr, if non-nil, receives the child process's stderr.
	Stderr io.Writer
}

// streamTransport speaks newline-delimited JSON-RPC over any reader/writer
// pair, correlating responses with requests by id. It backs both the
// subprocess stdio transport and in-process pipe-connected transports.
type streamTransport struct {
	w   io.WriteCloser
	enc *json.Encoder
	r   io.Reader

	// process control (nil for non-subprocess transports)
	proc   interface{ Kill() error }
	waitFn func() error

	mu      sync.Mutex
	counter int64
	pending map[string]chan *Message
	closed  bool
	exitErr error

	notify atomic.Pointer[func(*Message)]
	done   chan struct{}
}

func newStreamTransport(r io.Reader, w io.WriteCloser) *streamTransport {
	t := &streamTransport{
		w:       w,
		enc:     json.NewEncoder(w),
		r:       r,
		pending: make(map[string]chan *Message),
		done:    make(chan struct{}),
	}
	go t.readLoop()
	return t
}

// NewStreamTransport creates a transport over an existing reader/writer pair,
// e.g. an in-process pipe. It performs no process management.
func NewStreamTransport(r io.Reader, w io.WriteCloser) Transport {
	return newStreamTransport(r, w)
}

// NewStdioTransport spawns the configured command and serves MCP over its
// stdio. The process is terminated by Close.
func NewStdioTransport(cfg StdioConfig) (Transport, error) {
	if cfg.Command == "" {
		return nil, errors.New("stdio transport: empty command")
	}
	cmd := exec.Command(cfg.Command, cfg.Args...)
	cmd.Dir = cfg.Dir
	if len(cfg.Env) > 0 {
		cmd.Env = append(cmd.Environ(), cfg.Env...)
	}
	cmd.Stderr = cfg.Stderr

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdio transport: stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdio transport: stdout pipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("stdio transport: start %q: %w", cfg.Command, err)
	}

	t := &streamTransport{
		w:       stdin,
		enc:     json.NewEncoder(stdin),
		r:       stdout,
		waitFn:  cmd.Wait,
		pending: make(map[string]chan *Message),
		done:    make(chan struct{}),
	}
	if cmd.Process != nil {
		t.proc = cmd.Process
	}
	go t.readLoop()
	return t, nil
}

func (t *streamTransport) SetNotificationHandler(h func(*Message)) {
	if h == nil {
		t.notify.Store(nil)
		return
	}
	t.notify.Store(&h)
}

func (t *streamTransport) readLoop() {
	scanner := bufio.NewScanner(t.r)
	// MCP messages can be large (tool schemas); allow up to 16 MiB lines.
	scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
	for scanner.Scan() {
		line := scanner.Bytes()
		if len(bytes.TrimSpace(line)) == 0 {
			continue
		}
		var m Message
		if err := json.Unmarshal(line, &m); err != nil {
			continue // ignore non-JSON noise
		}
		if m.IsResponse() {
			key := normalizeID(m.ID)
			t.mu.Lock()
			ch := t.pending[key]
			delete(t.pending, key)
			t.mu.Unlock()
			if ch != nil {
				ch <- &m
			}
			continue
		}
		if m.Method != "" {
			if hp := t.notify.Load(); hp != nil {
				(*hp)(&m)
			}
		}
	}
	// input closed: peer is gone. Fail all in-flight calls.
	var exitErr error
	if t.waitFn != nil {
		exitErr = t.waitFn()
	}
	t.mu.Lock()
	t.closed = true
	t.exitErr = exitErr
	pending := t.pending
	t.pending = map[string]chan *Message{}
	t.mu.Unlock()
	for _, ch := range pending {
		close(ch)
	}
	close(t.done)
}

func (t *streamTransport) Call(ctx context.Context, method string, params any) (*Message, error) {
	id := atomic.AddInt64(&t.counter, 1)
	idRaw := json.RawMessage(strconv.FormatInt(id, 10))
	req, err := NewRequest(idRaw, method, params)
	if err != nil {
		return nil, err
	}
	key := normalizeID(idRaw)
	ch := make(chan *Message, 1)

	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return nil, t.closedErr()
	}
	t.pending[key] = ch
	err = t.enc.Encode(req)
	t.mu.Unlock()
	if err != nil {
		t.mu.Lock()
		delete(t.pending, key)
		t.mu.Unlock()
		return nil, fmt.Errorf("stdio write %s: %w", method, err)
	}

	select {
	case <-ctx.Done():
		t.mu.Lock()
		delete(t.pending, key)
		t.mu.Unlock()
		return nil, ctx.Err()
	case resp, ok := <-ch:
		if !ok {
			return nil, t.closedErr()
		}
		return resp, nil
	case <-t.done:
		return nil, t.closedErr()
	}
}

func (t *streamTransport) Notify(_ context.Context, method string, params any) error {
	note, err := NewNotification(method, params)
	if err != nil {
		return err
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.closed {
		return t.closedErr()
	}
	return t.enc.Encode(note)
}

func (t *streamTransport) closedErr() error {
	if t.exitErr != nil {
		return fmt.Errorf("mcp stdio transport closed: %w", t.exitErr)
	}
	return errors.New("mcp stdio transport closed")
}

func (t *streamTransport) Close() error {
	t.mu.Lock()
	closed := t.closed
	t.mu.Unlock()
	_ = t.w.Close()
	if t.proc != nil {
		_ = t.proc.Kill()
	}
	if !closed {
		<-t.done
	}
	return nil
}

// normalizeID produces a stable map key from a JSON-RPC id, tolerating
// whitespace differences in how peers echo the id.
func normalizeID(raw json.RawMessage) string {
	return string(bytes.TrimSpace(raw))
}