package engine

import (
	"bytes"
	"encoding/json"

	"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
	"git.yuntuo.hk/k2safe/orbit/weave/definition"
	weave "git.yuntuo.hk/k2safe/orbit/weave/model"
)

type planner struct {
	program   *Program
	states    map[string]State
	decision  Decision
	limit     int
	allowNew  bool
	verifying int
	visited   map[string]bool
}

// Plan 只计算本轮最多limit个步骤状态变更（1至32），不修改快照或重新执行已完成业务。
// 错误时丢弃全部本轮决策，存储适配层须保留原快照并记录核查原因。
func (p *Program) Plan(snapshot Snapshot, limit int) (Decision, error) {
	if limit < 1 || limit > 32 || snapshot.WorkflowID.Validate() != nil || !validBusinessKey(snapshot.BusinessKey) {
		return Decision{}, problem("", "INVALID_SNAPSHOT")
	}
	if err := definition.ValidateData(p.spec.InputType, snapshot.Input); err != nil {
		return Decision{}, problem("", "WORKFLOW_INPUT_INVALID")
	}
	if len(snapshot.Steps) > len(p.nodes) {
		return Decision{}, problem("", "INVALID_SNAPSHOT")
	}
	state := &planner{program: p, states: make(map[string]State, len(p.nodes)), visited: map[string]bool{}, limit: limit, allowNew: snapshot.AllowNew, decision: Decision{Changes: []Change{}}}
	for key, row := range snapshot.Steps {
		node, ok := p.nodes[key]
		if !ok {
			return Decision{}, problem(key, "UNKNOWN_STEP")
		}
		if row.Status == weave.StepFailed || row.Status == weave.StepNeedsReview || row.Status == weave.StepCancelled {
			return Decision{}, problem(key, "STEP_REQUIRES_REVIEW")
		}
		switch row.Status {
		case weave.StepBlocked, weave.StepSkipped:
			if len(row.Input) != 0 || len(row.Output) != 0 {
				return Decision{}, problem(key, "INVALID_STEP_STATE")
			}
		case weave.StepActive, weave.StepWaiting, weave.StepSucceeded:
			if err := definition.ValidateData(definition.Type{Kind: "JSON"}, row.Input); err != nil {
				return Decision{}, problem(key, "INVALID_STEP_INPUT")
			}
			if node.Kind == "TASK" || node.Kind == "POLL" {
				task := node.Task
				if node.Poll != nil {
					task = &node.Poll.Task
				}
				if null(row.Input) || definition.ValidateData(task.InputType, row.Input) != nil {
					return Decision{}, problem(key, "INVALID_STEP_INPUT")
				}
			}
			if row.Status == weave.StepSucceeded {
				if err := validateOutput(node, row.Output); err != nil {
					return Decision{}, problem(key, "STEP_OUTPUT_INVALID")
				}
			} else if len(row.Output) != 0 {
				return Decision{}, problem(key, "INVALID_STEP_STATE")
			}
		default:
			return Decision{}, problem(key, "INVALID_STEP_STATE")
		}
		state.states[key] = row
	}
	workflow, err := jsonwire.Marshal(map[string]string{"id": string(snapshot.WorkflowID), "business_key": snapshot.BusinessKey})
	if err != nil {
		return Decision{}, err
	}
	complete, output, err := state.block(p.spec.Body, values{"input": copyJSON(snapshot.Input), "workflow": workflow})
	if err != nil {
		return Decision{}, err
	}
	for key, row := range snapshot.Steps {
		if (row.Status == weave.StepActive || row.Status == weave.StepWaiting || row.Status == weave.StepSucceeded) && !state.visited[key] {
			return Decision{}, problem(key, "UNREACHABLE_EXECUTION")
		}
	}
	state.decision.Complete = complete
	state.decision.Output = copyJSON(output)
	return state.decision, nil
}
func null(raw []byte) bool { return bytes.Equal(bytes.TrimSpace(raw), []byte("null")) }
func outputType(node definition.Step) definition.Type {
	switch node.Kind {
	case "TASK":
		return node.Task.ResultType
	case "POLL":
		return node.Poll.Task.ResultType
	case "SIGNAL":
		return node.Signal.PayloadType
	case "TIMER":
		return definition.Type{Kind: "NULL"}
	case "BRANCH":
		return node.Branch.Default.OutputType
	case "PARALLEL":
		t := definition.Type{Kind: "OBJECT", Fields: map[string]definition.Type{}}
		for _, b := range node.Parallel.Branches {
			t.Fields[b.Name] = b.Body.OutputType
		}
		return t
	}
	return definition.Type{}
}
func validateOutput(node definition.Step, raw json.RawMessage) error {
	if (node.Kind == "TASK" || node.Kind == "POLL") && null(raw) {
		return problem(node.Key, "STEP_OUTPUT_INVALID")
	}
	if node.Poll != nil {
		outcome, err := classifyPoll(*node.Poll, raw)
		if err != nil {
			return err
		}
		if outcome != PollSucceeded {
			return problem(node.Key, "POLL_RESULT_NOT_SUCCESS")
		}
	}
	return definition.ValidateData(outputType(node), raw)
}
func (s *planner) add(change Change) bool {
	if len(s.decision.Changes) >= s.limit {
		s.decision.More = true
		return false
	}
	s.decision.Changes = append(s.decision.Changes, change)
	s.states[change.Key] = change.State
	return true
}
func (s *planner) block(block definition.Block, parent values) (bool, json.RawMessage, error) {
	env := copyValues(parent)
	for _, node := range block.Steps {
		complete, output, err := s.step(node, env)
		if err != nil {
			return false, nil, err
		}
		if !complete {
			return false, nil, nil
		}
		env[node.Key] = output
	}
	output, err := resolve(block.Output, block.OutputType, env)
	if err != nil {
		return false, nil, problem("", "BLOCK_OUTPUT_INVALID")
	}
	return true, output, nil
}
func (s *planner) step(node definition.Step, env values) (bool, json.RawMessage, error) {
	s.visited[node.Key] = true
	row, ok := s.states[node.Key]
	if !ok {
		row.Status = weave.StepBlocked
	}
	completed := row.Status == weave.StepSucceeded
	if s.verifying > 0 && !completed {
		return false, nil, problem(node.Key, "INCOMPLETE_COMPOSITE_CHILD")
	}
	if completed {
		if node.Kind != "BRANCH" && node.Kind != "PARALLEL" {
			return true, copyJSON(row.Output), nil
		}
		s.verifying++
		defer func() { s.verifying-- }()
	}
	if row.Status == weave.StepSkipped {
		return false, nil, problem(node.Key, "SELECTED_STEP_WAS_SKIPPED")
	}
	if row.Status == weave.StepBlocked {
		if !s.allowNew || s.decision.More {
			return false, nil, nil
		}
		change := Change{Key: node.Key, Kind: node.Kind, Operation: "ACTIVATE", State: State{Status: weave.StepActive}}
		switch node.Kind {
		case "TASK", "POLL":
			task := node.Task
			if node.Poll != nil {
				task = &node.Poll.Task
			}
			prepared, err := prepareTask(*task, env)
			if err != nil {
				return false, nil, problem(node.Key, "TASK_INPUT_INVALID")
			}
			change.State.Input = copyJSON(prepared.Payload)
			change.Task = &prepared
			if node.Poll != nil && node.Poll.FinalSignal != nil {
				wait, err := prepareSignal(*node.Poll.FinalSignal, "POLL", env)
				if err != nil {
					return false, nil, problem(node.Key, "SIGNAL_CORRELATION_INVALID")
				}
				change.Wait = wait
			}
		case "SIGNAL":
			wait, err := prepareSignal(*node.Signal, "SIGNAL", env)
			if err != nil {
				return false, nil, problem(node.Key, "SIGNAL_CORRELATION_INVALID")
			}
			change.Wait = wait
			change.State.Status = weave.StepWaiting
			change.State.Input, _ = jsonwire.Marshal(map[string]string{"name": wait.Name, "payload_version": wait.PayloadVersion, "correlation": wait.Correlation})
		case "TIMER":
			change.Wait = &WaitCommand{Kind: "TIMER", TimeoutMS: node.Timer.DelayMS}
			change.State.Status = weave.StepWaiting
			change.State.Input, _ = jsonwire.Marshal(map[string]uint64{"delay_ms": node.Timer.DelayMS})
		case "BRANCH":
			selected := -1
			for i, path := range node.Branch.Cases {
				matches, err := compare(path.When, env)
				if err != nil {
					return false, nil, problem(node.Key, "BRANCH_CONDITION_INVALID")
				}
				if matches {
					selected = i
					break
				}
			}
			change.State.Input, _ = jsonwire.Marshal(map[string]int{"case": selected})
		case "PARALLEL":
			change.State.Input = json.RawMessage(`{}`)
		default:
			return false, nil, problem(node.Key, "UNSUPPORTED_NODE")
		}
		if !s.add(change) {
			return false, nil, nil
		}
		row = change.State
	}
	if node.Kind != "BRANCH" && node.Kind != "PARALLEL" {
		return false, nil, nil
	}
	var output json.RawMessage
	switch node.Kind {
	case "BRANCH":
		contract := definition.Type{Kind: "OBJECT", Fields: map[string]definition.Type{"case": {Kind: "NUMBER"}}}
		var selection struct {
			Case int `json:"case"`
		}
		if definition.ValidateData(contract, row.Input) != nil || json.Unmarshal(row.Input, &selection) != nil || selection.Case < -1 || selection.Case >= len(node.Branch.Cases) {
			return false, nil, problem(node.Key, "BRANCH_DECISION_INVALID")
		}
		// 已冻结的选择不重算；未选择路径分批标记SKIPPED，不创建业务激活。
		for i, path := range node.Branch.Cases {
			if i != selection.Case {
				if err := s.skip(path.Body); err != nil {
					return false, nil, err
				}
			}
		}
		body := *node.Branch.Default
		if selection.Case >= 0 {
			if err := s.skip(*node.Branch.Default); err != nil {
				return false, nil, err
			}
			body = node.Branch.Cases[selection.Case].Body
		}
		// 预算只阻止新增变更；已有选中路径仍须遍历，避免遗漏并行分支的在途工作。
		done, raw, err := s.block(body, env)
		if err != nil {
			return false, nil, err
		}
		if !done {
			return false, nil, nil
		}
		output = raw
	case "PARALLEL":
		if definition.ValidateData(definition.Type{Kind: "OBJECT"}, row.Input) != nil {
			return false, nil, problem(node.Key, "INVALID_PARALLEL_STATE")
		}
		result := make(map[string]json.RawMessage, len(node.Parallel.Branches))
		all := true
		for _, branch := range node.Parallel.Branches {
			done, raw, err := s.block(branch.Body, env)
			if err != nil {
				return false, nil, err
			}
			if !done {
				all = false
			} else {
				result[branch.Name] = raw
			}
		}
		if !all {
			return false, nil, nil
		}
		var err error
		output, err = jsonwire.Marshal(result)
		if err != nil {
			return false, nil, problem(node.Key, "STEP_OUTPUT_INVALID")
		}
	}
	if validateOutput(node, output) != nil {
		return false, nil, problem(node.Key, "STEP_OUTPUT_INVALID")
	}
	if completed {
		original, err := jsonvalue.Canonical(row.Output)
		if err != nil {
			return false, nil, problem(node.Key, "STEP_OUTPUT_INVALID")
		}
		current, err := jsonvalue.Canonical(output)
		if err != nil || !bytes.Equal(original, current) {
			return false, nil, problem(node.Key, "COMPOSITE_OUTPUT_MISMATCH")
		}
		return true, copyJSON(row.Output), nil
	}
	row.Status = weave.StepSucceeded
	row.Output = copyJSON(output)
	if !s.add(Change{Key: node.Key, Kind: node.Kind, Operation: "COMPLETE", State: row}) {
		return false, nil, nil
	}
	return true, output, nil
}
func (s *planner) skip(block definition.Block) error {
	for _, node := range block.Steps {
		s.visited[node.Key] = true
		row, ok := s.states[node.Key]
		if ok && row.Status != weave.StepBlocked && row.Status != weave.StepSkipped {
			return problem(node.Key, "UNSELECTED_STEP_HAS_EXECUTION")
		}
		if !ok || row.Status == weave.StepBlocked {
			if s.verifying > 0 {
				return problem(node.Key, "INCOMPLETE_COMPOSITE_CHILD")
			}
			if !s.add(Change{Key: node.Key, Kind: node.Kind, Operation: "SKIP", State: State{Status: weave.StepSkipped}}) {
				return nil
			}
		}
		if node.Branch != nil {
			for _, c := range node.Branch.Cases {
				if err := s.skip(c.Body); err != nil {
					return err
				}
			}
			if err := s.skip(*node.Branch.Default); err != nil {
				return err
			}
		}
		if node.Parallel != nil {
			for _, b := range node.Parallel.Branches {
				if err := s.skip(b.Body); err != nil {
					return err
				}
			}
		}
	}
	return nil
}