"""Agent context builder for Cohort prompt injection.

Loads an agent's persistent memory (learned facts, collaborators, active
tasks) and builds compact context blocks for injection into the agent's
prompt.  Also provides channel context and collaboration awareness
helpers so that agents see the conversation history and know who else
is in the channel.

Ported from the Cohort enterprise agent_router pipeline, adapted for
lite-mode MCP where the host LLM *is* the agent.
"""

from __future__ import annotations

import logging
import os
import re
from datetime import datetime
from typing import Any

from cohort.context_window import truncate_context

logger = logging.getLogger(__name__)

# =====================================================================
# Constants
# =====================================================================

CONTEXT_HISTORY_LIMIT = int(os.environ.get("COHORT_HISTORY_LIMIT", "50"))

# Senders that should not appear in collaboration lists
_NON_AGENT_SENDERS = frozenset({
    "user", "system", "claude_code", "admin", "human",
})

GROUNDING_RULES = """
GROUNDING RULES (apply to ALL responses):
- You are an AI agent in the Cohort multi-agent system.
  Your team consists ONLY of other AI agents in this system.
- DO NOT invent, reference, or delegate to human team members
  (no fake names like "David", "Sarah", etc.).
- DO NOT reference external tools, channels, or platforms
  that are not part of the Cohort system.
- When delegating or coordinating, refer ONLY to agents
  by their actual agent IDs.
- If you do not know who should handle something, say so --
  do not fabricate organizational structures.

DESIGN SIMPLICITY:
- Prefer concrete implementations over abstractions.
- Solve the immediate problem simply. Generalize only when
  forced by a real second use case.
- Each agent's contribution should solve a real problem, not
  add scaffolding for hypothetical futures.
"""


# =====================================================================
# Fact scoring (contextual selection from agent memory)
# =====================================================================

_NOISE_WORDS = frozenset(
    "the a an is are was were in on at to for of and or but with from by "
    "this that these those be have has had do does did will would should "
    "could can may might must i you he she it we they".split()
)

MAX_FACTS = 10


def _normalize_query(query: str) -> list[str]:
    """Split query into lowercase search terms, dropping noise words."""
    terms = re.findall(r"[a-z0-9_]+", query.lower())
    return [t for t in terms if t not in _NOISE_WORDS and len(t) > 1]


def _score_text(text: str, terms: list[str]) -> float:
    """Score text against query terms (0.0-1.0 term overlap)."""
    if not text or not terms:
        return 0.0
    text_words = set(re.findall(r"[a-z0-9_]+", text.lower()))
    matched = sum(1 for t in terms if t in text_words)
    return matched / len(terms)


def _recency_boost(timestamp_str: str | None, max_boost: float = 0.2) -> float:
    """Add up to max_boost for recent entries. Decays over 90 days."""
    if not timestamp_str:
        return 0.0
    try:
        ts = timestamp_str.replace("Z", "+00:00")
        if "T" in ts:
            dt = datetime.fromisoformat(ts.split("+")[0])
        else:
            dt = datetime.strptime(ts[:10], "%Y-%m-%d")
        days_ago = max(0, (datetime.now() - dt).days)
        decay = max(0.0, 1.0 - (days_ago / 90))
        return decay * max_boost
    except (ValueError, TypeError):
        return 0.0


def _is_headline_only(fact: dict[str, Any]) -> bool:
    """Detect shallow temporal facts that are just article headlines."""
    text = fact.get("fact", "")
    source = fact.get("learned_from", "")
    if "temporal_facts_injector" in source:
        if len(text) < 150 and ("(via " in text or text.startswith("[")):
            return True
    return False


def _select_facts(
    facts: list[dict[str, Any]],
    query: str,
) -> list[dict[str, Any]]:
    """Select the most relevant learned facts for the current query.

    Uses term-overlap scoring + recency/confidence boosts.
    Falls back to most recent substantive facts if scoring produces
    fewer than 3 hits.
    """
    terms = _normalize_query(query)
    if not terms:
        substantive = [f for f in facts if not _is_headline_only(f)]
        return (substantive or facts)[-MAX_FACTS:]

    scored = []
    for fact in facts:
        fact_text = fact.get("fact", "")
        base = _score_text(fact_text, terms)
        if base > 0.0:
            conf_boost = {"high": 0.1, "medium": 0.0, "low": -0.05}.get(
                fact.get("confidence", "medium"), 0.0
            )
            rec = _recency_boost(fact.get("timestamp"))
            score = min(1.0, base + conf_boost + rec)
        else:
            score = 0.0
        scored.append((score, fact))

    scored.sort(key=lambda x: x[0], reverse=True)
    top = [f for s, f in scored[:MAX_FACTS * 2] if s > 0.0 and not _is_headline_only(f)]
    top = top[:MAX_FACTS]

    # Pad with most recent substantive facts if fewer than 3 scored hits
    if len(top) < 3:
        scored_ids = {id(f) for f in top}
        substantive_pad = [
            f for f in reversed(facts)
            if id(f) not in scored_ids and not _is_headline_only(f)
        ]
        top += substantive_pad[: MAX_FACTS - len(top)]

    return top


