package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"lingma-ipc-proxy/internal/httpapi"
	"lingma-ipc-proxy/internal/lingmaipc"
	"lingma-ipc-proxy/internal/service"
)

type fileConfig struct {
	Host                       string   `json:"host"`
	Port                       int      `json:"port"`
	Backend                    string   `json:"backend"`
	Transport                  string   `json:"transport"`
	Pipe                       string   `json:"pipe"`
	WebSocketURL               string   `json:"websocket_url"`
	RemoteBaseURL              string   `json:"remote_base_url"`
	RemoteAuthFile             string   `json:"remote_auth_file"`
	RemoteVersion              string   `json:"remote_version"`
	Cwd                        string   `json:"cwd"`
	CurrentFilePath            string   `json:"current_file_path"`
	Mode                       string   `json:"mode"`
	Model                      string   `json:"model"`
	ShellType                  string   `json:"shell_type"`
	SessionMode                string   `json:"session_mode"`
	TimeoutSeconds             int      `json:"timeout"`
	RemoteFallbackEnabled      *bool    `json:"remote_fallback_enabled"`
	RemoteFallbackModels       []string `json:"remote_fallback_models"`
	LingmaBootstrapEnabled     *bool    `json:"lingma_bootstrap_enabled"`
	LingmaSourceType           string   `json:"lingma_source_type"`
	LingmaVSIXURL              string   `json:"lingma_vsix_url"`
	LingmaMarketplacePublisher string   `json:"lingma_marketplace_publisher"`
	LingmaMarketplaceExtension string   `json:"lingma_marketplace_extension"`
	LingmaBootstrapOutputDir   string   `json:"lingma_bootstrap_output_dir"`
	LingmaBinaryPath           string   `json:"lingma_binary_path"`
	LingmaBootstrapAlways      *bool    `json:"lingma_bootstrap_always"`
	LingmaForceRefresh         *bool    `json:"lingma_force_refresh"`
	LingmaWorkDir              string   `json:"lingma_work_dir"`
	LingmaSessionBundle        string   `json:"lingma_session_bundle"`
	LingmaSessionBundleFile    string   `json:"lingma_session_bundle_file"`
}

func main() {
	cfg, configPath := loadConfig()
	addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)

	svc := service.New(cfg)
	if err := svc.PrepareRuntime(); err != nil {
		log.Fatalf("prepare runtime: %v", err)
	}
	warmupCtx, warmupCancel := context.WithTimeout(context.Background(), 10*time.Second)
	if err := svc.Warmup(warmupCtx); err != nil {
		log.Printf("warmup failed: %v", err)
	} else {
		log.Printf("Lingma IPC warmup completed")
	}
	warmupCancel()

	server := httpapi.NewServer(addr, svc)

	log.Printf("lingma-proxy listening on http://%s", addr)
	log.Printf("session mode: %s", cfg.SessionMode)
	log.Printf("transport: %s", cfg.Transport)
	log.Printf("mode: %s", cfg.Mode)
	if configPath != "" {
		log.Printf("config file: %s", configPath)
	}

	errCh := make(chan error, 1)
	go func() {
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			errCh <- err
		}
	}()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	select {
	case err := <-errCh:
		log.Fatal(err)
	case sig := <-sigCh:
		log.Printf("received %s, shutting down", sig.String())
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := server.Shutdown(ctx); err != nil {
		log.Fatal(err)
	}
}

