package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/creachadair/jrpc2"
	"github.com/creachadair/jrpc2/channel"
	"github.com/goccy/go-yaml"

	"mcp-bridge/src/logger"
)

// Version information set by build flags
var (
	Version    = "dev"
	BuildTime  = "unknown"
	CommitHash = "unknown"
)

type ServiceConfig struct {
	Config  string   `yaml:"config"`
	Command []string `yaml:"command"`
}

type EndpointConfig struct {
	Path      string `yaml:"path"`
	McpMethod string `yaml:"mcp_method"`
	ToolName  string `yaml:"tool_name"`
}

type ServiceDetails struct {
	ServiceName string           `yaml:"serviceName"`
	Endpoints   []EndpointConfig `yaml:"endpoints"`
}

type Config struct {
	Port     int                      `yaml:"port"`
	Services map[string]ServiceConfig `yaml:"services"`
}

// loadConfig loads the configuration from the config file
func loadConfig(configFile string) (*Config, error) {
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, fmt.Errorf("error reading %s: %v", configFile, err)
	}

	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling %s: %v", configFile, err)
	}

	return &config, nil
}

// loadServiceConfig loads the service configuration from the config file
func loadServiceConfig(configFile string) (*ServiceDetails, error) {
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, fmt.Errorf("error reading %s: %v", configFile, err)
	}

	var serviceDetails ServiceDetails
	err = yaml.Unmarshal(data, &serviceDetails)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling %s: %v", configFile, err)
	}

	return &serviceDetails, nil
}

