package plan

import (
	"fmt"
	"math"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestCompilerRejectsOversizedSourceBeforeParsing(t *testing.T) {
	t.Parallel()

	var diagnostics []Diagnostic
	require.NotPanics(t, func() {
		_, diagnostics, _ = Compile([]byte(strings.Repeat("x", MaxSourceBytes+1)))
	})
	require.True(t, hasDiagnosticCode(diagnostics, "resource.source_too_large"), diagnostics)
}

func TestCompilerRejectsDeepYAMLBeforeRecursiveSourceIndex(t *testing.T) {
	t.Parallel()

	source := "apiVersion: loop.plan/v1\nkind: PlanVersion\ngraph: " +
		strings.Repeat("[", MaxSourceYAMLNesting+2) + "0" +
		strings.Repeat("]", MaxSourceYAMLNesting+2)
	var diagnostics []Diagnostic
	require.NotPanics(t, func() {
		_, diagnostics, _ = Compile([]byte(source))
	})
	require.True(t, hasDiagnosticCode(diagnostics, "resource.yaml_nesting"), diagnostics)
}

func TestValidateRejectsGraphAndSchemaBudgetsBeforeAnalysis(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name     string
		document *Document
		code     string
	}{
		{
			name: "nodes",
			document: &Document{APIVersion: APIVersion, Kind: Kind, Graph: Graph{
				Nodes: make([]Node, MaxGraphNodes+1),
			}},
			code: "resource.too_many_nodes",
		},
		{
			name: "edges",
			document: &Document{APIVersion: APIVersion, Kind: Kind, Graph: Graph{
				Edges: make([]Edge, MaxGraphEdges+1),
			}},
			code: "resource.too_many_edges",
		},
		{
			name: "schemas",
			document: &Document{APIVersion: APIVersion, Kind: Kind,
				Input: schemaMapForBudget(MaxSchemaFields + 1),
			},
			code: "resource.too_many_schema_fields",
		},
	}
	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			var diagnostics []Diagnostic
			require.NotPanics(t, func() { diagnostics = Validate(test.document) })
			require.True(t, hasDiagnosticCode(diagnostics, test.code), diagnostics)
		})
	}
}

func TestValidateRejectsLiteralAndExpressionNestingBudgets(t *testing.T) {
	t.Parallel()

	value := any("leaf")
	for index := 0; index <= MaxLiteralNesting; index++ {
		value = []any{value}
	}
	document := &Document{
		APIVersion: APIVersion,
		Kind:       Kind,
		Input: map[string]Schema{
			"nested": {Type: TypeArray, Items: &Schema{Type: TypeString}, Default: value},
		},
	}
	diagnostics := Validate(document)
	require.True(t, hasDiagnosticCode(diagnostics, "resource.literal_budget"), diagnostics)

	require.Contains(t, expressionBudgetError(strings.Repeat("!", MaxExpressionNesting+1)+"true"), "nesting")
	require.Contains(t, expressionBudgetError(strings.Repeat("x", MaxExpressionBytes+1)), "bytes")
}

func TestValidateRejectsLoopAndRetryUpperBounds(t *testing.T) {
	t.Parallel()

	document := &Document{
		APIVersion: APIVersion,
		Kind:       Kind,
		Graph: Graph{Nodes: []Node{
			{
				ID: "loop", Type: NodeLoop,
				Loop: &AgentTeamConfig{
					Mode:       LoopCounted,
					Iterations: MaxLoopIterations + 1,
					State:      map[string]Schema{"value": {Type: TypeString}},
				},
			},
			{
				ID: "wait", Type: NodeWait,
				Wait: &WaitConfig{
					Until: ValueExpr{Ref: "input.until"}, Timeout: "1s",
					Retry: &RetryPolicy{MaxAttempts: MaxRetryAttempts + 1, Backoff: "1s"},
				},
				Input: map[string]Schema{"until": {Type: TypeDuration}},
			},
		}},
	}
	diagnostics := Validate(document)
	require.True(t, hasDiagnosticCode(diagnostics, "loop.iterations_too_large"), diagnostics)
	require.True(t, hasDiagnosticCode(diagnostics, "retry.too_many_attempts"), diagnostics)
}

