package openai

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"strings"

	edgeservice "iop/apps/edge/internal/service"
)

const singleRequestPlanPrompt = "Produce exactly one JSON object with non-empty string fields plan and verification. Keep both concise."

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

type singleRequestPlanStage struct{ provider *singleRequestProviderStage }

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

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

func singleRequestPlanResponseFormat() *singleRequestProviderResponseFormat {
	return &singleRequestProviderResponseFormat{
		Type: "json_schema",
		JSONSchema: singleRequestProviderResponseJSONSchema{
			Name:   "single_request_plan",
			Strict: true,
			Schema: singleRequestProviderOutputSchema{
				Type: "object",
				Properties: map[string]singleRequestProviderOutputProperty{
					"plan": {
						Type:        "string",
						Description: "A concise execution plan for the task.",
					},
					"verification": {
						Type:        "string",
						Description: "A concise verification procedure for the plan.",
					},
				},
				Required:             []string{"plan", "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 || ctrl == nil || req.RequestID == "" || req.Task == "" || req.Sequence == 0 {
		return nil, quality.validation(errSingleRequestPlanStage)
	}
	if err := ctrl.SubmitEnvelope(edgeservice.SingleRequestEnvelope{RequestID: req.RequestID, Sequence: req.Sequence, Stage: edgeservice.SingleRequestStatePlanning}); err != nil {
		return nil, quality.serviceFailure(ctx, err, errSingleRequestPlanStage)
	}
	response, err := s.provider.submit(ctx, singleRequestProviderStageRequest{
		StageBinding: req.StageBinding, Limits: req.Limits, NodeRef: req.NodeRef, SessionID: req.SessionID, UsageAttribution: req.UsageAttribution, Quality: quality,
		Messages:       []chatMessage{{Role: "system", Content: singleRequestPlanPrompt}, {Role: "user", Content: req.Task}},
		ResponseFormat: singleRequestPlanResponseFormat(),
	})
	if err != nil {
		return nil, quality.reclassify(err, errSingleRequestPlanStage)
	}
	content, err := renderSingleRequestPlan(response.Output, 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
}

type singleRequestPlanResult struct {
	Plan         string `json:"plan"`
	Verification string `json:"verification"`
}

type singleRequestPlanResultAlias singleRequestPlanResult

func (r *singleRequestPlanResult) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "plan", "verification"); err != nil {
		return err
	}
	var a singleRequestPlanResultAlias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	*r = singleRequestPlanResult(a)
	return nil
}

func renderSingleRequestPlan(raw string, maximum int) ([]byte, error) {
	if err := validateSingleRequestJSON([]byte(raw)); err != nil {
		return nil, errSingleRequestPlanStage
	}
	decoder := json.NewDecoder(strings.NewReader(raw))
	decoder.DisallowUnknownFields()
	var result singleRequestPlanResult
	if err := decoder.Decode(&result); err != nil {
		return nil, errSingleRequestPlanStage
	}
	if strings.TrimSpace(result.Plan) == "" || strings.TrimSpace(result.Verification) == "" {
		return nil, errSingleRequestPlanStage
	}
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		return nil, errSingleRequestPlanStage
	}
	content := []byte("# Plan\n\n" + strings.TrimSpace(result.Plan) + "\n\n## Verification\n\n" + strings.TrimSpace(result.Verification) + "\n")
	if maximum < 1 || len(content) > maximum {
		return nil, errSingleRequestPlanStage
	}
	return content, nil
}
