"""
Tests for Memory Orchestrator — PRD v1.0

Coverage:
    §6   TaskClassifier
    §7   MemoryRanker
    §8   BudgetAllocator
    §9   MemorySelector / SelectionPlan
    §10  PromptBuilder
    §5   MemoryOrchestrator (full pipeline)
    §14  Adapter integration
"""

import os
import sys
import pytest

OPENMEMO_ROOT = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, OPENMEMO_ROOT)

from openmemo.memory_orchestrator.task_classifier import TaskClassifier, TASK_TYPES
from openmemo.memory_orchestrator.ranker import MemoryRanker, MEMORY_LAYERS
from openmemo.memory_orchestrator.budget_allocator import BudgetAllocator
from openmemo.memory_orchestrator.selector import MemorySelector, SelectionPlan
from openmemo.memory_orchestrator.prompt_builder import PromptBuilder
from openmemo.memory_orchestrator import MemoryOrchestrator


# ══════════════════════════════════════════════════════════════════════════════
# §6 — TaskClassifier
# ══════════════════════════════════════════════════════════════════════════════

class TestTaskClassifier:
    def setup_method(self):
        self.clf = TaskClassifier()

    def test_deployment_keyword(self):
        assert self.clf.classify("Deploy my FastAPI backend") == "deployment"

    def test_deployment_docker(self):
        assert self.clf.classify("Build and push docker image to k8s") == "deployment"

    def test_debugging_keyword(self):
        assert self.clf.classify("Debug the error in my code") == "debugging"

    def test_debugging_fix(self):
        assert self.clf.classify("Fix the bug in auth.py") == "debugging"

    def test_coding_keyword(self):
        assert self.clf.classify("Write a function to parse JSON") == "coding"

    def test_coding_implement(self):
        assert self.clf.classify("Implement the REST API endpoint") == "coding"

    def test_research_keyword(self):
        assert self.clf.classify("Research best practices for caching") == "research"

    def test_research_what_is(self):
        assert self.clf.classify("What is vector database?") == "research"

    def test_planning_keyword(self):
        assert self.clf.classify("Plan the sprint for next week") == "planning"

    def test_planning_design(self):
        assert self.clf.classify("Design the architecture for microservices") == "planning"

    def test_empty_task_returns_default(self):
        result = self.clf.classify("")
        assert result in TASK_TYPES

    def test_agent_role_fallback(self):
        result = self.clf.classify("", agent_role="debugger")
        assert result == "debugging"

    def test_agent_role_devops(self):
        result = self.clf.classify("", agent_role="devops engineer")
        assert result == "deployment"

    def test_all_types_are_valid(self):
        tasks = [
            "deploy service",
            "debug error",
            "write code",
            "research topic",
            "plan roadmap",
        ]
        for task in tasks:
            result = self.clf.classify(task)
            assert result in TASK_TYPES, f"Invalid type {result} for '{task}'"

    def test_classify_prioritizes_debugging_over_coding(self):
        # "fix the bug" should be debugging, not coding
        result = self.clf.classify("fix the bug in my code")
        assert result == "debugging"


# ══════════════════════════════════════════════════════════════════════════════
# §7 — MemoryRanker
# ══════════════════════════════════════════════════════════════════════════════

class TestMemoryRanker:
    def setup_method(self):
        self.ranker = MemoryRanker()

    def test_rank_returns_all_layers(self):
        result = self.ranker.rank("coding")
        assert set(result) == set(MEMORY_LAYERS)
        assert len(result) == len(MEMORY_LAYERS)

    def test_user_memory_high_priority_default(self):
        result = self.ranker.rank("coding")
        assert result.index("user_memory") < result.index("experience")
        assert result.index("user_memory") < result.index("team_memory")

    def test_deployment_mode_experience_rises(self):
        result = self.ranker.rank("deployment")
        # Execution mode: experience ↑
        exp_idx = result.index("experience")
        ctx_idx = result.index("context")
        assert exp_idx < ctx_idx  # experience ranked above context

    def test_debugging_mode_context_rises(self):
        result = self.ranker.rank("debugging")
        # Debug mode: context ↑
        ctx_idx = result.index("context")
        exp_idx = result.index("experience")
        assert ctx_idx < exp_idx  # context ranked above experience

    def test_exploration_mode_context_rises(self):
        result = self.ranker.rank("planning")
        ctx_idx = result.index("context")
        exp_idx = result.index("experience")
        assert ctx_idx < exp_idx

    def test_get_mode(self):
        assert self.ranker.get_mode("deployment") == "execution"
        assert self.ranker.get_mode("debugging") == "debug"
        assert self.ranker.get_mode("planning") == "exploration"
        assert self.ranker.get_mode("research") == "exploration"
        assert self.ranker.get_mode("coding") == "execution"

    def test_mode_override(self):
        # Force exploration mode even for deployment task
        result = self.ranker.rank("deployment", mode="exploration")
        ctx_idx = result.index("context")
        exp_idx = result.index("experience")
        assert ctx_idx < exp_idx  # exploration: context before experience

    def test_priority_scores_keys(self):
        scores = self.ranker.priority_scores("coding")
        assert set(scores.keys()) == set(MEMORY_LAYERS)