func TestArtifactBudgetsRunBeforeDigestAndRecursiveValidation(t *testing.T) {
	t.Parallel()

	_, err := DecodeArtifact([]byte(strings.Repeat(" ", MaxArtifactBytes+1)))
	require.ErrorContains(t, err, "must not exceed")
	deepArtifact := strings.Repeat("[", MaxArtifactJSONNesting+1) + "0" + strings.Repeat("]", MaxArtifactJSONNesting+1)
	_, err = DecodeArtifact([]byte(deepArtifact))
	require.ErrorContains(t, err, "JSON nesting")

	schema := ArtifactSchema{Type: TypeString, Values: []string{}}
	for index := 0; index <= maxSchemaArrayDepth; index++ {
		child := schema
		schema = ArtifactSchema{Type: TypeArray, Values: []string{}, Items: &child}
	}
	artifact := &Artifact{
		ArtifactVersion: ArtifactVersion,
		CompilerVersion: CompilerVersion,
		APIVersion:      APIVersion,
		Digest:          "sha256:not-yet-checked",
		Contracts: ArtifactContracts{
			Input: ArtifactContract{{Name: "nested", Schema: schema}},
		},
		Nodes: []ArtifactNode{}, Edges: []ArtifactEdge{}, Scopes: []ArtifactScope{},
		Entry: ArtifactEntry{Bindings: []ArtifactBinding{}},
	}
	err = ValidateEnvelope(artifact)
	require.ErrorContains(t, err, "schema nesting")

	expression := ArtifactExpression{Op: ArtifactOpConst, ResultType: TypeBoolean, Value: ArtifactLiteral("true")}
	for index := 0; index <= MaxArtifactExprNesting; index++ {
		expression = ArtifactExpression{Op: ArtifactOpNot, ResultType: TypeBoolean, Args: []ArtifactExpression{expression}}
	}
	artifact.Contracts.Input = ArtifactContract{}
	artifact.Entry.Bindings = []ArtifactBinding{{Target: "value", Value: expression}}
	err = ValidateEnvelope(artifact)
	require.ErrorContains(t, err, "expression nesting")
}

func TestCompiledArtifactWireBudgetProducesStableDiagnostic(t *testing.T) {
	t.Parallel()

	artifact := &Artifact{
		ArtifactVersion: ArtifactVersion,
		CompilerVersion: CompilerVersion,
		APIVersion:      APIVersion,
		Nodes: []ArtifactNode{{
			ID: "message", Opcode: OpcodeSendMessageSent,
			Config: ArtifactNodeConfig{SendMessage: &ArtifactSendMessageConfig{
				Text: ArtifactTemplate{Segments: []ArtifactTemplateSegment{{
					Kind: ArtifactTemplateSegmentText,
					Text: strings.Repeat("x", MaxArtifactBytes),
				}}},
			}},
		}},
		Edges: []ArtifactEdge{}, Scopes: []ArtifactScope{},
		Entry: ArtifactEntry{Bindings: []ArtifactBinding{}},
	}
	diagnostics, err := artifactWireBudgetDiagnostics(artifact)
	require.NoError(t, err)
	require.True(t, hasDiagnosticCode(diagnostics, "resource.artifact_too_large"), diagnostics)
}

func TestCompilerRejectsProjectedArtifactAmplificationBeforeLowering(t *testing.T) {
	// Prove the generated graph itself is a valid authoring shape before using
	// the maximum-sized variant as an amplification input.
	control := projectedArtifactAmplificationSource(4, 4)
	artifact, diagnostics, err := Compile(control)
	require.NoError(t, err)
	require.False(t, HasErrors(diagnostics), diagnostics)
	require.NotNil(t, artifact)

	// One declared contract is copied to every fork output port. The source is
	// well below 1 MiB, but naïve lowering would allocate roughly 2,048 × 2,045
	// ArtifactField values before the final 8 MiB wire check could run.
	attack := projectedArtifactAmplificationSource(MaxSchemaFields, MaxGraphEdges-3)
	require.Less(t, len(attack), MaxSourceBytes)

	artifact = nil
	diagnostics = nil
	require.NotPanics(t, func() {
		artifact, diagnostics, err = Compile(attack)
	})
	require.NoError(t, err)
	require.Nil(t, artifact)
	require.Len(t, diagnostics, 1, diagnostics)
	require.Equal(t, "resource.artifact_budget", diagnostics[0].Code)
	require.Equal(t, "$", diagnostics[0].Path)
}

