// Package codeorch 在会话目录里跑一段 Python，并按行把工具调用交回宿主。
// 工作流脚本和 run_tool_script 共用这一套进程与限额。
package codeorch

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	_ "embed"
)

//go:embed agentwf.py
var agentwfSource string

const (
	defaultTimeout = 120 * time.Second
	defaultMaxWait = 600 * time.Second
	defaultCalls   = 50
	defaultResult  = 1 << 20
	defaultModel   = 16000
)

var (
	helperOnce sync.Once
	helperDir  string
	helperErr  error
)

// Call 是脚本发来的一行请求。Args 与 Choices 保持原始 JSON。
type Call struct {
	ID       string          `json:"id"`
	Op       string          `json:"op"`
	Name     string          `json:"name"`
	Args     json.RawMessage `json:"args"`
	Question string          `json:"question"`
	Choices  json.RawMessage `json:"choices"`
	Agent    string          `json:"agent"`
	Task     string          `json:"task"`
	Content  string          `json:"content"`
	Title    string          `json:"title"`
	Kind     string          `json:"kind"`
	Status   string          `json:"status"`
	StepID   string          `json:"stepId"`
}

// Limits 控制一次脚本运行。零值使用工具路径的缺省：120 秒、最长 600 秒、50 次、1 MiB、16k 字符。
type Limits struct {
	Timeout        time.Duration
	MaxTimeout     time.Duration
	MaxCalls       int
	MaxResultBytes int
	MaxModelChars  int
	FullDir        string
}

// Outcome 是给模型的文本、已成功的调用，以及超限时的全文路径。
type Outcome struct {
	Output   string
	FullPath string
	Calls    []string
	Err      error
}

// Starter 按工作目录和脚本路径准备命令。空则用 python3。
type Starter func(ctx context.Context, workDir, path string) (*exec.Cmd, error)

// Dispatch 处理一条回调。返回的错误会写回脚本并记入最终结果。
type Dispatch func(ctx context.Context, call Call) (string, error)

// ModelText 把截断后的输出、已执行调用和错误拼成给模型的文本。
func (o Outcome) ModelText() string {
	text := o.Output
	if o.Err != nil {
		msg := o.Err.Error()
		if len(o.Calls) > 0 {
			msg = "已执行：" + strings.Join(o.Calls, "、") + "\n" + msg
		}
		if text != "" {
			text += "\n"
		}
		text += msg
	}
	if o.FullPath != "" {
		if text != "" {
			text += "\n"
		}
		text += "全文在 " + o.FullPath
	}
	return text
}

// RunFile 执行已经落盘的脚本。
func RunFile(ctx context.Context, workDir, path string, lim Limits, start Starter, dispatch Dispatch) Outcome {
	if ctx == nil {
		ctx = context.Background()
	}
	timeout := lim.bounded()
	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	helper, err := helperPath()
	if err != nil {
		return Outcome{Err: err}
	}
	if start == nil {
		start = defaultStart
	}
	cmd, err := start(runCtx, workDir, path)
	if err != nil {
		return Outcome{Err: err}
	}
	cmd.Env = prependPythonPath(cmd.Env, helper)
	reqR, reqW, err := os.Pipe()
	if err != nil {
		return Outcome{Err: err}
	}
	respR, respW, err := os.Pipe()
	if err != nil {
		reqR.Close()
		reqW.Close()
		return Outcome{Err: err}
	}
	cmd.ExtraFiles = []*os.File{reqW, respR}
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		reqR.Close()
		reqW.Close()
		respR.Close()
		respW.Close()
		return Outcome{Err: err}
	}
	reqW.Close()
	respR.Close()
	defer reqR.Close()
	defer respW.Close()
	calls, callErr := serve(runCtx, reqR, respW, lim, dispatch)
	if callErr != nil {
		cancel()
	}
	waitErr := cmd.Wait()
	out := Outcome{Output: clipJoined(&stdout, &stderr), Calls: calls}
	out = clipModel(out, lim)
	switch {
	case callErr != nil:
		out.Err = callErr
	case runCtx.Err() == context.DeadlineExceeded:
		out.Err = fmt.Errorf("脚本超时")
	case runCtx.Err() == context.Canceled:
		out.Err = fmt.Errorf("脚本已取消")
	case waitErr != nil:
		out.Err = fmt.Errorf("脚本退出：%v", waitErr)
	}
	return out
}

// Run 把脚本体写入工作目录后再执行。
func Run(ctx context.Context, workDir, body string, lim Limits, start Starter, dispatch Dispatch) Outcome {
	dir := filepath.Join(workDir, ".agent")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return Outcome{Err: err}
	}
	path := filepath.Join(dir, fmt.Sprintf("orch-%d.py", time.Now().UnixNano()))
	if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
		return Outcome{Err: err}
	}
	return RunFile(ctx, workDir, path, lim, start, dispatch)
}

func (l Limits) bounded() time.Duration {
	d := l.Timeout
	if d <= 0 {
		d = defaultTimeout
	}
	max := l.MaxTimeout
	if max <= 0 {
		max = defaultMaxWait
	}
	if d > max {
		return max
	}
	return d
}

