package workspace

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

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

	targetDir := t.TempDir()

	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID:               "sol_test",
		SourceKind:           "agent",
		SourceRef:            "agt_source",
		ActiveDefinitionPath: "memory/system/identity.md",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: "memory/system/identity.md", Content: "# Identity Seed\nReusable reviewer\n"},
			{Path: ".agents/skills/reviewer/SKILL.md", Content: "name: reviewer\n"},
			{Path: ".agents/knowledge/project-conventions.md", Content: "# Project Conventions\n"},
		},
	})
	require.NoError(t, err)

	identityBytes, err := os.ReadFile(filepath.Join(targetDir, "memory", "system", "identity.md"))
	require.NoError(t, err)
	require.Contains(t, string(identityBytes), "Identity Seed")

	knowledgeBytes, err := os.ReadFile(filepath.Join(targetDir, ".agents/knowledge/project-conventions.md"))
	require.NoError(t, err)
	require.Contains(t, string(knowledgeBytes), "Project Conventions")

	markerBytes, err := os.ReadFile(filepath.Join(targetDir, materializationMarkerRelPath))
	require.NoError(t, err)
	var marker map[string]any
	require.NoError(t, json.Unmarshal(markerBytes, &marker))
	require.Equal(t, "sol_test", marker["roleId"])
	require.Equal(t, "agent", marker["sourceKind"])
	require.Equal(t, "agt_source", marker["sourceRef"])
	require.Equal(t, "memory/system/identity.md", marker["activeDefinitionPath"])
	_, hasSourceAgentID := marker["sourceAgentId"]
	require.False(t, hasSourceAgentID)
}

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

	targetDir := t.TempDir()

	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID: "sol_seed",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: "memory/system/identity.md", Content: "# Identity Seed\n"},
		},
	})
	require.NoError(t, err)

	_, err = os.Stat(filepath.Join(targetDir, "MEMORY.md"))
	require.Error(t, err)
	require.True(t, os.IsNotExist(err))

	_, err = os.Stat(filepath.Join(targetDir, "notes", "checkpoints.md"))
	require.Error(t, err)
	require.True(t, os.IsNotExist(err))
}

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

	targetDir := t.TempDir()
	require.NoError(t, os.MkdirAll(targetDir, 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(targetDir, "custom.md"), []byte("keep me"), 0o644))

	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID: "sol_existing",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: "memory/system/identity.md", Content: "seed"},
		},
	})
	require.NoError(t, err)

	customBytes, err := os.ReadFile(filepath.Join(targetDir, "custom.md"))
	require.NoError(t, err)
	require.Equal(t, "keep me", string(customBytes))

	_, err = os.Stat(filepath.Join(targetDir, "memory", "system", "identity.md"))
	require.Error(t, err)
	require.True(t, os.IsNotExist(err))
}

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

	targetDir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(targetDir, "MEMORY.md"), []byte("# Existing Memory\n"), 0o644))

	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID: "sol_memory_seeded",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: "memory/system/identity.md", Content: "# Identity Seed\n"},
		},
	})
	require.NoError(t, err)

	memoryBytes, err := os.ReadFile(filepath.Join(targetDir, "MEMORY.md"))
	require.NoError(t, err)
	require.Equal(t, "# Existing Memory\n", string(memoryBytes))

	identityBytes, err := os.ReadFile(filepath.Join(targetDir, "memory", "system", "identity.md"))
	require.NoError(t, err)
	require.Equal(t, "# Identity Seed\n", string(identityBytes))
}

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

	targetDir := t.TempDir()
	require.NoError(t, os.MkdirAll(filepath.Join(targetDir, "memory"), 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(targetDir, "memory", "MEMORY.md"), []byte("# Memory\n"), 0o644))

	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID:               "sol_remote",
		ActiveDefinitionPath: "memory/system/identity.md",
		SeedFiles: []types.WorkspaceSeedFile{{
			Path:    "memory/system/identity.md",
			Content: "# Remote Identity\n",
		}},
	})
	require.NoError(t, err)

	identityBytes, err := os.ReadFile(filepath.Join(targetDir, "memory", "system", "identity.md"))
	require.NoError(t, err)
	require.Equal(t, "# Remote Identity\n", string(identityBytes))

	_, err = os.Stat(filepath.Join(targetDir, "IDENTITY.md"))
	require.True(t, os.IsNotExist(err))

	markerBytes, err := os.ReadFile(filepath.Join(targetDir, materializationMarkerRelPath))
	require.NoError(t, err)
	var marker map[string]any
	require.NoError(t, json.Unmarshal(markerBytes, &marker))
	require.Equal(t, "memory/system/identity.md", marker["activeDefinitionPath"])
}

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

	targetDir := t.TempDir()
	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID:               "sol_portable",
		ActiveDefinitionPath: "memory/system/identity.md",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: "memory/system/identity.md", Content: "# Identity\n"},
			{Path: ".agents/knowledge/runbook.md", Content: "# Runbook\n"},
			{Path: "MEMORY.md", Content: "skip"},
			{Path: "state/checkpoint.json", Content: "skip"},
			{Path: "kb/private.md", Content: "skip"},
			{Path: "archive/log.txt", Content: "skip"},
		},
	})
	require.NoError(t, err)

	_, err = os.Stat(filepath.Join(targetDir, ".agents/knowledge/runbook.md"))
	require.NoError(t, err)
	_, err = os.Stat(filepath.Join(targetDir, "MEMORY.md"))
	require.True(t, os.IsNotExist(err))
	_, err = os.Stat(filepath.Join(targetDir, "state", "checkpoint.json"))
	require.True(t, os.IsNotExist(err))
	_, err = os.Stat(filepath.Join(targetDir, "kb", "private.md"))
	require.True(t, os.IsNotExist(err))
	_, err = os.Stat(filepath.Join(targetDir, "archive", "log.txt"))
	require.True(t, os.IsNotExist(err))
}

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

	targetDir := t.TempDir()
	err := MaterializeAgentWorkspace(targetDir, types.WorkspaceMaterialization{
		RoleID: "sol_large_seed",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: ".agents/knowledge/first.md", Content: "first"},
			{Path: ".agents/knowledge/large.md", Content: strings.Repeat("x", maxMaterializationSeedFileBytes+1)},
		},
	})
	require.ErrorContains(t, err, "exceeds")

	_, statErr := os.Stat(filepath.Join(targetDir, "first.txt"))
	require.True(t, os.IsNotExist(statErr), "seed file validation must happen before any write")
}

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

	tooMany := make([]types.WorkspaceSeedFile, maxMaterializationSeedFiles+1)
	for i := range tooMany {
		tooMany[i] = types.WorkspaceSeedFile{Path: filepath.Join(".agents/knowledge", string(rune('a'+i))+".md"), Content: "x"}
	}
	err := MaterializeAgentWorkspace(t.TempDir(), types.WorkspaceMaterialization{RoleID: "sol_many", SeedFiles: tooMany})
	require.ErrorContains(t, err, "seed files has")

	err = MaterializeAgentWorkspace(t.TempDir(), types.WorkspaceMaterialization{
		RoleID: "sol_long_path",
		SeedFiles: []types.WorkspaceSeedFile{{
			Path:    ".agents/knowledge/" + strings.Repeat("x", maxMaterializationSeedFilePathRunes+1) + ".md",
			Content: "x",
		}},
	})
	require.ErrorContains(t, err, "path[0] exceeds")

	totalTooLarge := make([]types.WorkspaceSeedFile, 5)
	for i := range totalTooLarge {
		totalTooLarge[i] = types.WorkspaceSeedFile{
			Path:    filepath.Join(".agents/knowledge", string(rune('a'+i))+".md"),
			Content: strings.Repeat("x", maxMaterializationSeedFileBytes),
		}
	}
	err = MaterializeAgentWorkspace(t.TempDir(), types.WorkspaceMaterialization{RoleID: "sol_total", SeedFiles: totalTooLarge})
	require.ErrorContains(t, err, "total bytes")

	err = MaterializeAgentWorkspace(t.TempDir(), types.WorkspaceMaterialization{
		RoleID: "sol_traversal",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: ".agents/knowledge/first.md", Content: "first"},
			{Path: ".agents/knowledge/../../../escape.md", Content: "x"},
		},
	})
	require.ErrorContains(t, err, "path traversal")

	err = MaterializeAgentWorkspace(t.TempDir(), types.WorkspaceMaterialization{
		RoleID: "sol_skill_traversal",
		SeedFiles: []types.WorkspaceSeedFile{
			{Path: ".agents/skills/../SKILL.md", Content: "x"},
		},
	})
	require.ErrorContains(t, err, "path traversal")
}