# ══════════════════════════════════════════════════════════════════════════════
# §8 — BudgetAllocator
# ══════════════════════════════════════════════════════════════════════════════

class TestBudgetAllocator:
    def setup_method(self):
        self.allocator = BudgetAllocator(max_tokens=2000)

    def test_allocate_sums_to_max(self):
        budget = self.allocator.allocate("coding")
        assert sum(budget.values()) == 2000

    def test_allocate_all_task_types(self):
        for task_type in TASK_TYPES:
            budget = self.allocator.allocate(task_type)
            assert sum(budget.values()) == 2000, f"Budget mismatch for {task_type}"

    def test_allocate_has_all_layers(self):
        budget = self.allocator.allocate("coding")
        assert "experience" in budget
        assert "context" in budget
        assert "user_memory" in budget
        assert "task_memory" in budget
        assert "team_memory" in budget

    def test_no_layer_below_50(self):
        for task_type in TASK_TYPES:
            budget = self.allocator.allocate(task_type)
            for layer, tokens in budget.items():
                assert tokens >= 50, f"{layer} below 50 for {task_type}"

    def test_deployment_experience_higher_than_default(self):
        deploy_budget = self.allocator.allocate("deployment")
        code_budget = self.allocator.allocate("coding")
        assert deploy_budget["experience"] >= code_budget["experience"]

    def test_debugging_context_higher_than_deployment(self):
        debug_budget = self.allocator.allocate("debugging")
        deploy_budget = self.allocator.allocate("deployment")
        assert debug_budget["context"] > deploy_budget["context"]

    def test_custom_max_tokens(self):
        allocator = BudgetAllocator(max_tokens=1000)
        budget = allocator.allocate("coding")
        assert sum(budget.values()) == 1000

    def test_remaining(self):
        budget = {"experience": 800, "context": 400}
        used = {"experience": 300, "context": 200}
        remaining = self.allocator.remaining(budget, used)
        assert remaining["experience"] == 500
        assert remaining["context"] == 200

    def test_utilization(self):
        budget = {"experience": 800, "context": 400}
        used = {"experience": 700, "context": 380}
        rate = self.allocator.utilization(budget, used)
        assert 0.0 < rate <= 1.0
        assert rate > 0.85  # should be high utilization


# ══════════════════════════════════════════════════════════════════════════════
# §9 — MemorySelector
# ══════════════════════════════════════════════════════════════════════════════

class TestMemorySelector:
    def setup_method(self):
        self.selector = MemorySelector()

    def test_select_returns_plan(self):
        plan = self.selector.select("coding", user_id="alice", session_id="s1")
        assert isinstance(plan, SelectionPlan)

    def test_experience_always_included(self):
        plan = self.selector.select("coding")
        assert plan.include_experience is True

    def test_user_memory_included_when_user_id(self):
        plan = self.selector.select("coding", user_id="alice")
        assert plan.include_user_memory is True

    def test_user_memory_excluded_when_no_user_id(self):
        plan = self.selector.select("coding", user_id="")
        assert plan.include_user_memory is False

    def test_context_included_when_session_id(self):
        plan = self.selector.select("coding", session_id="s1")
        assert plan.include_context is True

    def test_context_excluded_when_no_session(self):
        plan = self.selector.select("coding")
        assert plan.include_context is False

    def test_task_memory_included_when_agent_id(self):
        plan = self.selector.select("coding", agent_id="agent1")
        assert plan.include_task_memory is True

    def test_team_memory_included_for_relevant_tasks(self):
        plan = self.selector.select("deployment", team_id="team1")
        assert plan.include_team_memory is True

    def test_team_memory_excluded_for_research(self):
        plan = self.selector.select("research", team_id="team1")
        assert plan.include_team_memory is False

    def test_team_memory_excluded_when_no_team_id(self):
        plan = self.selector.select("deployment", team_id="")
        assert plan.include_team_memory is False

    def test_context_level_debug(self):
        plan = self.selector.select("debugging", session_id="s1")
        assert plan.context_level == 2   # summary + timeline

    def test_context_level_research(self):
        plan = self.selector.select("research", session_id="s1")
        assert plan.context_level == 1   # summary only

    def test_enabled_layers(self):
        plan = self.selector.select("deployment", user_id="alice",
                                     session_id="s1", agent_id="a1")
        layers = plan.enabled_layers()
        assert "experience" in layers
        assert "user_memory" in layers
        assert "context" in layers

    def test_selection_plan_to_dict(self):
        plan = self.selector.select("coding", user_id="alice")
        d = plan.to_dict()
        assert d["task_type"] == "coding"
        assert d["user_id"] == "alice"
        assert "include_experience" in d


