package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/alex-ilgayev/mcpspy/pkg/ebpf"
	mcpevents "github.com/alex-ilgayev/mcpspy/pkg/event"
	"github.com/alex-ilgayev/mcpspy/pkg/http"
	"github.com/alex-ilgayev/mcpspy/pkg/mcp"
	"github.com/alex-ilgayev/mcpspy/pkg/namespace"
	"github.com/alex-ilgayev/mcpspy/pkg/output"
	"github.com/alex-ilgayev/mcpspy/pkg/version"

	mcpspydebug "github.com/alex-ilgayev/mcpspy/pkg/debug"
)

// Command line flags
var (
	showBuffers bool
	verbose     bool
	outputFile  string
	logLevel    string
)

func main() {
	rootCmd := &cobra.Command{
		Use:   "mcpspy",
		Short: "Monitor Model Context Protocol communication",
		Long: `MCPSpy is a CLI utility that uses eBPF to monitor MCP (Model Context Protocol) 
communication by tracking stdio operations and analyzing JSON-RPC 2.0 messages.`,
		Version:      fmt.Sprintf("%s (commit: %s, built: %s)", version.Version, version.Commit, version.Date),
		RunE:         run,
		SilenceUsage: true,
	}

	// Add flags
	rootCmd.Flags().BoolVarP(&showBuffers, "buffers", "b", false, "Show raw message buffers")
	rootCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose logging (debug level)")
	rootCmd.Flags().StringVarP(&outputFile, "output", "o", "", "Output file (JSONL format will be written to file)")
	rootCmd.Flags().StringVarP(&logLevel, "log-level", "l", "info", "Set log level (trace, debug, info, warn, error, fatal, panic)")

	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func run(cmd *cobra.Command, args []string) error {
	// Set up logging
	// Handle verbose flag as shortcut for debug level
	if verbose {
		logLevel = "debug"
	}

	// Parse and set log level
	level, err := logrus.ParseLevel(logLevel)
	if err != nil {
		return fmt.Errorf("invalid log level '%s': %w", logLevel, err)
	}
	logrus.SetLevel(level)

	// Setup trace pipe to debug eBPF programs if debug or trace level
	if level >= logrus.DebugLevel {
		go mcpspydebug.PrintTracePipe(logrus.StandardLogger())
	}

	// Fetch current mount namespace
	mountNS, err := namespace.GetCurrentMountNamespace()
	if err != nil {
		return fmt.Errorf("failed to get current mount namespace: %w", err)
	}
	logrus.WithField("mount_ns", mountNS).Debug("Current mount namespace")

	// Set up console display (always show console output)
	consoleDisplay := output.NewConsoleDisplay(os.Stdout, showBuffers)
	consoleDisplay.PrintHeader()

	// Set up file output if specified
	var fileDisplay output.OutputHandler

	if outputFile != "" {
		file, err := os.Create(outputFile)
		if err != nil {
			return fmt.Errorf("failed to create output file '%s': %w", outputFile, err)
		}
		fileDisplay = output.NewJSONLDisplay(file)
		defer func() {
			if err := file.Close(); err != nil {
				logrus.WithError(err).Error("Failed to close output file")
			}
		}()
	}

	// Create and load eBPF program
	loader, err := ebpf.New(level >= logrus.DebugLevel)
	if err != nil {
		return fmt.Errorf("failed to create eBPF loader: %w", err)
	}
	defer loader.Close()

	// Process library events
	// and creates uprobe hooks for dynamically loaded libraries
	libManager := ebpf.NewLibraryManager(loader, mountNS)
	defer libManager.Close()

	// Manage HTTP sessions (1.1/2/chunked encoding/SSE)
	httpManager := http.NewSessionManager()
	defer httpManager.Close()

	consoleDisplay.PrintInfo("Loading eBPF programs...")
	if err := loader.Load(); err != nil {
		return fmt.Errorf("failed to load eBPF programs: %w", err)
	}

	// Set up signal handling
	ctx, cancel := context.WithCancel(context.Background())
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-sigChan
		cancel()
	}()

	// Start event processing
	if err := loader.Start(ctx); err != nil {
		return fmt.Errorf("failed to start event processing: %w", err)
	}

	// Enumerate all libraries for TLS inspection
	logrus.Debug("Doing initial enumeration of libraries for TLS inspection")
	if err := loader.RunIterLibEnum(); err != nil {
		return fmt.Errorf("failed to enumerate libraries: %w", err)
	}

	consoleDisplay.PrintInfo("Monitoring MCP communication... Press Ctrl+C to stop")
	consoleDisplay.PrintInfo("")

	// Create MCP parser and statistics
	parser := mcp.NewParser()
	stats := make(map[string]int)

	// Main event loop
	for {
		select {
		case <-ctx.Done():
			consoleDisplay.PrintStats(stats)
			return nil
		case event, ok := <-httpManager.HTTPEvents():
			if !ok {
				// Channel closed, exit
				consoleDisplay.PrintStats(stats)
				return nil
			}

			var allMessages []*mcp.Message

			switch event.Type() {
			case mcpevents.EventTypeHttpRequest:
				httpEvt := event.(*mcpevents.HttpRequestEvent)

				logrus.WithFields(logrus.Fields{
					"method": httpEvt.Method,
					"host":   httpEvt.Host,
					"path":   httpEvt.Path,
				}).Trace("HTTP request event")

				// Parse request payload for MCP messages
				requestMessages, err := parser.ParseDataHttp(httpEvt.RequestPayload, httpEvt.EventType, httpEvt.PID, httpEvt.Comm(), httpEvt.Host, true)
				if err != nil {
					logrus.WithError(err).Debugf("Failed to process %s event", httpEvt.Type())
				} else {
					allMessages = append(allMessages, requestMessages...)
				}

			case mcpevents.EventTypeHttpResponse:
				httpEvt := event.(*mcpevents.HttpResponseEvent)

				logrus.WithFields(logrus.Fields{
					"method":     httpEvt.Method,
					"host":       httpEvt.Host,
					"path":       httpEvt.Path,
					"code":       httpEvt.Code,
					"is_chunked": httpEvt.IsChunked,
				}).Trace("HTTP response event")

				// Parse response payload for MCP messages
				responseMessages, err := parser.ParseDataHttp(httpEvt.ResponsePayload, httpEvt.EventType, httpEvt.PID, httpEvt.Comm(), httpEvt.Host, false)
				if err != nil {
					logrus.WithError(err).Debugf("Failed to process %s event", httpEvt.Type())
				} else {
					allMessages = append(allMessages, responseMessages...)
				}

			case mcpevents.EventTypeHttpSSE:
				sseEvent := event.(*mcpevents.SSEEvent)

				logrus.WithFields(logrus.Fields{
					"method":    sseEvent.Method,
					"host":      sseEvent.Host,
					"path":      sseEvent.Path,
					"sse_event": sseEvent.SSEEventType,
				}).Trace("HTTP SSE event")

				// Parse SSE data as MCP messages
				messages, err := parser.ParseDataHttp(sseEvent.Data, sseEvent.EventType, sseEvent.PID, sseEvent.Comm(), sseEvent.Host, false)
				if err != nil {
					logrus.WithError(err).Debugf("Failed to process %s event", sseEvent.Type())
				} else {
					allMessages = append(allMessages, messages...)
				}
			}

			// Update statistics and display messages
			for _, msg := range allMessages {
				if msg.Method != "" {
					stats[msg.Method]++
				}
			}

			// Display messages to console
			if len(allMessages) > 0 {
				consoleDisplay.PrintMessages(allMessages)

				// Also write to file if specified
				if fileDisplay != nil {
					fileDisplay.PrintMessages(allMessages)
				}
			}
		case event, ok := <-loader.Events():
			if !ok {
				// Channel closed, exit
				consoleDisplay.PrintStats(stats)
				return nil
			}

			// Handle different event types
			switch e := event.(type) {
			case *mcpevents.FSDataEvent:
				buf := e.Buf[:e.BufSize]
				if len(buf) == 0 {
					continue
				}

				// Parse raw eBPF event data into MCP messages
				messages, err := parser.ParseDataStdio(buf, e.EventType, e.FromPID, e.FromCommStr(), e.ToPID, e.ToCommStr())
				if err != nil {
					logrus.WithError(err).Debugf("Failed to process %s event", e.Type())
					continue
				}

				// Update statistics
				for _, msg := range messages {
					if msg.Method != "" {
						stats[msg.Method]++
					}
				}

				// Display messages to console
				consoleDisplay.PrintMessages(messages)

				// Also write to file if specified
				if fileDisplay != nil {
					fileDisplay.PrintMessages(messages)
				}
			case *mcpevents.LibraryEvent:
				// Handle library events
				logrus.WithFields(logrus.Fields{
					"pid":     e.PID,
					"comm":    e.Comm(),
					"path":    e.Path(),
					"inode":   e.Inode,
					"mountNS": e.MntNSID,
				}).Trace("Library loaded")

				if err := libManager.ProcessLibraryEvent(e); err != nil {
					logrus.WithError(err).WithField("path", e.Path()).Warnf("Failed to process %s event", e.Type())
				}
			case *mcpevents.TlsPayloadEvent:
				// Handle TLS payloads
				logrus.WithFields(logrus.Fields{
					"type":     e.Type(),
					"pid":      e.PID,
					"comm":     e.Comm(),
					"size":     e.Size,
					"buf_size": e.BufSize,
					"version":  e.HttpVersion,
				}).Trace("TLS payload event")
				// Raw TLS event, we need to aggregate it into HTTP sessions
				if err := httpManager.ProcessTlsEvent(e); err != nil {
					logrus.WithError(err).Warnf("Failed to process %s event", e.Type())
				}
			case *mcpevents.TlsFreeEvent:
				// Handle TLS free event
				logrus.WithFields(logrus.Fields{
					"pid":     e.PID,
					"comm":    e.Comm(),
					"ssl_ctx": e.SSLContext,
				}).Trace("TLS free event")

				if err := httpManager.ProcessTlsFreeEvent(e); err != nil {
					logrus.WithError(err).Warnf("Failed to process %s event", e.Type())
				}
			default:
				logrus.WithField("type", event.Type()).Warn("Unknown event type")
			}
		}
	}
}