# =====================================================================
# Agent memory context
# =====================================================================

def load_agent_context(
    agent_id: str,
    query: str = "",
    agent_store: Any | None = None,
) -> str:
    """Build a context block from the agent's persistent memory.

    Includes contextually-selected learned facts, active tasks, and
    collaborator awareness.
    """
    if agent_store is None:
        return ""

    memory = agent_store.load_memory(agent_id)
    if memory is None:
        return ""

    parts: list[str] = []

    # Learned facts -- contextually selected
    facts_raw = [f.to_dict() if hasattr(f, "to_dict") else f for f in memory.learned_facts]
    if facts_raw:
        selected = _select_facts(facts_raw, query)
        if selected:
            lines = []
            for f in selected:
                conf = f.get("confidence", "medium")
                source = f.get("learned_from", "unknown")
                lines.append(f"- {f['fact']} ({source}, {conf} confidence)")
            parts.append("## Your Knowledge\n" + "\n".join(lines))

    # Active tasks
    if memory.active_tasks:
        lines = []
        for task in memory.active_tasks:
            status = task.get("status", "unknown")
            collab = task.get("collaborator", "")
            desc = task.get("task", task.get("description", ""))
            line = f"- {desc} (status: {status})"
            if collab:
                line += f" with @{collab}"
            lines.append(line)
        parts.append("## Your Active Tasks\n" + "\n".join(lines))

    # Collaborators
    if memory.collaborators:
        collabs = []
        for name, info in memory.collaborators.items():
            rel = info.get("relationship", "")
            collabs.append(f"@{name}" + (f" ({rel})" if rel else ""))
        if collabs:
            parts.append("## Your Collaborators\n" + ", ".join(collabs))

    if not parts:
        return ""

    return "=== AGENT MEMORY ===\n" + "\n\n".join(parts) + "\n=== END AGENT MEMORY ===\n"


# =====================================================================
# Distillation (compress channel context via local LLM)
# =====================================================================

# Cached router instance -- lazily created on first distill call
_local_router: Any | None = None
_local_router_checked: bool = False

# Optional external LLM backend for distillation (set by benchmark or other callers)
_external_llm_backend: Any | None = None


def set_distill_backend(backend: Any) -> None:
    """Set an external LLM backend for distillation.

    Allows callers (e.g. the benchmark harness) to provide a llama.cpp
    or other OpenAI-compatible backend for distillation when Ollama is
    not available.  The backend must implement
    ``generate(system: str, user: str, temperature: float) -> LLMResponse``.
    """
    global _external_llm_backend  # noqa: PLW0603
    _external_llm_backend = backend


def _get_local_router() -> Any | None:
    """Lazily create a LocalRouter if Ollama is available.

    Caches the result so the health check only runs once.
    Returns None if Ollama is not installed or not running.
    """
    global _local_router, _local_router_checked  # noqa: PLW0603
    if _local_router_checked:
        return _local_router
    _local_router_checked = True
    try:
        from cohort.local.router import LocalRouter
        router = LocalRouter()
        # Probe Ollama — if it's not running, don't attempt distillation
        if hasattr(router, "_ensure_client") and not router._ensure_client():
            logger.debug("[*] Ollama not available — distillation disabled")
            return None
        _local_router = router
        return _local_router
    except Exception:
        logger.debug("[*] LocalRouter unavailable — distillation disabled")
        return None


def _distill_via_external(raw_context: str) -> str | None:
    """Distill using the external LLM backend (llama.cpp, etc)."""
    if _external_llm_backend is None:
        return None
    try:
        from cohort.local.config import DISTILLATION_PROMPT
        prompt = DISTILLATION_PROMPT.format(qwen_output=raw_context)
        resp = _external_llm_backend.generate(
            system="You are a distillation engine. Compress the input into a structured briefing.",
            user=prompt,
            temperature=0.15,
        )
        if resp and resp.text.strip():
            return resp.text.strip()
    except Exception:
        logger.debug("[*] External LLM distillation failed")
    return None


def distill_context(raw_context: str) -> str:
    """Compress raw channel context through a local LLM distillation pipeline.

    Tries in order:
    1. Ollama via LocalRouter (original path)
    2. External LLM backend (llama.cpp, etc) if set via set_distill_backend()
    3. Falls back to raw text

    Parameters
    ----------
    raw_context:
        The formatted channel messages text to compress.
    """
    if not raw_context or len(raw_context) < 3000:
        # Too short to benefit from distillation
        return raw_context

    # Try Ollama first
    router = _get_local_router()
    if router is not None:
        briefing = router.distill(raw_context)
        if briefing:
            logger.info(
                "[OK] Distilled channel context (Ollama): %d -> %d chars (%.0f%% reduction)",
                len(raw_context), len(briefing),
                (1 - len(briefing) / len(raw_context)) * 100,
            )
            return briefing

    # Try external backend (llama.cpp, etc)
    briefing = _distill_via_external(raw_context)
    if briefing:
        logger.info(
            "[OK] Distilled channel context (external): %d -> %d chars (%.0f%% reduction)",
            len(raw_context), len(briefing),
            (1 - len(briefing) / len(raw_context)) * 100,
        )
        return briefing

    # Distillation failed — return raw context
    return raw_context


