package transport

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sync"

	"github.com/CalmoAI/mcp-server-signoz/pkg/mcp"
)

// StdioTransport implements MCP transport over stdin/stdout
type StdioTransport struct {
	stdin   io.Reader
	stdout  io.Writer
	scanner *bufio.Scanner
	writer  *bufio.Writer
	mu      sync.Mutex
	closed  bool
}

// NewStdioTransport creates a new STDIO transport
func NewStdioTransport() *StdioTransport {
	return &StdioTransport{
		stdin:   os.Stdin,
		stdout:  os.Stdout,
		scanner: bufio.NewScanner(os.Stdin),
		writer:  bufio.NewWriter(os.Stdout),
	}
}

// NewStdioTransportWithStreams creates a new STDIO transport with custom streams
func NewStdioTransportWithStreams(stdin io.Reader, stdout io.Writer) *StdioTransport {
	return &StdioTransport{
		stdin:   stdin,
		stdout:  stdout,
		scanner: bufio.NewScanner(stdin),
		writer:  bufio.NewWriter(stdout),
	}
}

// Send implements the Transport interface
func (t *StdioTransport) Send(message *mcp.MCPMessage) error {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed {
		return fmt.Errorf("transport is closed")
	}

	// Serialize the message to JSON
	data, err := json.Marshal(message)
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}

	// Write the JSON message followed by a newline
	if _, err := t.writer.Write(data); err != nil {
		return fmt.Errorf("failed to write message: %w", err)
	}

	if err := t.writer.WriteByte('\n'); err != nil {
		return fmt.Errorf("failed to write newline: %w", err)
	}

	// Flush the buffer
	if err := t.writer.Flush(); err != nil {
		return fmt.Errorf("failed to flush writer: %w", err)
	}

	return nil
}

// Receive implements the Transport interface
func (t *StdioTransport) Receive() (*mcp.MCPMessage, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed {
		return nil, fmt.Errorf("transport is closed")
	}

	// Read lines until we get a non-empty line or EOF
	var line []byte
	for {
		if !t.scanner.Scan() {
			if err := t.scanner.Err(); err != nil {
				return nil, fmt.Errorf("failed to read from stdin: %w", err)
			}
			return nil, io.EOF
		}

		line = t.scanner.Bytes()
		// Skip empty lines and continue reading
		if len(line) > 0 {
			break
		}
		// Continue loop to read next line if current line is empty
	}

	// Parse the JSON message
	var message mcp.MCPMessage
	if err := json.Unmarshal(line, &message); err != nil {
		return nil, fmt.Errorf("failed to unmarshal message: %w", err)
	}

	// Validate that it's a valid JSON-RPC 2.0 message
	if message.JSONRPC != "2.0" {
		return nil, fmt.Errorf("invalid JSON-RPC version: %s", message.JSONRPC)
	}

	return &message, nil
}

// Close implements the Transport interface
func (t *StdioTransport) Close() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed {
		return nil
	}

	t.closed = true

	// Flush any remaining data
	if t.writer != nil {
		_ = t.writer.Flush()
	}

	return nil
}

// IsClosed returns whether the transport is closed
func (t *StdioTransport) IsClosed() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.closed
}

// ReceiveWithContext receives a message with context cancellation support
func (t *StdioTransport) ReceiveWithContext(ctx context.Context) (*mcp.MCPMessage, error) {
	type result struct {
		message *mcp.MCPMessage
		err     error
	}

	resultCh := make(chan result, 1)

	go func() {
		message, err := t.Receive()
		resultCh <- result{message: message, err: err}
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case res := <-resultCh:
		return res.message, res.err
	}
}

// SendWithContext sends a message with context cancellation support
func (t *StdioTransport) SendWithContext(ctx context.Context, message *mcp.MCPMessage) error {
	type result struct {
		err error
	}

	resultCh := make(chan result, 1)

	go func() {
		err := t.Send(message)
		resultCh <- result{err: err}
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case res := <-resultCh:
		return res.err
	}
}

// ValidateMessage validates an MCP message structure
func ValidateMessage(message *mcp.MCPMessage) error {
	if message == nil {
		return fmt.Errorf("message is nil")
	}

	if message.JSONRPC != "2.0" {
		return fmt.Errorf("invalid JSON-RPC version: %s", message.JSONRPC)
	}

	// Check if it's a request, response, or notification
	if message.ID != nil {
		// Request or response
		if message.Method != "" && message.Result == nil && message.Error == nil {
			// Request
			return nil
		} else if message.Method == "" && (message.Result != nil || message.Error != nil) {
			// Response
			return nil
		}
		return fmt.Errorf("invalid request/response message structure")
	} else {
		// Notification
		if message.Method == "" {
			return fmt.Errorf("notification must have a method")
		}
		if message.Result != nil || message.Error != nil {
			return fmt.Errorf("notification cannot have result or error")
		}
		return nil
	}
}

// MessageBuilder provides a fluent interface for building MCP messages
type MessageBuilder struct {
	message *mcp.MCPMessage
}

// NewMessageBuilder creates a new message builder
func NewMessageBuilder() *MessageBuilder {
	return &MessageBuilder{
		message: &mcp.MCPMessage{
			JSONRPC: "2.0",
		},
	}
}

// Request sets up the message as a request
func (b *MessageBuilder) Request(id int64, method string) *MessageBuilder {
	b.message.ID = &id
	b.message.Method = method
	return b
}

// Response sets up the message as a response
func (b *MessageBuilder) Response(id int64) *MessageBuilder {
	b.message.ID = &id
	return b
}

// Notification sets up the message as a notification
func (b *MessageBuilder) Notification(method string) *MessageBuilder {
	b.message.Method = method
	return b
}

// WithParams sets the parameters for the message
func (b *MessageBuilder) WithParams(params interface{}) *MessageBuilder {
	if params != nil {
		data, err := json.Marshal(params)
		if err == nil {
			b.message.Params = data
		}
	}
	return b
}

// WithResult sets the result for the message
func (b *MessageBuilder) WithResult(result interface{}) *MessageBuilder {
	if result != nil {
		data, err := json.Marshal(result)
		if err == nil {
			b.message.Result = data
		}
	}
	return b
}

// WithError sets the error for the message
func (b *MessageBuilder) WithError(code int, message string, data interface{}) *MessageBuilder {
	mcpError := &mcp.MCPError{
		Code:    code,
		Message: message,
	}
	if data != nil {
		errorData, err := json.Marshal(data)
		if err == nil {
			mcpError.Data = errorData
		}
	}
	b.message.Error = mcpError
	return b
}

// Build returns the constructed message
func (b *MessageBuilder) Build() *mcp.MCPMessage {
	return b.message
}

// BuildAndValidate returns the constructed message after validation
func (b *MessageBuilder) BuildAndValidate() (*mcp.MCPMessage, error) {
	message := b.Build()
	if err := ValidateMessage(message); err != nil {
		return nil, err
	}
	return message, nil
}