func loadConfig() (service.Config, string) {
	cfg := service.Config{
		Host:                   "127.0.0.1",
		Port:                   8095,
		Backend:                service.BackendRemote,
		Transport:              lingmaipc.TransportAuto,
		Cwd:                    currentDir(),
		Mode:                   "agent",
		Model:                  "kmodel",
		ShellType:              defaultShellType(),
		SessionMode:            service.SessionModeAuto,
		Timeout:                0,
		RemoteFallbackEnabled:  true,
		RemoteFallbackModels:   service.DefaultRemoteFallbackModels(),
		LingmaBootstrapEnabled: false,
		LingmaSourceType:       "marketplace",
		LingmaBootstrapAlways:  true,
		LingmaForceRefresh:     false,
	}

	configPath, configLoaded := resolveConfigPath()
	if configLoaded {
		fileCfg, err := readFileConfig(configPath)
		if err != nil {
			log.Fatalf("load config file %q: %v", configPath, err)
		}
		overlayFileConfig(&cfg, fileCfg)
	}

	overlayEnvConfig(&cfg)

	host := flag.String("host", cfg.Host, "Listen host")
	port := flag.Int("port", cfg.Port, "Listen port")
	transport := flag.String("transport", string(cfg.Transport), "Lingma transport: auto, pipe, websocket")
	backend := flag.String("backend", string(cfg.Backend), "Backend mode: ipc or remote")
	pipe := flag.String("pipe", cfg.Pipe, "Explicit Lingma named pipe path")
	wsURL := flag.String("ws-url", cfg.WebSocketURL, "Explicit Lingma local websocket URL")
	remoteBaseURL := flag.String("remote-base-url", cfg.RemoteBaseURL, "Remote Lingma API base URL")
	remoteAuthFile := flag.String("remote-auth-file", cfg.RemoteAuthFile, "Remote Lingma credentials.json path; empty reads ~/.lingma cache")
	remoteVersion := flag.String("remote-version", cfg.RemoteVersion, "Remote Lingma cosy version")
	cwd := flag.String("cwd", cfg.Cwd, "Working directory used when creating Lingma sessions")
	currentFilePath := flag.String("current-file-path", cfg.CurrentFilePath, "Current file path sent through ACP meta")
	mode := flag.String("mode", cfg.Mode, "Lingma ACP mode value")
	model := flag.String("model", cfg.Model, "Default Lingma model when API request omits model")
	shellType := flag.String("shell-type", cfg.ShellType, "Shell type sent through ACP meta")
	timeoutSeconds := flag.Int("timeout", int(cfg.Timeout/time.Second), "Per-request timeout in seconds; 0 disables the proxy deadline")
	remoteFallbackEnabled := flag.Bool("remote-fallback", cfg.RemoteFallbackEnabled, "Enable remote timeout/5xx fallback to the next available model")
	remoteFallbackModels := flag.String("remote-fallback-models", strings.Join(cfg.RemoteFallbackModels, ","), "Comma-separated remote fallback model IDs")
	sessionMode := flag.String("session-mode", string(cfg.SessionMode), "Session mode: auto, fresh, reuse")
	lingmaBootstrap := flag.Bool("lingma-bootstrap", cfg.LingmaBootstrapEnabled, "Download/extract Lingma runtime assets before startup")
	lingmaSourceType := flag.String("lingma-source-type", cfg.LingmaSourceType, "Lingma bootstrap source: marketplace or vsix")
	lingmaVSIXURL := flag.String("lingma-vsix-url", cfg.LingmaVSIXURL, "Lingma VSIX URL used when bootstrap source is vsix or marketplace fallback")
	lingmaMarketplacePublisher := flag.String("lingma-marketplace-publisher", cfg.LingmaMarketplacePublisher, "VS Code marketplace publisher for Lingma bootstrap")
	lingmaMarketplaceExtension := flag.String("lingma-marketplace-extension", cfg.LingmaMarketplaceExtension, "VS Code marketplace extension name for Lingma bootstrap")
	lingmaBootstrapOutputDir := flag.String("lingma-bootstrap-output-dir", cfg.LingmaBootstrapOutputDir, "Lingma bootstrap release output directory")
	lingmaBinaryPath := flag.String("lingma-binary-path", cfg.LingmaBinaryPath, "Lingma binary output path")
	lingmaBootstrapAlways := flag.Bool("lingma-bootstrap-always", cfg.LingmaBootstrapAlways, "Re-check bootstrap source at startup")
	lingmaForceRefresh := flag.Bool("lingma-force-refresh", cfg.LingmaForceRefresh, "Force refresh Lingma bootstrap assets")
	lingmaWorkDir := flag.String("lingma-work-dir", cfg.LingmaWorkDir, "Lingma work/cache directory used for restored session bundles")
	lingmaSessionBundle := flag.String("lingma-session-bundle", cfg.LingmaSessionBundle, "Base64 tar.gz Lingma session bundle to restore before startup")
	lingmaSessionBundleFile := flag.String("lingma-session-bundle-file", cfg.LingmaSessionBundleFile, "File containing a base64 tar.gz Lingma session bundle")
	config := flag.String("config", valueOr(configPath, filepath.Join(currentDir(), "lingma-proxy.json")), "Path to JSON config file")
	flag.Parse()

	parsedSessionMode := parseSessionMode(*sessionMode)
	parsedTransport := parseTransport(*transport)
	finalConfigPath := strings.TrimSpace(*config)

	cfg.Host = strings.TrimSpace(*host)
	cfg.Port = *port
	cfg.Backend = parseBackend(*backend)
	cfg.Transport = parsedTransport
	cfg.Pipe = strings.TrimSpace(*pipe)
	cfg.WebSocketURL = strings.TrimSpace(*wsURL)
	cfg.RemoteBaseURL = strings.TrimSpace(*remoteBaseURL)
	cfg.RemoteAuthFile = strings.TrimSpace(*remoteAuthFile)
	cfg.RemoteVersion = strings.TrimSpace(*remoteVersion)
	cfg.Cwd = strings.TrimSpace(*cwd)
	cfg.CurrentFilePath = strings.TrimSpace(*currentFilePath)
	cfg.Mode = strings.TrimSpace(*mode)
	cfg.Model = strings.TrimSpace(*model)
	cfg.ShellType = strings.TrimSpace(*shellType)
	cfg.SessionMode = parsedSessionMode
	cfg.Timeout = time.Duration(*timeoutSeconds) * time.Second
	cfg.RemoteFallbackEnabled = *remoteFallbackEnabled
	cfg.RemoteFallbackModels = splitCSV(*remoteFallbackModels)
	cfg.LingmaBootstrapEnabled = *lingmaBootstrap
	cfg.LingmaSourceType = strings.TrimSpace(*lingmaSourceType)
	cfg.LingmaVSIXURL = strings.TrimSpace(*lingmaVSIXURL)
	cfg.LingmaMarketplacePublisher = strings.TrimSpace(*lingmaMarketplacePublisher)
	cfg.LingmaMarketplaceExtension = strings.TrimSpace(*lingmaMarketplaceExtension)
	cfg.LingmaBootstrapOutputDir = strings.TrimSpace(*lingmaBootstrapOutputDir)
	cfg.LingmaBinaryPath = strings.TrimSpace(*lingmaBinaryPath)
	cfg.LingmaBootstrapAlways = *lingmaBootstrapAlways
	cfg.LingmaForceRefresh = *lingmaForceRefresh
	cfg.LingmaWorkDir = strings.TrimSpace(*lingmaWorkDir)
	cfg.LingmaSessionBundle = strings.TrimSpace(*lingmaSessionBundle)
	cfg.LingmaSessionBundleFile = strings.TrimSpace(*lingmaSessionBundleFile)

	if configLoaded {
		configPath = finalConfigPath
	} else {
		configPath = ""
	}

	return cfg, configPath
}