func (l Limits) calls() int {
	if l.MaxCalls <= 0 {
		return defaultCalls
	}
	return l.MaxCalls
}

func (l Limits) resultBytes() int {
	if l.MaxResultBytes <= 0 {
		return defaultResult
	}
	return l.MaxResultBytes
}

func (l Limits) modelChars() int {
	if l.MaxModelChars <= 0 {
		return defaultModel
	}
	return l.MaxModelChars
}

func serve(ctx context.Context, reqR, respW *os.File, lim Limits, dispatch Dispatch) ([]string, error) {
	sc := bufio.NewScanner(reqR)
	sc.Buffer(make([]byte, 64*1024), lim.resultBytes())
	var calls []string
	n := 0
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		n++
		var call Call
		if err := json.Unmarshal([]byte(line), &call); err != nil {
			_ = writeResp(respW, call.ID, false, "", "请求不是合法 JSON")
			continue
		}
		if n > lim.calls() {
			_ = writeResp(respW, call.ID, false, "", "脚本回调超过 50 次")
			return calls, fmt.Errorf("脚本回调超过 50 次")
		}
		if ctx.Err() != nil {
			return calls, ctx.Err()
		}
		if dispatch == nil {
			_ = writeResp(respW, call.ID, false, "", "没有回调处理")
			continue
		}
		output, err := dispatch(ctx, call)
		if err != nil {
			_ = writeResp(respW, call.ID, false, "", err.Error())
			if ctx.Err() != nil {
				return calls, ctx.Err()
			}
			continue
		}
		if call.Op == "tool.call" && call.Name != "" {
			calls = append(calls, call.Name)
		}
		if err := writeResp(respW, call.ID, true, clipBytes(output, lim.resultBytes()), ""); err != nil {
			return calls, err
		}
	}
	if err := sc.Err(); err != nil {
		if err == bufio.ErrTooLong {
			return calls, fmt.Errorf("单行请求超过 1MB")
		}
		if ctx.Err() != nil {
			return calls, ctx.Err()
		}
		return calls, err
	}
	return calls, nil
}

func writeResp(w io.Writer, id string, ok bool, output, errText string) error {
	raw, err := json.Marshal(struct {
		ID     string `json:"id"`
		OK     bool   `json:"ok"`
		Output string `json:"output,omitempty"`
		Error  string `json:"error,omitempty"`
	}{ID: id, OK: ok, Output: output, Error: errText})
	if err != nil {
		return err
	}
	_, err = w.Write(append(raw, '\n'))
	return err
}

func clipBytes(s string, n int) string {
	if n <= 0 || len(s) <= n {
		return s
	}
	cut := s[:n]
	for !utf8.ValidString(cut) && len(cut) > 0 {
		cut = cut[:len(cut)-1]
	}
	return cut
}

func clipJoined(stdout, stderr *bytes.Buffer) string {
	text := strings.TrimSpace(stdout.String())
	errText := strings.TrimSpace(stderr.String())
	if errText != "" {
		if text != "" {
			text += "\n"
		}
		text += errText
	}
	return text
}

func clipModel(out Outcome, lim Limits) Outcome {
	limit := lim.modelChars()
	if utf8.RuneCountInString(out.Output) <= limit {
		return out
	}
	dir := lim.FullDir
	if dir == "" {
		dir = os.TempDir()
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		out.Output = clipRunes(out.Output, limit)
		return out
	}
	path := filepath.Join(dir, fmt.Sprintf("script-output-%d.txt", time.Now().UnixNano()))
	if err := os.WriteFile(path, []byte(out.Output), 0o644); err != nil {
		out.Output = clipRunes(out.Output, limit)
		return out
	}
	out.FullPath = path
	out.Output = clipRunes(out.Output, limit)
	return out
}

func clipRunes(s string, n int) string {
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	i := 0
	for count := 0; count < n && i < len(s); count++ {
		_, size := utf8.DecodeRuneInString(s[i:])
		i += size
	}
	return s[:i] + "…（超长截断）"
}

func defaultStart(ctx context.Context, workDir, path string) (*exec.Cmd, error) {
	bin, err := exec.LookPath("python3")
	if err != nil {
		return nil, fmt.Errorf("找不到 python3: %w", err)
	}
	cmd := exec.CommandContext(ctx, bin, path)
	cmd.Dir = workDir
	return cmd, nil
}

func helperPath() (string, error) {
	helperOnce.Do(func() {
		dir, err := os.MkdirTemp("", "agentwf-*")
		if err != nil {
			helperErr = err
			return
		}
		if err := os.WriteFile(filepath.Join(dir, "agentwf.py"), []byte(agentwfSource), 0o644); err != nil {
			helperErr = err
			return
		}
		helperDir = dir
	})
	return helperDir, helperErr
}

func prependPythonPath(env []string, dir string) []string {
	if env == nil {
		env = os.Environ()
	}
	for i, kv := range env {
		if strings.HasPrefix(kv, "PYTHONPATH=") {
			env[i] = "PYTHONPATH=" + dir + string(os.PathListSeparator) + strings.TrimPrefix(kv, "PYTHONPATH=")
			return env
		}
	}
	return append(env, "PYTHONPATH="+dir)
}