tests/test_project_context.py

from __future__ import annotations

from pathlib import Path
import shutil

import pytest

import kneiff.project_resources as project_resources_module
from kneiff.config import KNEIFF_RC
from kneiff.datasets.manifest.schema import manifest_schema_for_vocabulary
from kneiff.datasets.export.workflow import (
    build_dataset_sync_plan_from_config_path,
    validate_dataset_sync_plan_paths,
)
from kneiff.infer.comfy.workflow_presets import (
    project_workflow_files,
    resolve_showcase_workflow,
)
from kneiff.project import ProjectContext, ProjectContextError, load_project_context
from kneiff.project_layout import (
    DEFAULT_ANIMA_CONFIG_FILENAME,
    DEFAULT_FLUX2_KLEIN_9B_CONFIG_FILENAME,
    PROJECT_PROMPTS_FILENAME,
    PROJECT_VOCABULARY_FILENAME,
)
from kneiff.project_resources import (
    load_project_prompt_catalog,
    load_project_vocabulary,
    validate_project_resources,
    validate_project_scaffold,
)
from kneiff.project_scaffold import (
    ensure_project_resources,
    initialize_project_scaffold,
)
from kneiff.training.lora.config_templates import (
    read_default_flux2_klein_9b_starter_config,
)

AURORA_PROJECT_ROOT = Path(__file__).parent / "fixtures" / "projects" / "aurora"


def _write_project_config(path: Path) -> None:
    """Write the smallest config-first dataset mapping used by safety tests."""
    path.write_text(
        "mappings:\n  fullbody:\n    - 0-FULLBODY\n",
        encoding="utf-8",
    )


def test_project_context_derives_shallow_paths_without_loading_resources(
    tmp_path: Path,
) -> None:
    project = load_project_context(tmp_path, name="demo")

    assert project.name == "demo"
    assert project.root == tmp_path.resolve()
    assert project.vocabulary_path == tmp_path / PROJECT_VOCABULARY_FILENAME
    assert project.prompts_path == tmp_path / PROJECT_PROMPTS_FILENAME
    assert project.default_anima_config_path == (
        tmp_path / "configs" / DEFAULT_ANIMA_CONFIG_FILENAME
    )
    assert project.default_flux2_klein_9b_config_path == (
        tmp_path / "configs" / DEFAULT_FLUX2_KLEIN_9B_CONFIG_FILENAME
    )
    assert project.workflows_path == tmp_path / "workflows"


def test_project_context_reports_missing_vocabulary_at_consuming_boundary(
    tmp_path: Path,
) -> None:
    project = load_project_context(tmp_path)

    with pytest.raises(ProjectContextError, match=PROJECT_VOCABULARY_FILENAME):
        load_project_vocabulary(project)


def test_project_context_rejects_registered_root_symlink_before_resolving(
    tmp_path: Path,
) -> None:
    registered_target = tmp_path / "real-project"
    registered_target.mkdir()
    registered_root = tmp_path / "registered-project"
    registered_root.symlink_to(registered_target, target_is_directory=True)

    with pytest.raises(ProjectContextError, match="Project root must not be a symlink"):
        load_project_context(registered_root)


def test_project_context_rejects_symlinked_storage_environment(
    tmp_path: Path,
) -> None:
    project_root = tmp_path / "project"
    project_root.mkdir()
    external_env = tmp_path / "external.env"
    external_env.write_text("KNF_WORKERS=9\n", encoding="utf-8")
    (project_root / KNEIFF_RC.spec.storage_env_filename).symlink_to(external_env)

    with pytest.raises(
        ProjectContextError,
        match="AppRC storage environment must not be a symlink",
    ):
        load_project_context(project_root)


def test_project_context_rejects_storage_environment_directory(
    tmp_path: Path,
) -> None:
    project_root = tmp_path / "project"
    project_root.mkdir()
    (project_root / KNEIFF_RC.spec.storage_env_filename).mkdir()

    with pytest.raises(
        ProjectContextError,
        match="AppRC storage environment must be a file",
    ):
        load_project_context(project_root)