func resolveConfigPath() (string, bool) {
	if path := strings.TrimSpace(lookupArgValue("--config")); path != "" {
		return path, true
	}
	if path := strings.TrimSpace(os.Getenv("LINGMA_PROXY_CONFIG")); path != "" {
		return path, true
	}
	defaultPath := filepath.Join(currentDir(), "lingma-proxy.json")
	for _, candidate := range []string{defaultPath, filepath.Join(currentDir(), "lingma-ipc-proxy.json")} {
		if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
			return candidate, true
		}
	}
	return defaultPath, false
}

func readFileConfig(path string) (fileConfig, error) {
	var cfg fileConfig
	body, err := os.ReadFile(path)
	if err != nil {
		return cfg, err
	}
	if err := json.Unmarshal(body, &cfg); err != nil {
		return cfg, err
	}
	return cfg, nil
}

func overlayFileConfig(dst *service.Config, src fileConfig) {
	if strings.TrimSpace(src.Host) != "" {
		dst.Host = strings.TrimSpace(src.Host)
	}
	if src.Port > 0 {
		dst.Port = src.Port
	}
	if strings.TrimSpace(src.Transport) != "" {
		dst.Transport = parseTransport(src.Transport)
	}
	if strings.TrimSpace(src.Backend) != "" {
		dst.Backend = parseBackend(src.Backend)
	}
	if strings.TrimSpace(src.Pipe) != "" {
		dst.Pipe = strings.TrimSpace(src.Pipe)
	}
	if strings.TrimSpace(src.WebSocketURL) != "" {
		dst.WebSocketURL = strings.TrimSpace(src.WebSocketURL)
	}
	if strings.TrimSpace(src.RemoteBaseURL) != "" {
		dst.RemoteBaseURL = strings.TrimSpace(src.RemoteBaseURL)
	}
	if strings.TrimSpace(src.RemoteAuthFile) != "" {
		dst.RemoteAuthFile = strings.TrimSpace(src.RemoteAuthFile)
	}
	if strings.TrimSpace(src.RemoteVersion) != "" {
		dst.RemoteVersion = strings.TrimSpace(src.RemoteVersion)
	}
	if strings.TrimSpace(src.Cwd) != "" {
		dst.Cwd = strings.TrimSpace(src.Cwd)
	}
	if strings.TrimSpace(src.CurrentFilePath) != "" {
		dst.CurrentFilePath = strings.TrimSpace(src.CurrentFilePath)
	}
	if strings.TrimSpace(src.Mode) != "" {
		dst.Mode = strings.TrimSpace(src.Mode)
	}
	if strings.TrimSpace(src.Model) != "" {
		dst.Model = strings.TrimSpace(src.Model)
	}
	if strings.TrimSpace(src.ShellType) != "" {
		dst.ShellType = strings.TrimSpace(src.ShellType)
	}
	if strings.TrimSpace(src.SessionMode) != "" {
		dst.SessionMode = parseSessionMode(src.SessionMode)
	}
	if src.TimeoutSeconds >= 0 {
		dst.Timeout = time.Duration(src.TimeoutSeconds) * time.Second
	}
	if src.RemoteFallbackEnabled != nil {
		dst.RemoteFallbackEnabled = *src.RemoteFallbackEnabled
	}
	if len(src.RemoteFallbackModels) > 0 {
		dst.RemoteFallbackModels = cleanStringSlice(src.RemoteFallbackModels)
	}
	if src.LingmaBootstrapEnabled != nil {
		dst.LingmaBootstrapEnabled = *src.LingmaBootstrapEnabled
	}
	if strings.TrimSpace(src.LingmaSourceType) != "" {
		dst.LingmaSourceType = strings.TrimSpace(src.LingmaSourceType)
	}
	if strings.TrimSpace(src.LingmaVSIXURL) != "" {
		dst.LingmaVSIXURL = strings.TrimSpace(src.LingmaVSIXURL)
	}
	if strings.TrimSpace(src.LingmaMarketplacePublisher) != "" {
		dst.LingmaMarketplacePublisher = strings.TrimSpace(src.LingmaMarketplacePublisher)
	}
	if strings.TrimSpace(src.LingmaMarketplaceExtension) != "" {
		dst.LingmaMarketplaceExtension = strings.TrimSpace(src.LingmaMarketplaceExtension)
	}
	if strings.TrimSpace(src.LingmaBootstrapOutputDir) != "" {
		dst.LingmaBootstrapOutputDir = strings.TrimSpace(src.LingmaBootstrapOutputDir)
	}
	if strings.TrimSpace(src.LingmaBinaryPath) != "" {
		dst.LingmaBinaryPath = strings.TrimSpace(src.LingmaBinaryPath)
	}
	if src.LingmaBootstrapAlways != nil {
		dst.LingmaBootstrapAlways = *src.LingmaBootstrapAlways
	}
	if src.LingmaForceRefresh != nil {
		dst.LingmaForceRefresh = *src.LingmaForceRefresh
	}
	if strings.TrimSpace(src.LingmaWorkDir) != "" {
		dst.LingmaWorkDir = strings.TrimSpace(src.LingmaWorkDir)
	}
	if strings.TrimSpace(src.LingmaSessionBundle) != "" {
		dst.LingmaSessionBundle = strings.TrimSpace(src.LingmaSessionBundle)
	}
	if strings.TrimSpace(src.LingmaSessionBundleFile) != "" {
		dst.LingmaSessionBundleFile = strings.TrimSpace(src.LingmaSessionBundleFile)
	}
}

