package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/singlerequesttemplate"
)

const (
	singleRequestPlanPrompt  = "Inspect the workspace before planning. Use only the supplied read-only workspace tools and base the plan on the returned evidence. After at least one successful inspection, return exactly one JSON object with a non-empty one-line goal string, a steps array containing 2 to 6 non-empty one-line strings, and a verification array containing 1 to 3 non-empty one-line strings. Do not include Markdown bullets, headings, or template text; Edge owns all Plan document formatting."
	singleRequestPlanStageID = "plan"
)

var errSingleRequestPlanStage = errors.New("single-request plan stage: failed")

type singleRequestPlanStage struct {
	provider *singleRequestProviderStage
	bridge   *singleRequestWorkToolBridge
}

func newSingleRequestPlanStage(provider *singleRequestProviderStage, bridge *singleRequestWorkToolBridge) *singleRequestPlanStage {
	return &singleRequestPlanStage{provider: provider, bridge: bridge}
}

type singleRequestPlanStageRequest struct {
	RequestID        string
	Task             string
	StageBinding     edgeservice.SingleRequestStageBinding
	Limits           edgeservice.SingleRequestLimits
	NodeRef          string
	SessionID        string
	UsageAttribution string
	Sequence         uint64
	Quality          *singleRequestQualityGate
}

type singleRequestPlanFields struct {
	Goal         string   `json:"goal"`
	Steps        []string `json:"steps"`
	Verification []string `json:"verification"`
}

type singleRequestPlanProviderResponse struct {
	call *singleRequestReviewProviderToolCall
	plan *singleRequestPlanFields
}

func (v *singleRequestPlanFields) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "goal", "steps", "verification"); err != nil {
		return err
	}
	type alias singleRequestPlanFields
	var decoded alias
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*v = singleRequestPlanFields(decoded)
	return nil
}

func singleRequestPlanResponseFormat() *singleRequestProviderResponseFormat {
	return &singleRequestProviderResponseFormat{
		Type: "json_schema",
		JSONSchema: singleRequestProviderResponseJSONSchema{
			Name:   "single_request_plan",
			Strict: true,
			Schema: singleRequestProviderOutputSchema{
				Type: "object",
				Properties: map[string]singleRequestProviderOutputProperty{
					"goal": {Type: "string", MinLength: 1},
					"steps": {
						Type: "array", Items: &singleRequestProviderOutputProperty{Type: "string", MinLength: 1}, MinItems: 2, MaxItems: 6,
					},
					"verification": {
						Type: "array", Items: &singleRequestProviderOutputProperty{Type: "string", MinLength: 1}, MinItems: 1, MaxItems: 3,
					},
				},
				Required:             []string{"goal", "steps", "verification"},
				AdditionalProperties: false,
			},
		},
	}
}

