"""Configuration 4: Cohort-scored multi-agent orchestration.

Full 5-dimension scoring with stakeholder gating, recency penalty,
contribution bonus, phase detection, and topic-shift re-engagement.
This is the condition under test.
"""

from __future__ import annotations

import time
from typing import Any

from cohort.benchmark.code_runner import extract_code, run_code_against_test
from cohort.benchmark.config import StrategyResult
from cohort.benchmark.strategies.base import Strategy
from cohort.chat import ChatManager
from cohort.meeting import extract_keywords
from cohort.orchestrator import Orchestrator, TurnMode
from cohort.registry import JsonFileStorage

# Agent definitions with full scoring metadata
AGENT_CONFIGS = {
    "algorithm_designer": {
        "triggers": [
            "algorithm", "approach", "complexity", "design", "optimal",
            "dynamic", "greedy", "graph", "tree", "sort", "search",
            "math", "number", "formula", "pattern",
        ],
        "capabilities": [
            "algorithm design", "complexity analysis",
            "mathematical reasoning", "problem decomposition",
        ],
        "domain_expertise": [
            "competitive programming", "algorithms", "data structures",
        ],
        "complementary_agents": ["implementer"],
        "data_sources": [],
        "phase_roles": {"DISCOVER": "high", "PLAN": "high", "EXECUTE": "low"},
    },
    "implementer": {
        "triggers": [
            "implement", "code", "python", "function", "write",
            "build", "create", "solution", "program", "stdin", "stdout",
        ],
        "capabilities": [
            "Python implementation", "clean code",
            "stdin/stdout handling", "code optimization",
        ],
        "domain_expertise": [
            "python", "competitive programming", "code optimization",
        ],
        "complementary_agents": ["algorithm_designer", "debugger"],
        "data_sources": [],
        "phase_roles": {"EXECUTE": "high", "PLAN": "low"},
    },
    "debugger": {
        "triggers": [
            "test", "debug", "fix", "edge", "error", "wrong",
            "timeout", "overflow", "corner", "validate", "check",
            "review", "verify", "bug",
        ],
        "capabilities": [
            "edge case analysis", "debugging", "code review",
            "test case generation", "error diagnosis",
        ],
        "domain_expertise": [
            "testing", "debugging", "competitive programming",
        ],
        "complementary_agents": ["implementer"],
        "data_sources": ["test_results"],
        "phase_roles": {"VALIDATE": "high", "EXECUTE": "medium", "DISCOVER": "low"},
    },
}

AGENT_SYSTEMS = {
    "algorithm_designer": (
        "You are an algorithm designer on a competitive programming team. "
        "Analyze the problem, identify the optimal approach, and explain it clearly. "
        "Focus on correctness and complexity. If code exists, review the algorithm choice."
    ),
    "implementer": (
        "You are a Python implementer. Write clean, correct code that reads from stdin "
        "and writes to stdout. Follow the algorithm approach if one was proposed. "
        "Always output a complete ```python ... ``` code block."
    ),
    "debugger": (
        "You are a debugger and tester. Review the current solution for bugs, "
        "trace through edge cases, check boundary conditions. If you find issues, "
        "explain them and suggest fixes. If asked to fix, output corrected ```python ... ``` code."
    ),
}

MAX_TURNS = 12