func overlayEnvConfig(dst *service.Config) {
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_HOST")); value != "" {
		dst.Host = value
	}
	if value := envInt("LINGMA_PROXY_PORT", 0); value > 0 {
		dst.Port = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_TRANSPORT")); value != "" {
		dst.Transport = parseTransport(value)
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_BACKEND")); value != "" {
		dst.Backend = parseBackend(value)
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_IPC_PIPE")); value != "" {
		dst.Pipe = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_WS_URL")); value != "" {
		dst.WebSocketURL = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_REMOTE_BASE_URL")); value != "" {
		dst.RemoteBaseURL = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_REMOTE_AUTH_FILE")); value != "" {
		dst.RemoteAuthFile = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_REMOTE_VERSION")); value != "" {
		dst.RemoteVersion = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_CWD")); value != "" {
		dst.Cwd = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_CURRENT_FILE_PATH")); value != "" {
		dst.CurrentFilePath = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_MODE")); value != "" {
		dst.Mode = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_MODEL")); value != "" {
		dst.Model = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_SHELL_TYPE")); value != "" {
		dst.ShellType = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_PROXY_SESSION_MODE")); value != "" {
		dst.SessionMode = parseSessionMode(value)
	}
	if value := envInt("LINGMA_PROXY_TIMEOUT_SECONDS", -1); value >= 0 {
		dst.Timeout = time.Duration(value) * time.Second
	}
	if value, ok := envBool("LINGMA_REMOTE_FALLBACK_ENABLED"); ok {
		dst.RemoteFallbackEnabled = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_REMOTE_FALLBACK_MODELS")); value != "" {
		dst.RemoteFallbackModels = splitCSV(value)
	}
	if value, ok := envBool("LINGMA_BOOTSTRAP_ENABLED"); ok {
		dst.LingmaBootstrapEnabled = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_SOURCE_TYPE")); value != "" {
		dst.LingmaSourceType = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_VSIX_URL")); value != "" {
		dst.LingmaVSIXURL = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_MARKETPLACE_PUBLISHER")); value != "" {
		dst.LingmaMarketplacePublisher = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_MARKETPLACE_EXTENSION")); value != "" {
		dst.LingmaMarketplaceExtension = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_BOOTSTRAP_OUTPUT_DIR")); value != "" {
		dst.LingmaBootstrapOutputDir = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_BIN")); value != "" {
		dst.LingmaBinaryPath = value
	}
	if value, ok := envBool("LINGMA_BOOTSTRAP_ALWAYS"); ok {
		dst.LingmaBootstrapAlways = value
	}
	if value, ok := envBool("LINGMA_FORCE_REFRESH"); ok {
		dst.LingmaForceRefresh = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_WORK_DIR")); value != "" {
		dst.LingmaWorkDir = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_SESSION_BUNDLE")); value != "" {
		dst.LingmaSessionBundle = value
	}
	if value := strings.TrimSpace(os.Getenv("LINGMA_SESSION_BUNDLE_FILE")); value != "" {
		dst.LingmaSessionBundleFile = value
	}
}