def test_project_context_rejects_live_ancestor_symlink_redirection(
    tmp_path: Path,
) -> None:
    owner = tmp_path / "owner"
    project_root = owner / "project"
    project_root.mkdir(parents=True)
    project = load_project_context(project_root)

    original_owner = tmp_path / "original-owner"
    owner.rename(original_owner)
    redirected_owner = tmp_path / "redirected-owner"
    (redirected_owner / "project").mkdir(parents=True)
    owner.symlink_to(redirected_owner, target_is_directory=True)

    with pytest.raises(
        ProjectContextError,
        match="must remain at its canonical location after selection",
    ):
        project.optional_file(project.prompts_path, label="Project prompts")


def test_ensure_project_resources_creates_root_level_files(tmp_path: Path) -> None:
    project = ensure_project_resources(
        tmp_path / "demo",
        activation_token="Demo_Character",
        species_token="demo_species",
    )

    vocabulary = load_project_vocabulary(project)
    prompt_catalog = load_project_prompt_catalog(project)
    assert vocabulary.character_token == "Demo_Character"
    assert vocabulary.species_token == "demo_species"
    assert prompt_catalog.core_prompt_count > 0
    assert prompt_catalog.overlay_prompt_count == 0
    assert prompt_catalog.resolved_prompt_count == prompt_catalog.core_prompt_count
    assert prompt_catalog.prompts[0].id == "character_reference"
    assert project.prompts_path.read_text(encoding="utf-8") == (
        "version: 1\nprompts: []\n"
    )
    assert not (project.root / "kneiff.project.yaml").exists()
    assert not (project.root / "resources").exists()
    assert not project.workflows_path.exists()


def test_ensure_project_resources_requires_tokens_only_for_missing_vocabulary(
    tmp_path: Path,
) -> None:
    root = tmp_path / "demo"
    root.mkdir()
    source = AURORA_PROJECT_ROOT / PROJECT_VOCABULARY_FILENAME
    (root / PROJECT_VOCABULARY_FILENAME).write_bytes(source.read_bytes())

    project = ensure_project_resources(root)

    assert load_project_vocabulary(project).character_token == "Aurora_Fox"
    assert project.prompts_path.is_file()


def test_project_resources_allow_missing_prompt_overlay(tmp_path: Path) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    project.prompts_path.unlink()

    validation = validate_project_scaffold(project)

    assert validation.project_prompt_count == 0
    assert validation.resolved_prompt_count == validation.core_prompt_count


