package squad

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/stream"
)

// StepEvent 表示小队单步执行生命周期的事件
type StepEvent struct {
	StepID     string `json:"step_id"`
	Role       string `json:"role"`
	Title      string `json:"title"`
	Mode       string `json:"mode"`
	Status     string `json:"status"` // started, completed, failed, skipped
	DurationMs int64  `json:"duration_ms,omitempty"`
	Error      string `json:"error,omitempty"`
}

// ExecuteOptions 执行调度选项
type ExecuteOptions struct {
	Timeout    time.Duration
	StepNotify func(StepEvent)
}

// WorkflowExecutionResult 工作流执行结果
type WorkflowExecutionResult struct {
	Success      bool              `json:"success"`
	StepsRun     int               `json:"steps_run"`
	Deliverable  string            `json:"deliverable"`
	FinalContent []byte            `json:"-"`
	Duration     time.Duration     `json:"duration"`
	Fields       map[string]string `json:"fields,omitempty"`
	ErrorReason  string            `json:"error_reason,omitempty"`
}

// WorkflowDriver 负责动态工作流调度与沙箱隔离执行
type WorkflowDriver struct {
	engine    *WorkspaceEngine
	evaluator *StepEvaluator
}

func NewWorkflowDriver(engine *WorkspaceEngine) *WorkflowDriver {
	if engine == nil {
		engine = NewWorkspaceEngine()
	}
	return &WorkflowDriver{engine: engine, evaluator: NewStepEvaluator()}
}

// SetEvaluator 配置自定义单步评估器
func (d *WorkflowDriver) SetEvaluator(ev *StepEvaluator) {
	d.evaluator = ev
}

// LoadWorkflow 读取工作区 workflow.json
func (d *WorkflowDriver) LoadWorkflow(wsDir string) (*WorkflowPlan, error) {
	p := filepath.Join(wsDir, "workflow.json")
	data, err := os.ReadFile(p)
	if err != nil {
		return nil, fmt.Errorf("read workflow.json: %w", err)
	}
	var wf WorkflowPlan
	if err := json.Unmarshal(data, &wf); err != nil {
		return nil, fmt.Errorf("unmarshal workflow.json: %w", err)
	}
	return &wf, nil
}

// SaveWorkflow 持久化 workflow.json
func (d *WorkflowDriver) SaveWorkflow(wsDir string, wf *WorkflowPlan) error {
	p := filepath.Join(wsDir, "workflow.json")
	return writeJSON(p, wf)
}

// LoadBlackboard 读取工作区 blackboard.json
func (d *WorkflowDriver) LoadBlackboard(wsDir string) (*BlackboardState, error) {
	p := filepath.Join(wsDir, "blackboard.json")
	data, err := os.ReadFile(p)
	if err != nil {
		return nil, fmt.Errorf("read blackboard.json: %w", err)
	}
	var bb BlackboardState
	if err := json.Unmarshal(data, &bb); err != nil {
		return nil, fmt.Errorf("unmarshal blackboard.json: %w", err)
	}
	return &bb, nil
}

// SaveBlackboard 保存 blackboard.json
func (d *WorkflowDriver) SaveBlackboard(wsDir string, bb *BlackboardState) error {
	bb.UpdatedAt = time.Now()
	p := filepath.Join(wsDir, "blackboard.json")
	return writeJSON(p, bb)
}

// AddStep 动态向工作流追加任务步骤并自动构建其子工作区
func (d *WorkflowDriver) AddStep(wsDir string, step StepSpec) error {
	wf, err := d.LoadWorkflow(wsDir)
	if err != nil {
		return err
	}

	if step.SubWorkspace == "" {
		step.SubWorkspace = fmt.Sprintf("tasks/%02d_%s", len(wf.Steps)+1, step.Role)
	}
	if step.Status == "" {
		step.Status = "pending"
	}
	if step.Mode == "" {
		step.Mode = "execute"
	}

	// 准备子任务沙箱
	squadCfg, _ := d.loadSquadConfig(wsDir)
	tmpl, _ := os.ReadFile(filepath.Join(wsDir, "template.md"))
	singleStepPlan := WorkflowPlan{Steps: []StepSpec{step}}
	if err := d.engine.buildSubWorkspaces(wsDir, singleStepPlan, squadCfg, string(tmpl)); err != nil {
		return fmt.Errorf("build added sub workspace: %w", err)
	}

	wf.Steps = append(wf.Steps, step)
	return d.SaveWorkflow(wsDir, wf)
}