# ══════════════════════════════════════════════════════════════════════════════
# §10 — PromptBuilder
# ══════════════════════════════════════════════════════════════════════════════

class TestPromptBuilder:
    def setup_method(self):
        self.builder = PromptBuilder()

    def test_build_basic(self):
        result = self.builder.build(user_input="Deploy my backend")
        assert "Deploy my backend" in result["context"]
        assert result["token_usage"] > 0

    def test_build_with_all_blocks(self):
        result = self.builder.build(
            user_input="Deploy backend",
            user_memory_block="USER PROFILE\n- prefers FastAPI",
            context_block="CONTEXT SUMMARY\n- Goal: deploy",
            task_memory_block="TASK MEMORY\n- service ready",
            experience_block="EXPERIENCE\n- deploy|fastapi",
            budget={"user_memory": 300, "context": 400,
                    "task_memory": 300, "experience": 800},
        )
        assert "USER PROFILE" in result["context"]
        assert "CONTEXT SUMMARY" in result["context"]
        assert "TASK MEMORY" in result["context"]
        assert "EXPERIENCE" in result["context"]
        assert "Deploy backend" in result["context"]

    def test_build_ordering(self):
        """PRD §10: USER PROFILE → CONTEXT SUMMARY → TASK MEMORY → EXPERIENCE → USER INPUT"""
        result = self.builder.build(
            user_input="USER INPUT",
            user_memory_block="USER PROFILE",
            context_block="CONTEXT SUMMARY",
            task_memory_block="TASK MEMORY",
            experience_block="EXPERIENCE",
        )
        ctx = result["context"]
        assert ctx.index("USER PROFILE") < ctx.index("CONTEXT SUMMARY")
        assert ctx.index("CONTEXT SUMMARY") < ctx.index("TASK MEMORY")
        assert ctx.index("TASK MEMORY") < ctx.index("EXPERIENCE")
        assert ctx.index("EXPERIENCE") < ctx.index("USER INPUT")

    def test_build_skips_empty_blocks(self):
        result = self.builder.build(
            user_input="Do something",
            user_memory_block="",
            context_block="",
            experience_block="",
        )
        assert "Do something" in result["context"]
        assert result["user_memory"] == ""

    def test_token_usage_reported(self):
        result = self.builder.build(
            user_input="task",
            experience_block="EXPERIENCE\n- some experience",
        )
        assert result["token_usage"] > 0

    def test_layer_tokens_reported(self):
        result = self.builder.build(
            user_input="task",
            user_memory_block="USER PROFILE\n- pref1\n- pref2",
            experience_block="EXPERIENCE\n- workflow",
        )
        assert "user_memory" in result["layer_tokens"]
        assert "experience" in result["layer_tokens"]

    def test_budget_truncation(self):
        """Blocks exceeding budget should be truncated."""
        long_experience = "EXPERIENCE\n" + "word " * 1000
        result = self.builder.build(
            user_input="task",
            experience_block=long_experience,
            budget={"experience": 20},  # very tight
        )
        # Should be truncated
        assert len(result["experience"].split()) < 30

    def test_system_block_prepended(self):
        result = self.builder.build(
            user_input="Do task",
            system_block="SYSTEM SOUL",
        )
        assert result["context"].startswith("SYSTEM SOUL")

    def test_output_keys(self):
        result = self.builder.build(user_input="task")
        assert "context" in result
        assert "user_memory" in result
        assert "experience" in result
        assert "token_usage" in result
        assert "layer_tokens" in result


# ══════════════════════════════════════════════════════════════════════════════
# §5 — MemoryOrchestrator (Full Pipeline)
# ══════════════════════════════════════════════════════════════════════════════

