package windsurf

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/sync/singleflight"
)

const (
	DefaultLSBinary  = "/opt/windsurf/language_server_linux_x64"
	DefaultLSPort    = 42100
	DefaultCSRF      = "windsurf-api-csrf-fixed-token"
	DefaultAPIServer = "https://server.self-serve.windsurf.com"
)

type LSPoolConfig struct {
	Binary       string
	BasePort     int
	CSRFToken    string
	APIServerURL string
	DataDir      string
}

func (c *LSPoolConfig) defaults() {
	if c.Binary == "" {
		c.Binary = os.Getenv("LS_BINARY_PATH")
		if c.Binary == "" {
			c.Binary = DefaultLSBinary
		}
	}
	if c.BasePort <= 0 {
		c.BasePort = DefaultLSPort
	}
	if c.CSRFToken == "" {
		c.CSRFToken = DefaultCSRF
	}
	if c.APIServerURL == "" {
		c.APIServerURL = os.Getenv("CODEIUM_API_URL")
		if c.APIServerURL == "" {
			c.APIServerURL = DefaultAPIServer
		}
	}
	if c.DataDir == "" {
		c.DataDir = "/opt/windsurf/data"
	}
}

type LSEntry struct {
	Cmd       *exec.Cmd
	Port      int
	CSRFToken string
	Client    *LocalLSClient
	ProxyKey  string
	Ready     atomic.Bool
	StartedAt time.Time
	done      chan struct{} // closed when the process exits
}

type LSPool struct {
	pool     map[string]*LSEntry
	mu       sync.RWMutex
	sf       singleflight.Group
	nextPort atomic.Int32
	config   LSPoolConfig
	logFunc  func(format string, args ...any)
}

func NewLSPool(cfg LSPoolConfig, logFn func(string, ...any)) *LSPool {
	cfg.defaults()
	p := &LSPool{
		pool:    make(map[string]*LSEntry),
		config:  cfg,
		logFunc: logFn,
	}
	p.nextPort.Store(int32(cfg.BasePort + 1))
	return p
}

func (p *LSPool) log(format string, args ...any) {
	if p.logFunc != nil {
		p.logFunc(format, args...)
	}
}

var nonAlphaNum = regexp.MustCompile(`[^a-zA-Z0-9]`)

// proxyKey produces a pool key from a proxy URL.
// Includes auth hash so different credentials on the same host get separate LS instances.
func proxyKey(proxyURL string) string {
	proxyURL = strings.TrimSpace(proxyURL)
	if proxyURL == "" {
		return "default"
	}
	u, err := url.Parse(proxyURL)
	if err != nil {
		return "px_" + nonAlphaNum.ReplaceAllString(proxyURL, "_")
	}
	key := u.Hostname()
	if u.Port() != "" {
		key += "_" + u.Port()
	}
	if u.User != nil {
		key += "_" + nonAlphaNum.ReplaceAllString(u.User.Username(), "_")
	}
	return "px_" + nonAlphaNum.ReplaceAllString(key, "_")
}

// redactProxyURL strips credentials from a proxy URL for safe logging.
func redactProxyURL(proxyURL string) string {
	if proxyURL == "" {
		return "none"
	}
	u, err := url.Parse(proxyURL)
	if err != nil {
		return "<invalid>"
	}
	u.User = nil
	return u.String()
}

func (p *LSPool) Ensure(ctx context.Context, proxyURL string) (*LSEntry, error) {
	key := proxyKey(proxyURL)

	p.mu.RLock()
	if e, ok := p.pool[key]; ok && e.Ready.Load() {
		p.mu.RUnlock()
		return e, nil
	}
	p.mu.RUnlock()

	val, err, _ := p.sf.Do(key, func() (any, error) {
		p.mu.RLock()
		if e, ok := p.pool[key]; ok && e.Ready.Load() {
			p.mu.RUnlock()
			return e, nil
		}
		p.mu.RUnlock()
		return p.spawnLS(ctx, key, proxyURL)
	})
	if err != nil {
		return nil, err
	}
	return val.(*LSEntry), nil
}

func (p *LSPool) Get(proxyURL string) *LSEntry {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.pool[proxyKey(proxyURL)]
}

func (p *LSPool) Restart(ctx context.Context, proxyURL string) (*LSEntry, error) {
	key := proxyKey(proxyURL)
	p.mu.Lock()
	if old, ok := p.pool[key]; ok {
		p.stopEntry(old)
		delete(p.pool, key)
	}
	p.mu.Unlock()
	return p.Ensure(ctx, proxyURL)
}

func (p *LSPool) Shutdown() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for key, entry := range p.pool {
		p.stopEntry(entry)
		p.log("LS instance %s stopped", key)
	}
	p.pool = make(map[string]*LSEntry)
}

func (p *LSPool) stopEntry(e *LSEntry) {
	e.Ready.Store(false)
	if e.Cmd == nil || e.Cmd.Process == nil {
		return
	}
	_ = e.Cmd.Process.Signal(os.Interrupt)
	select {
	case <-e.done:
	case <-time.After(5 * time.Second):
		_ = e.Cmd.Process.Kill()
		<-e.done
	}
}

type LSStatus struct {
	Running   bool
	Instances []LSInstanceStatus
}

type LSInstanceStatus struct {
	Key       string
	Port      int
	PID       int
	ProxyKey  string
	StartedAt time.Time
	Ready     bool
}