// RemoveStep 动态删除未执行的任务步骤
func (d *WorkflowDriver) RemoveStep(wsDir string, stepID string) error {
	wf, err := d.LoadWorkflow(wsDir)
	if err != nil {
		return err
	}

	idx := -1
	for i, s := range wf.Steps {
		if s.ID == stepID {
			idx = i
			break
		}
	}
	if idx == -1 {
		return fmt.Errorf("step %s not found in workflow", stepID)
	}

	wf.Steps = append(wf.Steps[:idx], wf.Steps[idx+1:]...)
	return d.SaveWorkflow(wsDir, wf)
}

// UpdateStep 动态修改现有任务步骤（如重写验收标准或切换模式）
func (d *WorkflowDriver) UpdateStep(wsDir string, step StepSpec) error {
	wf, err := d.LoadWorkflow(wsDir)
	if err != nil {
		return err
	}

	found := false
	for i, s := range wf.Steps {
		if s.ID == step.ID {
			wf.Steps[i] = step
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("step %s not found to update", step.ID)
	}

	// 同步更新子工作区 TASK.md
	squadCfg, _ := d.loadSquadConfig(wsDir)
	tmpl, _ := os.ReadFile(filepath.Join(wsDir, "template.md"))
	_ = d.engine.buildSubWorkspaces(wsDir, WorkflowPlan{Steps: []StepSpec{step}}, squadCfg, string(tmpl))

	return d.SaveWorkflow(wsDir, wf)
}

// Execute 调度并执行工作区内的全套步骤
func (d *WorkflowDriver) Execute(ctx context.Context, wsDir string, backend agent.Backend, opts ExecuteOptions) (*WorkflowExecutionResult, error) {
	start := time.Now()

	wf, err := d.LoadWorkflow(wsDir)
	if err != nil {
		return nil, err
	}
	wf.Status = "in_progress"
	_ = d.SaveWorkflow(wsDir, wf)

	_ = d.updateWorkspaceStatus(wsDir, "running")

	timeout := opts.Timeout
	if timeout <= 0 {
		timeout = 30 * time.Minute
	}
	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var mu sync.Mutex
	stepsRun := 0
	var producedFiles []string

	i := 0
	for i < len(wf.Steps) {
		step := wf.Steps[i]
		if step.Status == "skipped" {
			i++
			continue
		}

		if step.Parallel {
			var parallelGroup []StepSpec
			for i < len(wf.Steps) && wf.Steps[i].Parallel {
				if wf.Steps[i].Status != "skipped" {
					parallelGroup = append(parallelGroup, wf.Steps[i])
				}
				i++
			}

			if len(parallelGroup) > 0 {
				var wg sync.WaitGroup
				var groupErr error
				for _, s := range parallelGroup {
					sCopy := s
					wg.Add(1)
					go func() {
						defer wg.Done()
						if err := d.runOneStep(runCtx, wsDir, sCopy, backend, opts.StepNotify, &mu, &producedFiles); err != nil {
							mu.Lock()
							if groupErr == nil {
								groupErr = err
							}
							mu.Unlock()
						}
					}()
				}
				wg.Wait()
				if groupErr != nil {
					wf.Status = "failed"
					_ = d.SaveWorkflow(wsDir, wf)
					_ = d.updateWorkspaceStatus(wsDir, "failed")
					return &WorkflowExecutionResult{
						Success:     false,
						ErrorReason: groupErr.Error(),
						Duration:    time.Since(start),
					}, groupErr
				}
				stepsRun += len(parallelGroup)
				for idx := range wf.Steps {
					for _, pg := range parallelGroup {
						if wf.Steps[idx].ID == pg.ID {
							wf.Steps[idx].Status = "completed"
						}
					}
				}
				_ = d.SaveWorkflow(wsDir, wf)
			}
		} else {
			stepIndex := i
			i++
			if err := d.runOneStep(runCtx, wsDir, step, backend, opts.StepNotify, &mu, &producedFiles); err != nil {
				wf.Status = "failed"
				_ = d.SaveWorkflow(wsDir, wf)
				_ = d.updateWorkspaceStatus(wsDir, "failed")
				return &WorkflowExecutionResult{
					Success:     false,
					ErrorReason: err.Error(),
					Duration:    time.Since(start),
				}, err
			}
			stepsRun++
			wf.Steps[stepIndex].Status = "completed"
			_ = d.SaveWorkflow(wsDir, wf)

			if step.AllowDynamicAdjustment && d.evaluator != nil {
				d.applyDynamicAdjustment(ctx, wsDir, wf, step, &mu, opts.StepNotify)
			}
		}
	}

	// 收集最终交付物并汇总到工作区根目录 output.md
	finalContent, _ := d.resolveFinalDeliverable(wsDir, wf)
	if len(finalContent) == 0 {
		// 容错：若未提取到内容，读取 template.md
		finalContent, _ = os.ReadFile(filepath.Join(wsDir, "template.md"))
	}
	rootOut := filepath.Join(wsDir, "output.md")
	_ = os.WriteFile(rootOut, finalContent, 0o644)

	wf.Status = "completed"
	_ = d.SaveWorkflow(wsDir, wf)
	_ = d.updateWorkspaceStatus(wsDir, "completed")

	bb, _ := d.LoadBlackboard(wsDir)
	var finalFields map[string]string
	if bb != nil {
		finalFields = bb.Fields
	}

	return &WorkflowExecutionResult{
		Success:      true,
		StepsRun:     stepsRun,
		Deliverable:  rootOut,
		FinalContent: finalContent,
		Duration:     time.Since(start),
		Fields:       finalFields,
	}, nil
}

func (d *WorkflowDriver) runOneStep(
	ctx context.Context,
	wsDir string,
	step StepSpec,
	backend agent.Backend,
	notify func(StepEvent),
	mu *sync.Mutex,
	producedFiles *[]string,
) error {
	if mu == nil {
		var localMu sync.Mutex
		mu = &localMu
	}
	if producedFiles == nil {
		var localFiles []string
		producedFiles = &localFiles
	}

	subDir := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace))
	_ = os.MkdirAll(subDir, 0o755)

	deliverableName := "output.md"
	if step.Mode == "plan" || step.Mode == "replan" {
		deliverableName = "output_plan.md"
	} else if step.Mode == "review_plan" {
		deliverableName = "review_notes.md"
	}

	if notify != nil {
		notify(StepEvent{
			StepID: step.ID,
			Role:   step.Role,
			Title:  step.Title,
			Mode:   step.Mode,
			Status: "started",
		})
	}
	streamer := stream.FromContext(ctx)
	roleTitle := FriendlyRoleName(step.Role)
	if streamer != nil {
		streamer.EmitPhase("step_executing", fmt.Sprintf("%s正在进行：%s...", roleTitle, step.Title), step.Role, 0, 0, 0)
	}
	stepStart := time.Now()

	// 若当前步骤为 leader 审核定稿，且依赖前序专员稿件，先预置 baseline 草稿，供审查与润色
	if step.Role == "leader" && deliverableName == "output.md" {
		targetFile := filepath.Join(subDir, deliverableName)
		if _, err := os.Stat(targetFile); os.IsNotExist(err) && len(step.Needs) > 0 {
			srcFile := filepath.Join(wsDir, filepath.FromSlash(step.Needs[0]))
			if srcData, err := os.ReadFile(srcFile); err == nil && len(srcData) > 0 {
				_ = os.WriteFile(targetFile, srcData, 0o644)
			}
		}
	}

	// 1. 读取 TASK.md 与 AGENTS.md
	taskData, _ := os.ReadFile(filepath.Join(subDir, "TASK.md"))
	agentsData, _ := os.ReadFile(filepath.Join(subDir, "AGENTS.md"))

	targets := []string{deliverableName}
	if step.Role == "organizer" {
		targets = append(targets, "fields.json")
	}

	req := agent.Request{
		WorkDir:      subDir,
		Prompt:       string(taskData),
		Instructions: string(agentsData),
		Targets:      targets,
		Timeout:      30 * time.Minute,
	}

	res, err := backend.Execute(ctx, req)
	duration := time.Since(stepStart).Milliseconds()

	if err != nil {
		if streamer != nil {
			streamer.EmitPhase("step_failed", fmt.Sprintf("%s执行异常：%s", roleTitle, err.Error()), step.Role, 0, 0, 0)
		}
		if notify != nil {
			notify(StepEvent{
				StepID:     step.ID,
				Role:       step.Role,
				Title:      step.Title,
				Mode:       step.Mode,
				Status:     "failed",
				DurationMs: duration,
				Error:      err.Error(),
			})
		}
		return fmt.Errorf("step %s (%s) backend execute: %w", step.ID, step.Role, err)
	}

	if res != nil && !res.Success {
		errMsg := res.ErrorReason
		if streamer != nil {
			streamer.EmitPhase("step_failed", fmt.Sprintf("%s执行未成功：%s", roleTitle, errMsg), step.Role, 0, 0, 0)
		}
		if notify != nil {
			notify(StepEvent{
				StepID:     step.ID,
				Role:       step.Role,
				Title:      step.Title,
				Mode:       step.Mode,
				Status:     "failed",
				DurationMs: duration,
				Error:      errMsg,
			})
		}
		return fmt.Errorf("step %s failed: %s", step.ID, errMsg)
	}

	// 2. 确认 deliverable 文件落盘
	targetPath := filepath.Join(subDir, deliverableName)
	if _, err := os.Stat(targetPath); err != nil {
		// 若后端补丁在内存中返回，写入磁盘
		if res != nil && len(res.Patches[deliverableName]) > 0 {
			_ = os.WriteFile(targetPath, res.Patches[deliverableName], 0o644)
		} else {
			return fmt.Errorf("step %s missing target deliverable %s", step.ID, deliverableName)
		}
	}

	// 3. 检查是否有 fields.json 并合并入 blackboard.json
	fieldsPath := filepath.Join(subDir, "fields.json")
	if fData, err := os.ReadFile(fieldsPath); err == nil {
		var extractedFields map[string]string
		if err := json.Unmarshal(fData, &extractedFields); err == nil {
			mu.Lock()
			if bb, err := d.LoadBlackboard(wsDir); err == nil {
				if bb.Fields == nil {
					bb.Fields = make(map[string]string)
				}
				for k, v := range extractedFields {
					bb.Fields[k] = v
				}
				_ = d.SaveBlackboard(wsDir, bb)
			}
			mu.Unlock()
		}
	}

	// 4. 记录日志
	logEntry := fmt.Sprintf("[%s] 步骤 %s (%s) 执行完成，耗时 %dms\n",
		time.Now().Format(time.RFC3339), step.ID, step.Role, duration)
	_ = os.WriteFile(filepath.Join(subDir, "execution.log"), []byte(logEntry), 0o644)

	mu.Lock()
	*producedFiles = append(*producedFiles, filepath.ToSlash(filepath.Join(step.SubWorkspace, deliverableName)))
	mu.Unlock()

	if streamer != nil {
		streamer.EmitPhase("step_completed", fmt.Sprintf("%s完成任务：%s (耗时 %dms)", roleTitle, step.Title, duration), step.Role, 0, 0, 0)
		if stepData, err := os.ReadFile(targetPath); err == nil && len(stepData) > 0 {
			streamer.EmitDetail("step_output", step.ID, step.Role, fmt.Sprintf("%s阶段性产出: %s", roleTitle, step.Title), string(stepData), true)
		}
	}

	if notify != nil {
		notify(StepEvent{
			StepID:     step.ID,
			Role:       step.Role,
			Title:      step.Title,
			Mode:       step.Mode,
			Status:     "completed",
			DurationMs: duration,
		})
	}

	return nil
}