func (s *singleRequestPlanStage) run(ctx context.Context, req singleRequestPlanStageRequest, ctrl edgeservice.SingleRequestController) ([]byte, error) {
	quality := singleRequestQualityGateOrNew(req.Quality)
	if s == nil || s.provider == nil || s.provider.service == nil || s.bridge == nil || ctrl == nil || req.RequestID == "" || req.Task == "" || req.Sequence == 0 || req.StageBinding.Dispatch == nil {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	if req.StageBinding.Options["reasoning_effort"] != "high" {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	binding := ctrl.Binding()
	if binding == nil || binding.Workspace == nil || binding.Workspace.NodeID == "" || binding.Templates.Plan == "" || req.NodeRef != binding.Workspace.NodeID {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	if err := singlerequesttemplate.ValidatePlanTemplate(binding.Templates.Plan); err != nil {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	tools, err := singleRequestPlanTools(binding.Workspace)
	if err != nil {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	sequence := req.Sequence
	if err := ctrl.SubmitEnvelope(edgeservice.SingleRequestEnvelope{RequestID: req.RequestID, Sequence: sequence, Stage: edgeservice.SingleRequestStatePlanning}); err != nil {
		return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
	}
	messages := []chatMessage{
		{Role: "system", Content: singleRequestPlanPrompt},
		{Role: "user", Content: strings.TrimSpace(req.Task)},
	}
	inspectionEligible := false
	toolAttempts := 0
	for {
		response, err := s.submit(ctx, req, messages, tools, inspectionEligible)
		if err != nil {
			return nil, quality.reclassify(err, errSingleRequestPlanStage)
		}
		if response.plan != nil {
			if !inspectionEligible {
				return nil, quality.malformed(errSingleRequestPlanStage)
			}
			content, err := singlerequesttemplate.RenderPlan(binding.Templates.Plan, singlerequesttemplate.PlanFields{
				Goal: response.plan.Goal, Steps: response.plan.Steps, Verification: response.plan.Verification,
			}, req.Limits.MaxOutputBytes)
			if err != nil {
				return nil, quality.malformed(errSingleRequestPlanStage)
			}
			if err := ctrl.WriteInternalArtifact(ctx, edgeservice.SingleRequestArtifactPlan, content); err != nil {
				return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
			}
			return content, nil
		}
		if response.call == nil || toolAttempts >= req.Limits.MaxToolIterations {
			if response.call != nil {
				return nil, quality.budget(errSingleRequestPlanStage)
			}
			return nil, quality.malformed(errSingleRequestPlanStage)
		}
		if response.call.Function.Name != edgeservice.InternalWorkspaceToolRead && response.call.Function.Name != edgeservice.InternalWorkspaceToolList {
			return nil, quality.malformed(errSingleRequestPlanStage)
		}
		arguments, err := decodeSingleRequestWorkToolArguments(response.call.Function.Arguments)
		if err != nil {
			return nil, quality.malformed(errSingleRequestPlanStage)
		}
		arguments = normalizeSingleRequestProviderToolArguments(response.call.Function.Name, arguments)
		key := singleRequestWorkToolKey{requestID: req.RequestID, stageID: singleRequestPlanStageID, toolCallID: response.call.ID}
		call := &edgeservice.InternalWorkspaceToolCall{RequestID: req.RequestID, StageID: key.stageID, ToolCallID: key.toolCallID, Name: response.call.Function.Name, Arguments: arguments}
		if err := edgeservice.ValidateInternalWorkspaceToolCall(call); err != nil {
			return nil, quality.malformed(errSingleRequestPlanStage)
		}
		toolAttempts++
		resultCh, err := s.bridge.register(key)
		if err != nil {
			return nil, quality.internalTool(errSingleRequestPlanStage)
		}
		sequence++
		if err := ctrl.SubmitEnvelope(edgeservice.SingleRequestEnvelope{RequestID: req.RequestID, Sequence: sequence, Stage: edgeservice.SingleRequestStateInternalTool, SavedStage: edgeservice.SingleRequestStatePlanning, ToolCall: call}); err != nil {
			s.bridge.unregister(key)
			return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
		}
		result, err := s.bridge.wait(ctx, key, resultCh)
		if err != nil {
			return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
		}
		if err := quality.observeToolCycle(singleRequestPlanStageID, response.call.Function.Name, arguments, result, errSingleRequestPlanStage); err != nil {
			return nil, err
		}
		inspectionEligible = result.Status == "success" && result.ErrorCode == ""
		messages = append(messages,
			chatMessage{Role: "assistant", ToolCalls: []any{response.call.asChatToolCall()}},
			chatMessage{Role: "tool", ToolCallID: response.call.ID, ToolName: response.call.Function.Name, Content: singleRequestWorkToolResultContent(result, req.Limits.MaxOutputBytes)},
		)
		sequence++
		if err := ctrl.SubmitEnvelope(edgeservice.SingleRequestEnvelope{RequestID: req.RequestID, Sequence: sequence, Stage: edgeservice.SingleRequestStatePlanning, SavedStage: edgeservice.SingleRequestStatePlanning}); err != nil {
			return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
		}
	}
}

func singleRequestPlanTools(workspace *edgeservice.SingleRequestWorkspaceBinding) ([]any, error) {
	if workspace == nil {
		return nil, errSingleRequestPlanStage
	}
	has := func(operation string) bool {
		for _, candidate := range workspace.OperationIDs {
			if candidate == operation {
				return true
			}
		}
		return false
	}
	pathProperty := map[string]any{"type": "string", "minLength": 1, "description": singleRequestCanonicalRelativePathDescription}
	parameters := map[string]any{"type": "object", "additionalProperties": false, "required": []string{"relative_path"}, "properties": map[string]any{"relative_path": pathProperty}}
	tools := make([]any, 0, 2)
	if has("read") {
		tools = append(tools, singleRequestWorkToolSchema(edgeservice.InternalWorkspaceToolRead, parameters))
	}
	if has("list") {
		tools = append(tools, singleRequestWorkToolSchema(edgeservice.InternalWorkspaceToolList, parameters))
	}
	if len(tools) == 0 {
		return nil, errSingleRequestPlanStage
	}
	return tools, nil
}

func (s *singleRequestPlanStage) submit(ctx context.Context, req singleRequestPlanStageRequest, messages []chatMessage, tools []any, inspectionEligible bool) (*singleRequestPlanProviderResponse, error) {
	quality := singleRequestQualityGateOrNew(req.Quality)
	dispatch := req.StageBinding.Dispatch
	stageCtx, cancel := providerStageContext(ctx, req.Limits.StageTimeoutMS)
	defer cancel()
	poolReq := edgeservice.ProviderPoolDispatchRequest{
		Run: edgeservice.SubmitRunRequest{NodeRef: req.NodeRef, ModelGroupKey: dispatch.ModelGroupKey, ProviderID: dispatch.ProviderID, UsageAttribution: req.UsageAttribution, SessionID: req.SessionID, TimeoutSec: dispatch.TimeoutSec, MaxQueue: dispatch.MaxQueue, QueueTimeoutMS: dispatch.QueueTimeoutMS, ProviderPool: true},
		Tunnel: edgeservice.SubmitProviderTunnelRequest{CredentialBinding: dispatch.CredentialBindingSnapshot(), NodeRef: req.NodeRef, ModelGroupKey: dispatch.ModelGroupKey, ProviderID: dispatch.ProviderID, UsageAttribution: req.UsageAttribution, Adapter: "openai_compat", Target: dispatch.UpstreamModel, SessionID: req.SessionID, Method: http.MethodPost, Path: "/v1/chat/completions", Stream: false, TimeoutSec: dispatch.TimeoutSec, MaxQueue: dispatch.MaxQueue, QueueTimeoutMS: dispatch.QueueTimeoutMS, ProviderPool: true, BuildBody: func(target string) ([]byte, error) {
			return buildSingleRequestPlanBody(messages, req.StageBinding.Options, tools, target, inspectionEligible)
		}},
	}
	var responseFormat *singleRequestProviderResponseFormat
	if inspectionEligible {
		responseFormat = singleRequestPlanResponseFormat()
	}
	requirements := singleRequestProviderRequirements(req.StageBinding.Options, tools, responseFormat)
	poolReq.AcceptCandidate = singleRequestProviderCandidatePredicate(dispatch, requirements)
	poolReq.PrepareProtocolTunnel = singleRequestProviderTunnelPreparer(requirements, func(target string) ([]byte, error) {
		return buildSingleRequestPlanBody(messages, req.StageBinding.Options, tools, target, inspectionEligible)
	})
	result, err := s.provider.service.SubmitProviderPool(stageCtx, poolReq)
	if err != nil || result == nil || result.Tunnel == nil {
		return nil, quality.providerFailure(stageCtx, err, errSingleRequestPlanStage)
	}
	if !providerStageDispatchMatches(result, dispatch) {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	defer result.Tunnel.Close()
	body, err := collectProviderStageFrames(stageCtx, result.Tunnel.Stream().Frames, req.Limits.MaxOutputBytes)
	if err != nil {
		return nil, quality.providerFailure(stageCtx, err, errSingleRequestPlanStage)
	}
	body, err = normalizeSingleRequestProviderResponse(body, result.DispatchInfo)
	if err != nil {
		return nil, quality.providerFailure(stageCtx, err, errSingleRequestPlanStage)
	}
	response, err := decodeSingleRequestPlanProviderResponse(body, req.Limits.MaxOutputBytes)
	if err != nil {
		if !errors.Is(err, errProviderStageOutputLimit) && !errors.Is(err, errProviderStageContextLimit) {
			return nil, quality.malformed(errSingleRequestPlanStage)
		}
		return nil, quality.providerFailure(stageCtx, err, errSingleRequestPlanStage)
	}
	return response, nil
}

func buildSingleRequestPlanBody(messages []chatMessage, options map[string]any, tools []any, target string, inspectionEligible bool) ([]byte, error) {
	if target == "" || len(messages) == 0 || len(tools) == 0 || options["reasoning_effort"] != "high" {
		return nil, errSingleRequestPlanStage
	}
	toolChoice := "required"
	if inspectionEligible {
		toolChoice = "auto"
	}
	body := map[string]any{"model": target, "messages": messages, "tools": tools, "tool_choice": toolChoice, "parallel_tool_calls": false, "stream": false}
	if inspectionEligible {
		body["response_format"] = singleRequestPlanResponseFormat()
	}
	for key, value := range options {
		folded := strings.ToLower(key)
		if isSingleRequestWorkReservedOption(folded) && key != folded {
			return nil, errSingleRequestPlanStage
		}
		if isSingleRequestWorkReservedOption(folded) {
			continue
		}
		body[key] = value
	}
	body["reasoning_effort"] = "high"
	return json.Marshal(body)
}

func decodeSingleRequestPlanProviderResponse(body []byte, maximum int) (*singleRequestPlanProviderResponse, error) {
	if len(body) == 0 || len(body) > maximum || validateSingleRequestJSON(body) != nil {
		return nil, errSingleRequestPlanStage
	}
	var envelope singleRequestReviewProviderEnvelope
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&envelope); err != nil || len(envelope.Choices) != 1 {
		return nil, errSingleRequestPlanStage
	}
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		return nil, errSingleRequestPlanStage
	}
	choice := envelope.Choices[0]
	message := choice.Message
	if choice.Index != 0 || message.Role != "assistant" {
		return nil, errSingleRequestPlanStage
	}
	if choice.FinishReason == "length" {
		return nil, errors.Join(errSingleRequestPlanStage, errProviderStageOutputLimit)
	}
	if choice.FinishReason == "context_length" || choice.FinishReason == "context_length_exceeded" {
		return nil, errors.Join(errSingleRequestPlanStage, errProviderStageContextLimit)
	}
	if choice.FinishReason == "tool_calls" && message.Content == nil && len(message.ToolCalls) == 1 {
		call := message.ToolCalls[0]
		if call.ID == "" || call.Type != "function" || call.Function.Name == "" || call.Function.Arguments == "" {
			return nil, errSingleRequestPlanStage
		}
		return &singleRequestPlanProviderResponse{call: &call}, nil
	}
	if choice.FinishReason == "stop" && message.Content != nil && len(message.ToolCalls) == 0 {
		plan, err := decodeSingleRequestPlanFields(*message.Content, maximum)
		if err != nil {
			return nil, errSingleRequestPlanStage
		}
		return &singleRequestPlanProviderResponse{plan: plan}, nil
	}
	return nil, errSingleRequestPlanStage
}

func decodeSingleRequestPlanFields(raw string, maximum int) (*singleRequestPlanFields, error) {
	if len(raw) == 0 || len(raw) > maximum || validateSingleRequestJSON([]byte(raw)) != nil {
		return nil, errSingleRequestPlanStage
	}
	decoder := json.NewDecoder(strings.NewReader(raw))
	decoder.DisallowUnknownFields()
	var fields singleRequestPlanFields
	if err := decoder.Decode(&fields); err != nil {
		return nil, errSingleRequestPlanStage
	}
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		return nil, errSingleRequestPlanStage
	}
	return &fields, nil
}