func parseSessionMode(value string) service.SessionMode {
	mode := service.SessionMode(strings.ToLower(strings.TrimSpace(value)))
	switch mode {
	case service.SessionModeAuto, service.SessionModeFresh, service.SessionModeReuse:
		return mode
	default:
		log.Fatalf("invalid session mode %q; expected auto, fresh, or reuse", value)
		return service.SessionModeAuto
	}
}

func parseBackend(value string) service.BackendMode {
	mode := service.BackendMode(strings.ToLower(strings.TrimSpace(value)))
	switch mode {
	case "":
		return service.BackendRemote
	case service.BackendIPC:
		return service.BackendIPC
	case service.BackendRemote:
		return service.BackendRemote
	default:
		log.Fatalf("invalid backend %q; expected ipc or remote", value)
		return service.BackendIPC
	}
}

func parseTransport(value string) lingmaipc.Transport {
	transport, err := lingmaipc.ParseTransport(value)
	if err != nil {
		log.Fatal(err)
	}
	return transport
}

func lookupArgValue(flagName string) string {
	for i := 1; i < len(os.Args); i++ {
		arg := os.Args[i]
		if arg == flagName {
			if i+1 < len(os.Args) {
				return os.Args[i+1]
			}
			return ""
		}
		prefix := flagName + "="
		if strings.HasPrefix(arg, prefix) {
			return strings.TrimPrefix(arg, prefix)
		}
	}
	return ""
}

func envInt(key string, fallback int) int {
	if value := strings.TrimSpace(os.Getenv(key)); value != "" {
		if n, err := strconv.Atoi(value); err == nil {
			return n
		}
	}
	return fallback
}

func envBool(key string) (bool, bool) {
	value := strings.ToLower(strings.TrimSpace(os.Getenv(key)))
	switch value {
	case "1", "true", "yes", "on":
		return true, true
	case "0", "false", "no", "off":
		return false, true
	default:
		return false, false
	}
}

func splitCSV(value string) []string {
	return cleanStringSlice(strings.Split(value, ","))
}

func cleanStringSlice(values []string) []string {
	out := make([]string, 0, len(values))
	seen := map[string]bool{}
	for _, value := range values {
		item := strings.TrimSpace(value)
		if item == "" || seen[item] {
			continue
		}
		seen[item] = true
		out = append(out, item)
	}
	return out
}

func currentDir() string {
	if wd, err := os.Getwd(); err == nil {
		return wd
	}
	return "."
}

func valueOr(value string, fallback string) string {
	if strings.TrimSpace(value) != "" {
		return strings.TrimSpace(value)
	}
	return fallback
}

func defaultShellType() string {
	if runtime.GOOS == "windows" {
		return "powershell"
	}
	return "zsh"
}