func (d *WorkflowDriver) applyDynamicAdjustment(
	ctx context.Context,
	wsDir string,
	wf *WorkflowPlan,
	step StepSpec,
	mu *sync.Mutex,
	notify func(StepEvent),
) {
	deliverableName := "output.md"
	if step.Mode == "plan" || step.Mode == "replan" {
		deliverableName = "output_plan.md"
	}
	outPath := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), deliverableName)
	outContent, _ := os.ReadFile(outPath)
	bb, _ := d.LoadBlackboard(wsDir)

	evalRes, err := d.evaluator.EvaluateStepOutput(wsDir, step, string(outContent), bb)
	if err != nil || evalRes == nil {
		return
	}

	mu.Lock()
	defer mu.Unlock()

	// 1. 合并新推断的基准假设
	if len(evalRes.NewAssumptions) > 0 && bb != nil {
		for _, na := range evalRes.NewAssumptions {
			exists := false
			for _, ea := range bb.Assumptions {
				if ea.Topic == na.Topic {
					exists = true
					break
				}
			}
			if !exists {
				bb.Assumptions = append(bb.Assumptions, na)
			}
		}
		_ = d.SaveBlackboard(wsDir, bb)
	}

	// 2. 处理动态跳步
	if len(evalRes.StepsToSkip) > 0 {
		for _, skipID := range evalRes.StepsToSkip {
			target := strings.ToLower(skipID)
			for idx := range wf.Steps {
				curID := strings.ToLower(wf.Steps[idx].ID)
				curRole := strings.ToLower(wf.Steps[idx].Role)
				isMatch := (curID == target || curRole == target || strings.Contains(curID, target) || strings.Contains(target, curID))
				if isMatch && (wf.Steps[idx].Status == "pending" || wf.Steps[idx].Status == "") {
					wf.Steps[idx].Status = "skipped"
					if notify != nil {
						notify(StepEvent{
							StepID: wf.Steps[idx].ID,
							Role:   wf.Steps[idx].Role,
							Title:  wf.Steps[idx].Title,
							Status: "skipped",
						})
					}
				}
			}
		}
		_ = d.SaveWorkflow(wsDir, wf)
	}

	// 3. 处理动态加步
	if len(evalRes.DynamicStepsToAdd) > 0 {
		squadCfg, _ := d.loadSquadConfig(wsDir)
		tmpl, _ := os.ReadFile(filepath.Join(wsDir, "template.md"))

		for _, addStep := range evalRes.DynamicStepsToAdd {
			alreadyExists := false
			for _, s := range wf.Steps {
				if s.ID == addStep.ID {
					alreadyExists = true
					break
				}
			}
			if alreadyExists {
				continue
			}

			if addStep.SubWorkspace == "" {
				addStep.SubWorkspace = fmt.Sprintf("tasks/%02d_%s", len(wf.Steps)+1, addStep.Role)
			}
			if addStep.Status == "" {
				addStep.Status = "pending"
			}
			if addStep.Mode == "" {
				addStep.Mode = "execute"
			}

			// 构建子沙箱
			_ = d.engine.buildSubWorkspaces(wsDir, WorkflowPlan{Steps: []StepSpec{addStep}}, squadCfg, string(tmpl))

			// 插入在 writer 或 leader 之前，确保专家成果在成文前产生
			insertIdx := len(wf.Steps)
			for idx, s := range wf.Steps {
				if s.Role == "writer" || s.Role == "leader" {
					insertIdx = idx
					break
				}
			}

			newSteps := make([]StepSpec, 0, len(wf.Steps)+1)
			newSteps = append(newSteps, wf.Steps[:insertIdx]...)
			newSteps = append(newSteps, addStep)
			newSteps = append(newSteps, wf.Steps[insertIdx:]...)
			wf.Steps = newSteps

			streamer := stream.FromContext(ctx)
			if streamer != nil {
				streamer.EmitPhase("dynamic_adjustment", fmt.Sprintf("单步评估发现缺口，动态追加任务步骤：%s (%s)", addStep.Title, FriendlyRoleName(addStep.Role)), addStep.Role, 0, 0, 0)
			}
		}
		_ = d.SaveWorkflow(wsDir, wf)
	}
}