class CohortScored(Strategy):
    """Cohort-scored: full orchestration with 5-dimension scoring."""

    NAME = "cohort_scored"

    # Set externally before run() to use a separate server for distillation
    distill_backend: Any = None

    def run(self) -> StrategyResult:
        t0 = time.time()
        failures: list[str] = []

        # Register distill backend (separate server or fallback to main LLM)
        from cohort.agent_context import set_distill_backend
        set_distill_backend(self.distill_backend or self.llm)

        # Set up isolated Cohort infrastructure
        import tempfile
        from pathlib import Path

        tmp_dir = Path(tempfile.mkdtemp(prefix="cohort_bench_"))
        storage = JsonFileStorage(tmp_dir)
        chat = ChatManager(storage)

        channel_id = f"bench_{self.task.task_id}"
        chat.create_channel(name=channel_id, description=self.task.title)

        # Create orchestrator with agent configs
        orchestrator = Orchestrator(
            chat=chat,
            agents=AGENT_CONFIGS,
        )

        # Start scored session
        session = orchestrator.start_session(
            channel_id=channel_id,
            topic=f"Solve: {self.task.title} - {self.task.description[:200]}",
            initial_agents=list(AGENT_CONFIGS.keys()),
            turn_mode=TurnMode.GUIDED,
            max_turns=MAX_TURNS,
        )

        # Post the problem
        chat.post_message(
            channel_id=channel_id,
            sender="user",
            content=self.task.to_prompt(),
        )

        # Scored turn loop
        conversation: list[str] = [f"[user]: {self.task.to_prompt()}"]
        code = ""
        all_pass = False

        for turn_num in range(MAX_TURNS):
            # Ask Cohort who should speak next
            recommendation = orchestrator.get_next_speaker(session.session_id)
            if recommendation is None:
                # No agent passes threshold -- discussion naturally concluded
                break

            speaker = recommendation["recommended_speaker"]
            score = recommendation["relevance_score"]
            phase = recommendation.get("phase", "")

            # Build context via Cohort's channel context pipeline (with distill)
            from cohort.agent_context import build_channel_context

            history = build_channel_context(chat, channel_id, distill=True)
            user_msg = (
                f"## Problem\n{self.task.to_prompt()}\n\n"
                f"## Discussion So Far\n{history}\n\n"
                "Your turn. You were selected because: "
                f"{recommendation.get('reason', 'highest score')}. "
                f"Current phase: {phase}. Contribute to solving this problem."
            )

            # Add test feedback if we have code and this is the debugger
            if code and speaker == "debugger":
                errors = []
                for i, tc in enumerate(self.task.sample_cases):
                    result = run_code_against_test(code, tc, self.task.time_limit_seconds)
                    if not result.passed:
                        errors.append(
                            f"Test {i + 1}: expected={tc.expected_output!r}, "
                            f"got={result.actual!r}"
                            + (f" ({result.error})" if result.error else "")
                        )
                if errors:
                    user_msg += "\n\n## Test Results (FAILING)\n" + "\n".join(errors)
                else:
                    user_msg += "\n\n## Test Results: ALL SAMPLE TESTS PASSING"

            # Generate response
            response = self._call_llm(
                agent_id=speaker,
                system=AGENT_SYSTEMS[speaker],
                user=user_msg,
            )

            # Update the turn log with scoring data
            if self.turns:
                self.turns[-1].relevance_score = score
                self.turns[-1].phase = phase

            # Post to channel and record turn
            msg = chat.post_message(
                channel_id=channel_id,
                sender=speaker,
                content=response,
            )
            orchestrator.record_turn(
                session_id=session.session_id,
                speaker=speaker,
                message_id=msg.id,
                was_recommended=True,
            )
            conversation.append(f"[{speaker}]: {response}")

            # Check for code
            new_code = extract_code(response)
            if new_code:
                code = new_code

            # Early exit: if all sample tests pass after implementer/debugger turn
            if code and speaker in ("implementer", "debugger"):
                all_pass = True
                for tc in self.task.sample_cases:
                    result = run_code_against_test(code, tc, self.task.time_limit_seconds)
                    if not result.passed:
                        all_pass = False
                        break
                if all_pass:
                    break

        # End session
        orchestrator.end_session(session.session_id)

        # Detect failure modes
        if not code:
            failures.append("no_code_produced")
        loop_count = self._detect_loops(conversation)
        if loop_count > 0:
            failures.append(f"loops:{loop_count}")

        elapsed = time.time() - t0
        return self._build_result(code, elapsed, failures)

    def _detect_loops(self, conversation: list[str]) -> int:
        loops = 0
        prev_keywords: set[str] = set()
        for msg in conversation:
            keywords = set(extract_keywords(msg))
            if prev_keywords and keywords:
                overlap = len(keywords & prev_keywords) / max(len(keywords), 1)
                if overlap > 0.8:
                    loops += 1
            prev_keywords = keywords
        return loops