class TestMemoryOrchestrator:
    @pytest.fixture
    def orchestrator(self, tmp_path):
        return MemoryOrchestrator(
            db_path=str(tmp_path / "test.db"),
            max_tokens=2000,
        )

    def test_build_prompt_basic(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Deploy my FastAPI backend",
        )
        assert "context" in result
        assert result["task_type"] == "deployment"
        assert result["token_usage"] >= 0

    def test_build_prompt_returns_all_keys(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Write a function",
            user_id="alice",
            session_id="s1",
        )
        required_keys = {"context", "task_type", "mode", "token_usage",
                         "latency_ms", "plan", "budget", "ranked_layers"}
        assert required_keys.issubset(result.keys())

    def test_build_prompt_deployment_mode(self, orchestrator):
        result = orchestrator.build_prompt(task="Deploy to kubernetes")
        assert result["task_type"] == "deployment"
        assert result["mode"] == "execution"

    def test_build_prompt_debug_mode(self, orchestrator):
        result = orchestrator.build_prompt(task="Debug the error in my code")
        assert result["task_type"] == "debugging"
        assert result["mode"] == "debug"

    def test_build_prompt_research_mode(self, orchestrator):
        result = orchestrator.build_prompt(task="Research vector databases")
        assert result["task_type"] == "research"
        assert result["mode"] == "exploration"

    def test_task_type_override(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Deploy something",
            task_type="debugging",  # explicit override
        )
        assert result["task_type"] == "debugging"

    def test_mode_override(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Deploy something",
            mode="exploration",  # explicit mode override
        )
        assert result["mode"] == "exploration"

    def test_latency_under_threshold(self, orchestrator):
        """PRD §15: Latency < 80ms."""
        result = orchestrator.build_prompt(task="Deploy backend")
        assert result["latency_ms"] < 500  # generous threshold for CI

    def test_budget_allocated(self, orchestrator):
        result = orchestrator.build_prompt(task="Deploy backend")
        budget = result["budget"]
        assert sum(budget.values()) == 2000
        assert "experience" in budget
        assert "context" in budget

    def test_ranked_layers_complete(self, orchestrator):
        result = orchestrator.build_prompt(task="Deploy backend")
        assert set(result["ranked_layers"]) == set(MEMORY_LAYERS)

    def test_plan_summary(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Deploy backend",
            user_id="alice",
            session_id="s1",
        )
        plan = result["plan"]
        assert "task_type" in plan
        assert "layers" in plan

    def test_classify_shortcut(self, orchestrator):
        assert orchestrator.classify("Deploy to kubernetes") == "deployment"
        assert orchestrator.classify("Debug the error") == "debugging"

    def test_get_mode_shortcut(self, orchestrator):
        assert orchestrator.get_mode("deployment") == "execution"
        assert orchestrator.get_mode("debugging") == "debug"

    def test_user_input_in_context(self, orchestrator):
        """User input must always appear in context."""
        result = orchestrator.build_prompt(
            task="Deploy backend",
            user_input="My specific deploy request",
        )
        assert "My specific deploy request" in result["context"]

    def test_system_block_in_context(self, orchestrator):
        result = orchestrator.build_prompt(
            task="Deploy backend",
            system_block="SYSTEM SOUL",
        )
        assert result["context"].startswith("SYSTEM SOUL")


# ══════════════════════════════════════════════════════════════════════════════
# §14 — Adapter Integration
# ══════════════════════════════════════════════════════════════════════════════

class TestAdapterOrchestratorIntegration:
    @pytest.fixture
    def adapter(self, tmp_path):
        from openmemo.adapters.base_adapter import BaseMemoryAdapter
        return BaseMemoryAdapter(
            db_path=str(tmp_path / "test.db"),
            agent_id="agent1",
            user_id="alice",
        )

    @pytest.fixture
    def openclaw(self, tmp_path):
        from openmemo.adapters.openclaw import OpenClawMemoryBackend
        return OpenClawMemoryBackend(
            db_path=str(tmp_path / "test.db"),
            agent_id="claw1",
            conversation_id="session1",
        )

    def test_adapter_has_orchestrator(self, adapter):
        assert adapter._orchestrator is not None

    def test_orchestrate_prompt_returns_dict(self, adapter):
        result = adapter.orchestrate_prompt(task="Deploy backend")
        assert isinstance(result, dict)
        assert "context" in result
        assert "task_type" in result

    def test_orchestrate_prompt_deployment(self, adapter):
        result = adapter.orchestrate_prompt(task="Deploy my FastAPI app")
        assert result["task_type"] == "deployment"

    def test_orchestrate_prompt_debug(self, adapter):
        result = adapter.orchestrate_prompt(task="Debug the error in auth")
        assert result["task_type"] == "debugging"

    def test_orchestrate_prompt_with_session(self, adapter):
        result = adapter.orchestrate_prompt(
            task="Deploy backend",
            session_id="session_abc",
        )
        assert isinstance(result["context"], str)

    def test_openclaw_has_orchestrator(self, openclaw):
        assert openclaw._orchestrator is not None

    def test_openclaw_build_orchestrated_prompt(self, openclaw):
        result = openclaw.build_orchestrated_prompt(
            task="Deploy my service",
            session_id="session1",
        )
        assert isinstance(result, dict)
        assert "context" in result
        assert result["task_type"] == "deployment"

    def test_openclaw_build_orchestrated_prompt_debug(self, openclaw):
        result = openclaw.build_orchestrated_prompt(
            task="Fix the connection error",
        )
        assert result["task_type"] == "debugging"