func (p *LSPool) Status() LSStatus {
	p.mu.RLock()
	defer p.mu.RUnlock()
	s := LSStatus{Running: len(p.pool) > 0}
	for key, e := range p.pool {
		pid := 0
		if e.Cmd != nil && e.Cmd.Process != nil {
			pid = e.Cmd.Process.Pid
		}
		s.Instances = append(s.Instances, LSInstanceStatus{
			Key: key, Port: e.Port, PID: pid,
			ProxyKey: e.ProxyKey, StartedAt: e.StartedAt, Ready: e.Ready.Load(),
		})
	}
	return s
}

func (p *LSPool) allocPort(isDefault bool) (int, error) {
	if isDefault {
		return p.config.BasePort, nil
	}
	for i := 0; i < 50; i++ {
		port := int(p.nextPort.Add(1)) - 1
		if !isPortInUse(port) {
			return port, nil
		}
		p.log("LS port %d busy, advancing", port)
	}
	return 0, fmt.Errorf("no free port for LS in 50 attempts starting from %d", p.config.BasePort+1)
}

func isPortInUse(port int) bool {
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), time.Second)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

func waitPortReady(port int, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	h2t := &http2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{Timeout: 2 * time.Second}).DialContext(ctx, network, addr)
		},
	}
	defer h2t.CloseIdleConnections()

	for time.Now().Before(deadline) {
		conn, err := h2t.DialTLSContext(context.Background(), "tcp", fmt.Sprintf("127.0.0.1:%d", port), nil)
		if err == nil {
			conn.Close()
			return nil
		}
		time.Sleep(500 * time.Millisecond)
	}
	return fmt.Errorf("LS port %d not ready after %v", port, timeout)
}

func (p *LSPool) spawnLS(ctx context.Context, key, proxyURL string) (*LSEntry, error) {
	isDefault := key == "default"

	if isDefault && isPortInUse(p.config.BasePort) {
		p.log("LS default port %d already in use — adopting existing instance", p.config.BasePort)
		entry := &LSEntry{
			Port:      p.config.BasePort,
			CSRFToken: p.config.CSRFToken,
			Client:    NewLocalLSClient(p.config.BasePort, p.config.CSRFToken),
			ProxyKey:  key,
			StartedAt: time.Now(),
			done:      make(chan struct{}),
		}
		entry.Ready.Store(true)
		close(entry.done)
		p.mu.Lock()
		p.pool[key] = entry
		p.mu.Unlock()
		return entry, nil
	}

	port, err := p.allocPort(isDefault)
	if err != nil {
		return nil, err
	}

	dataDir := filepath.Join(p.config.DataDir, key)
	if err := os.MkdirAll(filepath.Join(dataDir, "db"), 0o755); err != nil {
		return nil, fmt.Errorf("mkdirAll %s/db: %w", dataDir, err)
	}

	args := []string{
		fmt.Sprintf("--api_server_url=%s", p.config.APIServerURL),
		fmt.Sprintf("--server_port=%d", port),
		fmt.Sprintf("--csrf_token=%s", p.config.CSRFToken),
		"--register_user_url=https://api.codeium.com/register_user/",
		fmt.Sprintf("--codeium_dir=%s", dataDir),
		fmt.Sprintf("--database_dir=%s/db", dataDir),
		"--enable_local_search=false",
		"--enable_index_service=false",
		"--enable_lsp=false",
		"--detect_proxy=false",
	}

	// Don't bind LS process lifetime to request context — use background context for the process.
	cmd := exec.Command(p.config.Binary, args...)
	cmd.Env = append(os.Environ(), "HOME=/root")
	if proxyURL != "" {
		cmd.Env = append(cmd.Env,
			"HTTPS_PROXY="+proxyURL,
			"HTTP_PROXY="+proxyURL,
			"https_proxy="+proxyURL,
			"http_proxy="+proxyURL,
		)
	}

	cmd.Stdout = nil
	cmd.Stderr = nil

	p.log("Starting LS instance key=%s port=%d proxy=%s", key, port, redactProxyURL(proxyURL))

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("spawn LS %s: %w", key, err)
	}

	entry := &LSEntry{
		Cmd:       cmd,
		Port:      port,
		CSRFToken: p.config.CSRFToken,
		Client:    NewLocalLSClient(port, p.config.CSRFToken),
		ProxyKey:  key,
		StartedAt: time.Now(),
		done:      make(chan struct{}),
	}

	p.mu.Lock()
	p.pool[key] = entry
	p.mu.Unlock()

	go p.monitorProcess(key, entry)

	if err := waitPortReady(port, 25*time.Second); err != nil {
		p.log("LS instance %s failed to become ready: %v", key, err)
		_ = cmd.Process.Kill()
		p.mu.Lock()
		delete(p.pool, key)
		p.mu.Unlock()
		<-entry.done
		return nil, err
	}

	entry.Ready.Store(true)
	p.log("LS instance %s ready on port %d", key, port)
	return entry, nil
}

// monitorProcess is the sole reaper for the LS process.
func (p *LSPool) monitorProcess(key string, entry *LSEntry) {
	err := entry.Cmd.Wait()
	close(entry.done)
	entry.Ready.Store(false)

	exitMsg := "nil"
	if err != nil {
		exitMsg = err.Error()
	}
	p.log("LS instance %s exited: %s", key, exitMsg)

	p.mu.Lock()
	if cur, ok := p.pool[key]; ok && cur == entry {
		delete(p.pool, key)
	}
	p.mu.Unlock()
}