func (d *WorkflowDriver) resolveFinalDeliverable(wsDir string, wf *WorkflowPlan) ([]byte, string) {
	// 查找优先级：leader 的 output.md > writer 的 output.md > 最后一个步骤的 output.md
	for i := len(wf.Steps) - 1; i >= 0; i-- {
		step := wf.Steps[i]
		if step.Role == "leader" && (step.Mode == "execute" || step.Mode == "") {
			p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
			if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
				return data, p
			}
		}
	}

	for i := len(wf.Steps) - 1; i >= 0; i-- {
		step := wf.Steps[i]
		if step.Role == "writer" && (step.Mode == "execute" || step.Mode == "") {
			p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
			if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
				return data, p
			}
		}
	}

	for i := len(wf.Steps) - 1; i >= 0; i-- {
		step := wf.Steps[i]
		p := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace), "output.md")
		if data, err := os.ReadFile(p); err == nil && len(strings.TrimSpace(string(data))) > 0 {
			return data, p
		}
	}

	return nil, ""
}

func (d *WorkflowDriver) loadSquadConfig(wsDir string) (SquadConfig, error) {
	p := filepath.Join(wsDir, "squad.json")
	var sc SquadConfig
	data, err := os.ReadFile(p)
	if err != nil {
		return sc, err
	}
	err = json.Unmarshal(data, &sc)
	return sc, err
}

func (d *WorkflowDriver) updateWorkspaceStatus(wsDir string, status string) error {
	p := filepath.Join(wsDir, "workspace.json")
	data, err := os.ReadFile(p)
	if err != nil {
		return err
	}
	var meta WorkspaceMetadata
	if err := json.Unmarshal(data, &meta); err != nil {
		return err
	}
	meta.Status = status
	meta.UpdatedAt = time.Now()
	if status == "completed" {
		now := time.Now()
		meta.CompletedAt = &now
	}
	return writeJSON(p, meta)
}