func main() {
	// Define command-line flags
	port := flag.Int("port", 0, "port to listen on")
	verbosity := flag.String("v", "info", "log level (debug, info, warn, error)")
	configFile := flag.String("config", "config.yaml", "path to the configuration file")
	showVersion := flag.Bool("version", false, "show version information and exit")
	flag.Parse()

	// Show version information if requested
	if *showVersion {
		fmt.Printf("mcp-bridge version %s\n", Version)
		fmt.Printf("Build time: %s\n", BuildTime)
		fmt.Printf("Commit: %s\n", CommitHash)
		return
	}

	// Initialize the logger
	logger.Initialize(*verbosity)

	// Log version information
	logger.WithFields(logger.Fields{
		"version":    Version,
		"build_time": BuildTime,
		"commit":     CommitHash,
	}).Info("Starting mcp-bridge")

	// Store child processes and a WaitGroup to wait for them
	childProcesses := make(map[string]*exec.Cmd)
	jrpcClients := make(map[string]*jrpc2.Client)
	var wg sync.WaitGroup

	// Create a channel for SIGHUP signal
	sighupCh := make(chan os.Signal, 1)

	// Create a context that can be canceled for reload operations
	ctx, cancelCtx := context.WithCancel(context.Background())
	defer cancelCtx()

	// Defer the cleanup of child processes. This will run when main exits.
	defer func() {
		logger.Info("Closing JSON-RPC clients...")
		for name, client := range jrpcClients {
			if err := client.Close(); err != nil {
				logger.WithFields(logger.Fields{
					"service": name,
					"error":   err,
				}).Error("Error closing client")
			}
		}

		logger.Info("Shutting down adapters...")
		for name, cmd := range childProcesses {
			if cmd.Process == nil {
				continue
			}
			logger.WithFields(logger.Fields{
				"service": name,
				"pid":     cmd.Process.Pid,
			}).Info("Sending SIGTERM to process group")
			if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM); err != nil {
				logger.WithFields(logger.Fields{
					"service": name,
					"pid":     cmd.Process.Pid,
					"error":   err,
				}).Error("Failed to send SIGTERM to process group")
			}
		}

		done := make(chan struct{})
		go func() {
			wg.Wait()
			close(done)
		}()

		select {
		case <-done:
			logger.Info("All adapter processes shut down gracefully.")
		case <-time.After(2 * time.Second):
			logger.Warn("Timed out waiting for adapter processes to exit. Forcing shutdown...")
			for name, cmd := range childProcesses {
				if cmd.ProcessState == nil || !cmd.ProcessState.Exited() {
					logger.WithFields(logger.Fields{
						"service": name,
						"pid":     cmd.Process.Pid,
					}).Info("Sending SIGKILL to process group")
					if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
						logger.WithFields(logger.Fields{
							"service": name,
							"pid":     cmd.Process.Pid,
							"error":   err,
						}).Error("Failed to kill process group")
					}
				}
			}
			wg.Wait() // Wait again after killing
			logger.Info("All adapter processes shut down.")
		}
	}()

	// Load the initial configuration
	config, err := loadConfig(*configFile)
	if err != nil {
		logger.WithFields(logger.Fields{
			"error": err,
		}).Fatal("Error loading configuration")
		return
	}

	// Override port with the flag if it was provided
	if *port != 0 {
		config.Port = *port
	}

	logger.WithFields(logger.Fields{
		"port": config.Port,
	}).Info("Configuration loaded")

	// Function to initialize or reload services
	initializeServices := func(config *Config) error {
		// Clear existing HTTP handlers
		http.DefaultServeMux = http.NewServeMux()

		// Initialize services, start child processes, and set up HTTP handlers
		for name, service := range config.Services {
			logger.WithFields(logger.Fields{
				"service":     name,
				"config_file": service.Config,
				"command":     service.Command,
			}).Info("Initializing service")

			// Check if the service is already running
			if cmd, exists := childProcesses[name]; exists && cmd.Process != nil {
				logger.WithFields(logger.Fields{
					"service": name,
					"pid":     cmd.Process.Pid,
				}).Info("Service already running, reloading configuration")

				// Close the existing JSON-RPC client
				if client, ok := jrpcClients[name]; ok {
					if err := client.Close(); err != nil {
						logger.WithFields(logger.Fields{
							"service": name,
							"error":   err,
						}).Error("Error closing client during reload")
					}
					delete(jrpcClients, name)
				}

				// Send SIGTERM to the process group
				if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM); err != nil {
					logger.WithFields(logger.Fields{
						"service": name,
						"pid":     cmd.Process.Pid,
						"error":   err,
					}).Error("Failed to send SIGTERM to process group during reload")
				}

				// Wait for the process to exit
				done := make(chan struct{})
				go func() {
					cmd.Wait()
					close(done)
				}()

				select {
				case <-done:
					logger.WithFields(logger.Fields{
						"service": name,
					}).Info("Service shut down gracefully during reload")
				case <-time.After(2 * time.Second):
					logger.WithFields(logger.Fields{
						"service": name,
						"pid":     cmd.Process.Pid,
					}).Warn("Timed out waiting for service to exit during reload, forcing shutdown")
					if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
						logger.WithFields(logger.Fields{
							"service": name,
							"pid":     cmd.Process.Pid,
							"error":   err,
						}).Error("Failed to kill process group during reload")
					}
					cmd.Wait()
				}

				delete(childProcesses, name)
			}

			// Start the MCP adapter as a child process
			if len(service.Command) > 0 {
				cmd := exec.Command(service.Command[0], service.Command[1:]...)
				cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
				stdin, err := cmd.StdinPipe()
				if err != nil {
					logger.WithFields(logger.Fields{
						"service": name,
						"error":   err,
					}).Error("Error getting stdin pipe")
					continue
				}
				stdout, err := cmd.StdoutPipe()
				if err != nil {
					logger.WithFields(logger.Fields{
						"service": name,
						"error":   err,
					}).Error("Error getting stdout pipe")
					continue
				}

				cmd.Stderr = os.Stderr
				err = cmd.Start()
				if err != nil {
					logger.WithFields(logger.Fields{
						"service": name,
						"error":   err,
					}).Error("Error starting service")
					continue // Skip this service if it fails to start
				}

				logger.WithFields(logger.Fields{
					"service": name,
					"pid":     cmd.Process.Pid,
				}).Info("Service started")
				childProcesses[name] = cmd
				wg.Add(1)

				// Set up the JSON-RPC client.
				// The client will run in the background and stop automatically when its
				// underlying channel is closed.
				opts := &jrpc2.ClientOptions{
					OnStop: func(_ *jrpc2.Client, err error) {
						if err != nil {
							logger.WithFields(logger.Fields{
								"service": name,
								"error":   err,
							}).Error("JSON-RPC client disconnected with error")
						} else {
							logger.WithFields(logger.Fields{
								"service": name,
							}).Info("JSON-RPC client disconnected gracefully")
						}
					},
				}
				ch := channel.Line(stdout, stdin)
				jrpcClients[name] = jrpc2.NewClient(ch, opts)

				// Goroutine to wait for the process to exit
				go func(name string, cmd *exec.Cmd) {
					defer wg.Done()
					if err := cmd.Wait(); err != nil {
						logger.WithFields(logger.Fields{
							"service": name,
							"error":   err,
						}).Error("Service exited with error")
					} else {
						logger.WithFields(logger.Fields{
							"service": name,
						}).Info("Service exited gracefully")
					}
				}(name, cmd)
			}

			// Read the service config file
			serviceDetails, err := loadServiceConfig(service.Config)
			if err != nil {
				logger.WithFields(logger.Fields{
					"service":     name,
					"config_file": service.Config,
					"error":       err,
				}).Error("Error loading service config")
				continue
			}

			logger.WithFields(logger.Fields{
				"service":      name,
				"service_name": serviceDetails.ServiceName,
			}).Info("Service details loaded")

			for _, endpoint := range serviceDetails.Endpoints {
				logger.WithFields(logger.Fields{
					"service":    name,
					"path":       endpoint.Path,
					"mcp_method": endpoint.McpMethod,
					"tool_name":  endpoint.ToolName,
				}).Info("Registering endpoint")

				// Use a closure to capture the endpoint and service name correctly
				ep := endpoint
				serviceName := name
				http.HandleFunc(ep.Path, func(w http.ResponseWriter, r *http.Request) {
					reqLogger := logger.WithFields(logger.Fields{
						"path":       ep.Path,
						"service":    serviceName,
						"method":     r.Method,
						"remote_ip":  r.RemoteAddr,
						"user_agent": r.UserAgent(),
						"request_id": r.Header.Get("X-Request-ID"),
					})
					reqLogger.Info("Handling request")

					// Get the client for this service
					client, ok := jrpcClients[serviceName]
					if !ok {
						reqLogger.Error("Service not found in jrpcClients map")
						http.Error(w, fmt.Sprintf("Service %s not found", serviceName), http.StatusInternalServerError)
						return
					}

					body, err := ioutil.ReadAll(r.Body)
					if err != nil {
						reqLogger.WithFields(logger.Fields{
							"error": err,
						}).Error("Error reading request body")
						http.Error(w, "Error reading request body", http.StatusInternalServerError)
						return
					}
					defer r.Body.Close()

					var params json.RawMessage
					if len(body) > 0 {
						if err := json.Unmarshal(body, &params); err != nil {
							reqLogger.WithFields(logger.Fields{
								"error": err,
							}).Error("Error parsing JSON body")
							http.Error(w, "Error parsing JSON body", http.StatusBadRequest)
							return
						}
					}

					// Call the adapter
					var result json.RawMessage
					reqCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
					defer cancel()

					startTime := time.Now()
					if err := client.CallResult(reqCtx, ep.McpMethod, params, &result); err != nil {
						reqLogger.WithFields(logger.Fields{
							"error":       err,
							"duration_ms": time.Since(startTime).Milliseconds(),
						}).Error("RPC call failed")

						// Check for JSON-RPC specific errors
						if jrpcErr, ok := err.(*jrpc2.Error); ok {
							w.WriteHeader(http.StatusBadRequest)
							json.NewEncoder(w).Encode(jrpcErr)
						} else {
							http.Error(w, fmt.Sprintf("RPC call failed: %v", err), http.StatusInternalServerError)
						}
						return
					}

					reqLogger.WithFields(logger.Fields{
						"duration_ms": time.Since(startTime).Milliseconds(),
					}).Info("Request completed successfully")

					// Write the successful response back
					w.Header().Set("Content-Type", "application/json")
					w.Write(result)
				})
			}
		}

		// Register the /reload endpoint
		http.HandleFunc("/reload", func(w http.ResponseWriter, r *http.Request) {
			if r.Method != http.MethodPost {
				http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
				return
			}

			logger.Info("Received reload request via HTTP endpoint")

			// Send SIGHUP to self
			process, err := os.FindProcess(os.Getpid())
			if err != nil {
				logger.WithFields(logger.Fields{
					"error": err,
				}).Error("Failed to find own process")
				http.Error(w, "Failed to initiate reload", http.StatusInternalServerError)
				return
			}

			if err := process.Signal(syscall.SIGHUP); err != nil {
				logger.WithFields(logger.Fields{
					"error": err,
				}).Error("Failed to send SIGHUP to self")
				http.Error(w, "Failed to initiate reload", http.StatusInternalServerError)
				return
			}

			w.WriteHeader(http.StatusAccepted)
			w.Write([]byte(`{"status":"reload initiated"}`))
		})

		return nil
	}

	// Initialize services with the initial configuration
	if err := initializeServices(config); err != nil {
		logger.WithFields(logger.Fields{
			"error": err,
		}).Fatal("Failed to initialize services")
		return
	}

	// Create a new HTTP server with the logger middleware
	portStr := strconv.Itoa(config.Port)
	server := &http.Server{
		Addr:    ":" + portStr,
		Handler: logger.HTTPMiddleware(http.DefaultServeMux),
	}

	// Start the HTTP server
	serverErr := make(chan error, 1)
	go func() {
		logger.WithFields(logger.Fields{
			"port": portStr,
		}).Info("Starting server")
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			serverErr <- err
		}
	}()

	// Set up signal handling for graceful shutdown and reload
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	signal.Notify(sighupCh, syscall.SIGHUP)

	// Handle SIGHUP for configuration reload
	go func() {
		for {
			select {
			case <-sighupCh:
				logger.Info("Received SIGHUP signal, reloading configuration")

				// Load the updated configuration
				newConfig, err := loadConfig(*configFile)
				if err != nil {
					logger.WithFields(logger.Fields{
						"error": err,
					}).Error("Failed to reload configuration, continuing with existing configuration")
					continue
				}

				// Keep the port from the current configuration
				newConfig.Port = config.Port

				// Update the global configuration
				config = newConfig

				// Reload services
				if err := initializeServices(config); err != nil {
					logger.WithFields(logger.Fields{
						"error": err,
					}).Error("Failed to reload services")
				} else {
					logger.Info("Configuration reload completed successfully")
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	// Block until a shutdown signal or server error
	select {
	case err := <-serverErr:
		logger.WithFields(logger.Fields{
			"error": err,
		}).Error("HTTP server error. Shutting down.")
	case sig := <-c:
		logger.WithFields(logger.Fields{
			"signal": sig,
		}).Info("Received signal. Shutting down.")
	}

	// Gracefully shut down the HTTP server
	logger.Info("Shutting down server...")
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownCtx); err != nil {
		logger.WithFields(logger.Fields{
			"error": err,
		}).Error("Server shutdown failed")
	}

	logger.Info("Main function finished. Deferred cleanup will now run.")
}
