package gobridge

import (
	"context"
	"fmt"
	"io"
	"net"
	"sync/atomic"
)

// poolConfig 是进程池内部配置，通过 Option 函数填充
type poolConfig struct {
	workers           int
	maxConnsPerWorker int
	pythonExe         string
	scriptArgs        []string
	workDir           string
	env               []string
	socketDir         string
	stdout            io.Writer
	stderr            io.Writer
}

// Option 是 NewPool 的函数选项
type Option func(*poolConfig)

// WithWorkers 设置 Python 进程数量（默认 2）
func WithWorkers(n int) Option {
	return func(c *poolConfig) { c.workers = n }
}

// WithMaxConns 设置每个进程的最大连接数（默认 4）
func WithMaxConns(n int) Option {
	return func(c *poolConfig) { c.maxConnsPerWorker = n }
}

// WithPythonExe 设置 Python 可执行文件（默认 "python3"）
// uv 模式：WithPythonExe("uv"), WithScriptArgs("run")
func WithPythonExe(exe string) Option {
	return func(c *poolConfig) { c.pythonExe = exe }
}

// WithScriptArgs 设置脚本路径之后的附加参数
// uv 模式示例：WithScriptArgs("run")  → 执行 uv run <script>
func WithScriptArgs(args ...string) Option {
	return func(c *poolConfig) { c.scriptArgs = args }
}

// WithWorkDir 设置子进程工作目录（默认继承当前进程）
func WithWorkDir(workDir string) Option {
	return func(c *poolConfig) { c.workDir = workDir }
}

// WithEnv 设置附加环境变量，格式为 "KEY=VALUE"
// 与当前进程环境合并，同名时以此处为准
func WithEnv(env ...string) Option {
	return func(c *poolConfig) { c.env = env }
}

// WithSocketDir 设置 UDS socket 文件目录（默认 /tmp）
func WithSocketDir(dir string) Option {
	return func(c *poolConfig) { c.socketDir = dir }
}

// WithStdout 设置子进程标准输出目标（默认 os.Stdout，传 io.Discard 可静默）
func WithStdout(w io.Writer) Option {
	return func(c *poolConfig) { c.stdout = w }
}

// WithStderr 设置子进程标准错误目标（默认 os.Stderr，传 io.Discard 可静默）
func WithStderr(w io.Writer) Option {
	return func(c *poolConfig) { c.stderr = w }
}

// Pool 管理多个 Python worker 进程及其连接池
type Pool struct {
	workers []*worker
	idx     atomic.Uint64
	reqID   atomic.Uint64
}

// NewPool 创建并启动进程池
//
//	pool, err := gobridge.NewPool("worker.py")
//	pool, err := gobridge.NewPool("worker.py", gobridge.WithWorkers(4))
//	pool, err := gobridge.NewPool("run",
//	    gobridge.WithPythonExe("uv"),
//	    gobridge.WithScriptArgs("worker.py"),
//	    gobridge.WithWorkDir("./worker"),
//	)
func NewPool(script string, opts ...Option) (*Pool, error) {
	cfg := &poolConfig{
		workers:           2,
		maxConnsPerWorker: 4,
		pythonExe:         "python3",
		socketDir:         "/tmp",
	}
	for _, o := range opts {
		o(cfg)
	}
	if script == "" {
		return nil, fmt.Errorf("NewPool: script must not be empty")
	}
	cfg.scriptArgs = append([]string{script}, cfg.scriptArgs...)

	workers := make([]*worker, cfg.workers)
	for i := range workers {
		w, err := newWorker(cfg, i)
		if err != nil {
			for j := 0; j < i; j++ {
				workers[j].stop()
			}
			return nil, fmt.Errorf("create worker %d: %w", i, err)
		}
		workers[i] = w
	}
	return &Pool{workers: workers}, nil
}

// acquire 以轮询方式从进程池取出一个可用连接
func (p *Pool) acquire(ctx context.Context) (net.Conn, *worker, error) {
	idx := p.idx.Add(1) % uint64(len(p.workers))
	w := p.workers[idx]
	conn, err := w.acquire(ctx)
	return conn, w, err
}

// Close 关闭所有 worker 进程和连接
func (p *Pool) Close() {
	for _, w := range p.workers {
		w.stop()
	}
}