def test_project_resource_validation_loads_vocabulary_once(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    original_loader = project_resources_module.load_project_vocabulary
    load_count = 0

    def tracked_loader(project_context: ProjectContext):
        nonlocal load_count
        load_count += 1
        return original_loader(project_context)

    monkeypatch.setattr(
        project_resources_module,
        "load_project_vocabulary",
        tracked_loader,
    )

    validation = project_resources_module.validate_project_resources(project)

    assert validation.vocabulary.character_token == "Demo_Character"
    assert load_count == 1


def test_project_init_keeps_existing_flux2_config_on_repeated_initialization(
    tmp_path: Path,
) -> None:
    project_root = tmp_path / "project"
    first = initialize_project_scaffold(
        project_root,
        activation_token="Demo_Character",
        species_token="demo_species",
    )
    config_path = first.project.default_flux2_klein_9b_config_path
    config_path.write_text(
        "# Existing user customization.\n"
        + read_default_flux2_klein_9b_starter_config(),
        encoding="utf-8",
    )
    expected_bytes = config_path.read_bytes()

    second = initialize_project_scaffold(project_root)

    assert config_path.read_bytes() == expected_bytes
    assert config_path not in second.created_files


def test_project_init_adopts_overlay_before_creating_vocabulary(tmp_path: Path) -> None:
    root = tmp_path / "project"
    root.mkdir()
    (root / PROJECT_PROMPTS_FILENAME).write_bytes(
        (AURORA_PROJECT_ROOT / PROJECT_PROMPTS_FILENAME).read_bytes()
    )

    project = ensure_project_resources(
        root,
        activation_token="Demo_Character",
        species_token="demo_species",
    )
    prompt_catalog = load_project_prompt_catalog(project)

    assert prompt_catalog.overlay_prompt_count == 1
    assert "aurora_portrait" in {prompt.id for prompt in prompt_catalog.prompts}


def test_non_rook_project_resources_drive_selector_boundaries() -> None:
    project = load_project_context(AURORA_PROJECT_ROOT, name="aurora")
    vocabulary = load_project_vocabulary(project)
    schema = manifest_schema_for_vocabulary(vocabulary)
    prompt_catalog = load_project_prompt_catalog(project)
    workflow = resolve_showcase_workflow("demo", project=project)

    assert vocabulary.character_token == "Aurora_Fox"
    assert vocabulary.species_token == "foxkin"
    assert schema.manual_columns == (
        "subject",
        "appearance",
        "composition",
        "pose_and_behavior",
        "face",
        "anatomy",
        "sexual_content",
        "scene",
        "presentation",
        "fallback",
    )
    aurora_prompt = next(
        prompt for prompt in prompt_catalog.prompts if prompt.id == "aurora_portrait"
    )
    assert aurora_prompt.positive_prompt("nlg").startswith("Aurora_Fox.")
    assert list(project_workflow_files(project)) == ["demo"]
    assert (
        workflow.source.path
        == (AURORA_PROJECT_ROOT / "workflows" / "showcase_demo.workflow.json").resolve()
    )


def test_project_workflow_discovery_rejects_symlinked_directory(
    tmp_path: Path,
) -> None:
    project_root = tmp_path / "project"
    project_root.mkdir()
    external_workflows = tmp_path / "external-workflows"
    external_workflows.mkdir()
    (project_root / "workflows").symlink_to(
        external_workflows,
        target_is_directory=True,
    )
    project = load_project_context(project_root)

    with pytest.raises(ValueError, match="directory must not be a symlink"):
        project_workflow_files(project)


def test_project_workflow_discovery_rejects_symlinked_file(tmp_path: Path) -> None:
    project_root = tmp_path / "project"
    workflows_path = project_root / "workflows"
    workflows_path.mkdir(parents=True)
    external_workflow = tmp_path / "external.workflow.json"
    external_workflow.write_text("{}", encoding="utf-8")
    (workflows_path / "showcase_demo.workflow.json").symlink_to(external_workflow)
    project = load_project_context(project_root)

    with pytest.raises(ValueError, match="workflow path must not be a symlink"):
        project_workflow_files(project)


def test_project_workflow_discovery_rejects_normalized_duplicate_presets(
    tmp_path: Path,
) -> None:
    project_root = tmp_path / "project"
    workflows_path = project_root / "workflows"
    workflows_path.mkdir(parents=True)
    (workflows_path / "showcase_NOOB_Nova.workflow.json").write_text(
        "{}",
        encoding="utf-8",
    )
    (workflows_path / "showcase_noob-nova.workflow.json").write_text(
        "{}",
        encoding="utf-8",
    )
    project = load_project_context(project_root)

    with pytest.raises(ValueError, match="duplicate preset id 'noob-nova'"):
        project_workflow_files(project)


def test_project_resource_validation_loads_showcase_workflow_contract(
    tmp_path: Path,
) -> None:
    project = ensure_project_resources(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    )
    project.workflows_path.mkdir()
    (project.workflows_path / "showcase_incomplete.workflow.json").write_text(
        """
{
  "1": {
    "class_type": "CLIPTextEncode",
    "inputs": {"text": "__KNF_PROMPT__"}
  }
}
""",
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="missing required placeholders"):
        validate_project_resources(project)


@pytest.mark.parametrize(
    "dirname",
    ("configs", "SOURCE", "HF", "TRAINING", ".old_manifests"),
)
@pytest.mark.parametrize(
    ("mutation", "error"),
    (
        ("missing", "is missing"),
        ("wrong_type", "must be a directory"),
        ("symlink", "must not be a symlink"),
    ),
)
def test_project_scaffold_validation_rejects_required_directory_mutations(
    tmp_path: Path,
    dirname: str,
    mutation: str,
    error: str,
) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    path = project.root / dirname
    shutil.rmtree(path)
    if mutation == "wrong_type":
        path.write_text("wrong type\n", encoding="utf-8")
    elif mutation == "symlink":
        external = tmp_path / f"external-{dirname}"
        external.mkdir()
        path.symlink_to(external, target_is_directory=True)

    with pytest.raises(ProjectContextError, match=error):
        validate_project_scaffold(project)


@pytest.mark.parametrize(
    "mutated_path",
    (
        "configs",
        "SOURCE",
        "HF",
        "TRAINING",
        "MANIFEST.knf.xlsx",
        ".old_manifests",
        ".old_manifests/_thumbnail_cache",
    ),
)
def test_dataset_plan_rejects_fixed_path_symlink_added_after_project_init(
    tmp_path: Path,
    mutated_path: str,
) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    config_path = project.configs_path / "demo.knf.yaml"
    _write_project_config(config_path)
    plan = build_dataset_sync_plan_from_config_path(config_path)
    path = project.root / mutated_path
    external = tmp_path / f"external-{path.name}"

    if path.exists() and path.is_dir():
        path.rename(external)
        path.symlink_to(external, target_is_directory=True)
    elif path.suffix:
        external.write_text("external", encoding="utf-8")
        path.symlink_to(external)
    else:
        external.mkdir()
        path.symlink_to(external, target_is_directory=True)

    with pytest.raises(ProjectContextError, match="must not be a symlink"):
        validate_dataset_sync_plan_paths(plan)


@pytest.mark.parametrize(
    "mutated_path",
    (
        "configs",
        "SOURCE",
        "HF",
        "TRAINING",
        "MANIFEST.knf.xlsx",
        ".old_manifests",
        ".old_manifests/_thumbnail_cache",
    ),
)
def test_dataset_plan_rejects_fixed_path_wrong_type_after_project_init(
    tmp_path: Path,
    mutated_path: str,
) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    config_path = project.configs_path / "demo.knf.yaml"
    _write_project_config(config_path)
    plan = build_dataset_sync_plan_from_config_path(config_path)
    path = project.root / mutated_path

    if path.is_dir():
        shutil.rmtree(path)
        path.write_text("wrong type", encoding="utf-8")
    elif path.suffix:
        path.mkdir()
    else:
        path.write_text("wrong type", encoding="utf-8")

    with pytest.raises(ProjectContextError, match="must be"):
        validate_dataset_sync_plan_paths(plan)


def test_project_source_boundary_allows_nested_source_symlinks(tmp_path: Path) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    external_image = tmp_path / "external.png"
    external_image.write_bytes(b"image")
    nested_dir = project.source_path / "nested"
    nested_dir.mkdir()
    (nested_dir / "linked.png").symlink_to(external_image)

    assert (
        project.require_directory(
            project.source_path,
            label="Project SOURCE directory",
        )
        == project.source_path
    )


def test_project_prompt_catalog_rejects_post_init_symlink(tmp_path: Path) -> None:
    project = initialize_project_scaffold(
        tmp_path / "project",
        activation_token="Demo_Character",
        species_token="demo_species",
    ).project
    external_prompts = tmp_path / "external-prompts.yaml"
    project.prompts_path.replace(external_prompts)
    project.prompts_path.symlink_to(external_prompts)

    with pytest.raises(ValueError, match="must not be a symlink"):
        load_project_prompt_catalog(project)