# =====================================================================
# Channel context
# =====================================================================

def build_channel_context(
    chat: Any,
    channel_id: str,
    *,
    distill: bool = True,
) -> str:
    """Build a context string with recent channel messages.

    Fetches up to CONTEXT_HISTORY_LIMIT messages, applies sliding-window
    truncation, and optionally distills through the local LLM.

    When Ollama is running, the raw messages are compressed into a
    structured briefing (matching the enterprise smartest pipeline).
    When Ollama is not available, the raw messages are returned as-is.

    Parameters
    ----------
    chat:
        :class:`~cohort.chat.ChatManager` instance.
    channel_id:
        Channel to read from.
    distill:
        If True (default), attempt to compress via local LLM.
    """
    channel = chat.get_channel(channel_id)
    if not channel:
        return ""

    header_parts = [
        f"Channel: #{channel.name}",
    ]
    if channel.description:
        header_parts.append(f"Description: {channel.description}")
    if channel.topic:
        header_parts.append(f"Topic: {channel.topic}")

    messages = chat.get_channel_messages(channel_id, limit=CONTEXT_HISTORY_LIMIT)
    if not messages:
        return ""

    messages = truncate_context(messages)

    # Format raw messages
    msg_lines: list[str] = []
    for msg in messages:
        ts = msg.timestamp.split("T")[1][:8] if "T" in msg.timestamp else msg.timestamp
        limit = 6000 if msg.sender == "system" else 2000
        msg_lines.append(f"[{ts}] {msg.sender}: {msg.content[:limit]}")

    raw_messages = "\n".join(msg_lines)

    # Attempt distillation if enabled
    if distill:
        distilled = distill_context(raw_messages)
        if distilled != raw_messages:
            # Distillation succeeded — use briefing format
            parts = [
                "=== CHANNEL BRIEFING ===",
                *header_parts,
                "",
                "--- Distilled Conversation Summary ---",
                distilled,
                "=== END CHANNEL BRIEFING ===\n",
            ]
            return "\n".join(parts)

    # No distillation — return raw messages
    parts = [
        "=== CHANNEL CONTEXT ===",
        *header_parts,
        "\n--- Recent Messages ---",
        raw_messages,
        "=== END CHANNEL CONTEXT ===\n",
    ]
    return "\n".join(parts)


# =====================================================================
# Collaboration awareness
# =====================================================================

def build_collaboration_section(
    chat: Any,
    channel_id: str,
    agent_id: str,
) -> str:
    """Build a collaboration section listing other agents in the channel.

    Encourages the agent to @mention peers for coordination.
    """
    channel = chat.get_channel(channel_id)
    if not channel:
        return ""

    members = getattr(channel, "members", None) or []
    other_agents = [
        m for m in members
        if m != agent_id and m not in _NON_AGENT_SENDERS
    ]

    if not other_agents:
        # Fall back to recent senders if channel has no explicit member list
        messages = chat.get_channel_messages(channel_id, limit=50)
        senders = {m.sender for m in messages} if messages else set()
        other_agents = sorted(
            s for s in senders
            if s != agent_id and s not in _NON_AGENT_SENDERS
        )

    if not other_agents:
        return ""

    member_list = ", ".join(f"@{m}" for m in other_agents)
    return (
        "=== COLLABORATION ===\n"
        f"Other agents in this channel: {member_list}\n"
        "When your response would benefit from another agent's expertise, "
        "tag them with @agent_name to bring them into the conversation.\n"
        "=== END COLLABORATION ===\n"
    )


# =====================================================================
# Full context assembler
# =====================================================================

def build_full_agent_context(
    chat: Any,
    agent_store: Any,
    agent_id: str,
    channel_id: str | None = None,
    query: str = "",
) -> str:
    """Assemble the complete agent context stack.

    Combines agent memory, grounding rules, collaboration awareness,
    and channel context into a single string suitable for appending
    to an agent persona prompt.

    Parameters
    ----------
    chat:
        :class:`~cohort.chat.ChatManager` instance.
    agent_store:
        :class:`~cohort.agent_store.AgentStore` instance.
    agent_id:
        Agent receiving the context.
    channel_id:
        Channel to pull context from (optional).
    query:
        Current message/topic for contextual fact selection.
    """
    parts: list[str] = []

    # 1. Agent memory (facts, tasks, collaborators)
    memory_block = load_agent_context(agent_id, query=query, agent_store=agent_store)
    if memory_block:
        parts.append(memory_block)

    # 2. Grounding rules
    parts.append(GROUNDING_RULES.strip())

    # 3. Collaboration awareness
    if channel_id:
        collab = build_collaboration_section(chat, channel_id, agent_id)
        if collab:
            parts.append(collab)

    # 4. Channel context (recent messages)
    if channel_id:
        channel_ctx = build_channel_context(chat, channel_id)
        if channel_ctx:
            parts.append(channel_ctx)

    return "\n\n".join(parts)