func TestArtifactProjectionArithmeticSaturates(t *testing.T) {
	t.Parallel()

	require.Equal(t, uint64(MaxArtifactMaterializationUnits+1), saturatedAdd(math.MaxUint64, 1, MaxArtifactMaterializationUnits+1))
	require.Equal(t, uint64(MaxArtifactProjectedBytes+1), saturatedMultiply(math.MaxUint64, math.MaxUint64, MaxArtifactProjectedBytes+1))

	projection := artifactMaterializationProjection{}
	projection.addContract(artifactSchemaProjection{
		fields: math.MaxUint64,
		units:  math.MaxUint64,
		bytes:  math.MaxUint64,
	}, math.MaxUint64)
	require.Equal(t, uint64(MaxArtifactSchemaFields+1), projection.schemaFields)
	require.Equal(t, uint64(MaxArtifactMaterializationUnits+1), projection.units)
	require.Equal(t, uint64(MaxArtifactProjectedBytes+1), projection.projectedBytes)
}

func TestArtifactProjectionAccountsForJSONEscapeAmplification(t *testing.T) {
	t.Parallel()

	// encoding/json expands every '&' to "\\u0026". With a raw-byte
	// estimate this contract appears to fit below the 64 MiB preflight budget,
	// while copying it to every fork port would produce hundreds of MiB of
	// deterministic JSON. Exercise the projection directly so the test can never
	// accidentally lower that artifact.
	defaultValue := strings.Repeat("&", 30_000)
	branches := make([]string, MaxGraphEdges-3)
	for index := range branches {
		branches[index] = fmt.Sprintf("branch_%04d", index)
	}
	document := &Document{
		APIVersion: APIVersion,
		Kind:       Kind,
		Graph: Graph{Nodes: []Node{{
			ID:    "fan",
			Type:  NodeFork,
			Input: map[string]Schema{"value": {Type: TypeString, Default: defaultValue}},
			Fork:  &ForkConfig{Branches: branches},
		}}},
	}

	contractCopies := uint64(len(branches) + 1) // one input port plus every output port
	require.Less(t, uint64(len(defaultValue))*contractCopies, uint64(MaxArtifactProjectedBytes))

	projection := projectArtifactMaterialization(document, nil)
	require.LessOrEqual(t, projection.schemaFields, uint64(MaxArtifactSchemaFields))
	require.LessOrEqual(t, projection.expressions, uint64(MaxArtifactExpressions))
	require.LessOrEqual(t, projection.units, uint64(MaxArtifactMaterializationUnits))
	require.Greater(t, projection.projectedBytes, uint64(MaxArtifactProjectedBytes))

	diagnostics := artifactMaterializationBudgetDiagnostics(document, nil)
	require.Len(t, diagnostics, 1)
	require.Equal(t, "resource.artifact_budget", diagnostics[0].Code)
}

func projectedArtifactAmplificationSource(fieldCount, branchCount int) []byte {
	var source strings.Builder
	source.WriteString("apiVersion: loop.plan/v1\nkind: PlanVersion\ngraph:\n  entry: {nodeId: fan, port: in}\n  nodes:\n")
	source.WriteString("    - id: fan\n      type: fork\n      input:\n")
	for index := 0; index < fieldCount; index++ {
		fmt.Fprintf(&source, "        field_%04d: {type: string}\n", index)
	}
	source.WriteString("      fork:\n        branches:\n")
	for index := 0; index < branchCount; index++ {
		fmt.Fprintf(&source, "          - branch_%04d\n", index)
	}
	source.WriteString("    - id: join\n      type: join\n      join: {forkId: fan, mode: all}\n")
	source.WriteString("    - id: success\n      type: end\n      end: {outcome: succeeded}\n")
	source.WriteString("    - id: failure\n      type: end\n      end: {outcome: failed, code: test.failure}\n")
	source.WriteString("  edges:\n")
	for index := 0; index < branchCount; index++ {
		fmt.Fprintf(&source, "    - id: edge_%04d\n      source: {nodeId: fan, port: branch_%04d}\n      target: {nodeId: join, port: branch_%04d}\n", index, index, index)
	}
	source.WriteString("    - id: join_completed\n      source: {nodeId: join, port: completed}\n      target: {nodeId: success, port: in}\n")
	source.WriteString("    - id: join_failed\n      source: {nodeId: join, port: failed}\n      target: {nodeId: failure, port: in}\n")
	return []byte(source.String())
}

func schemaMapForBudget(count int) map[string]Schema {
	result := make(map[string]Schema, count)
	for index := 0; index < count; index++ {
		result[fmt.Sprintf("field_%04d", index)] = Schema{Type: TypeString}
	}
	return result
}

func hasDiagnosticCode(diagnostics []Diagnostic, code string) bool {
	for _, diagnostic := range diagnostics {
		if diagnostic.Code == code {
			return true
		}
	}
	return false
}