"""Context Engine — the central intelligence of the system.

Builds structured prompts from session state. Never includes raw tool
outputs. Handles compaction, artifact summarization, and selective
rehydration.
"""

from __future__ import annotations

import logging
import time
from collections import defaultdict
from typing import Any

from ..config import settings

logger = logging.getLogger(__name__)
from ..models.agent import AgentProfile
from ..models.artifacts import ArtifactSummary
from ..memory.embeddings import EmbeddingService
from ..memory.store import MemoryStore
from ..models.context import (
    ContextPackage,
    ContextSection,
    ContextSectionType,
    MemoryDocument,
    MemoryType,
)
from ..models.session import SessionState, TaskState
from .compactor import ContextCompactor, estimate_tokens


class ContextEngine:
    """Assembles the context package that gets sent to the model.

    The engine enforces a strict contract:
    - Raw tool outputs NEVER appear in the context
    - Each section has a priority for compaction
    - Immutable rules are always included in full
    """

    def __init__(
        self,
        compactor: ContextCompactor | None = None,
        memory_store: MemoryStore | None = None,
    ) -> None:
        self.compactor = compactor or ContextCompactor(
            max_tokens=settings.context_max_tokens
        )
        self.memory = memory_store
        self._embed_service: EmbeddingService | None = None
        # Debug history: last N context builds per session
        self._history: dict[str, list[dict[str, Any]]] = defaultdict(list)
        self._max_history = 20
        # Full context of the LAST build per session (not accumulated)
        self._last_full_context: dict[str, dict[str, Any]] = {}

    # ------------------------------------------------------------------
    # Public — build context for a model call
    # ------------------------------------------------------------------

    async def build_context(
        self,
        session: SessionState,
        agent: AgentProfile,
        artifacts: list[ArtifactSummary] | None = None,
        conversation: list[dict[str, Any]] | None = None,
        extra_instructions: str = "",
    ) -> ContextPackage:
        """Build a full ContextPackage for the given agent and session.

        The conversation parameter contains real assistant/tool messages
        with complete tool results. These go into the messages array,
        not the system prompt — like professional agentic tools.
        """

        sections: list[ContextSection] = []
        allowed = set(agent.context_sections)

        # 1. Immutable rules — highest priority, never trimmed
        if "immutable_rules" in allowed:
            sections.append(self._build_immutable_rules(session, agent))

        # 2. Project profile
        if "project_profile" in allowed:
            sections.append(self._build_project_profile(session))

        # 3. Knowledge base — loaded from memory store
        if "knowledge_base" in allowed and self.memory:
            kb_section = await self._build_knowledge_base(session)
            if kb_section:
                sections.append(kb_section)

        # 4. Task history — compact summaries of past tasks in this session
        if "task_state" in allowed and session.task_history:
            sections.append(self._build_task_history(session))

        # 5. Task state — current task (includes compacted previous steps)
        if "task_state" in allowed and session.current_task:
            sections.append(self._build_task_state(session.current_task))

        # Compact to fit budget
        sections = self.compactor.compact_sections(sections)

        # Assemble system prompt from sections
        system_prompt = self._assemble_system_prompt(sections)

        # Build messages with real conversation history
        messages = self._build_messages(session, conversation)

        total_tokens = estimate_tokens(system_prompt) + sum(
            estimate_tokens(m.get("content", "")) for m in messages
        )

        package = ContextPackage(
            sections=sections,
            system_prompt=system_prompt,
            messages=messages,
            total_token_estimate=total_tokens,
        )

        # Guardar contexto completo del último build (solo el último por sesión)
        self._last_full_context[session.session_id] = {
            "system_prompt": system_prompt,
            "messages": messages,
            "total_tokens": total_tokens,
            "timestamp": time.time(),
        }

        # --- Debug: log and store context build ---
        section_summary = []
        for s in sections:
            section_summary.append({
                "type": s.section_type.value,
                "priority": s.priority,
                "tokens": s.token_estimate,
                "chars": len(s.content),
                "preview": s.content[:150].replace("\n", " "),
            })

        conv_len = len(conversation) if conversation else 0
        debug_entry = {
            "timestamp": time.time(),
            "agent": agent.role,
            "agent_name": agent.name,
            "total_tokens": total_tokens,
            "sections": section_summary,
            "sections_count": len(sections),
            "compacted": len(sections) < len(allowed),
            "system_prompt_tokens": estimate_tokens(system_prompt),
            "user_message_preview": messages[0]["content"][:200] if messages else "",
            "artifacts_count": len(artifacts) if artifacts else 0,
            "conversation_messages": conv_len,
        }

        history = self._history[session.session_id]
        history.append(debug_entry)
        if len(history) > self._max_history:
            self._history[session.session_id] = history[-self._max_history:]

        logger.info(
            "Context built for [%s/%s] — %d sections, ~%d tokens, artifacts=%d, conversation=%d msgs",
            session.session_id[:8],
            agent.role,
            len(sections),
            total_tokens,
            len(artifacts) if artifacts else 0,
            conv_len,
        )

        return package

    def get_debug_history(self, session_id: str) -> list[dict[str, Any]]:
        """Return the context build history for a session."""
        return list(self._history.get(session_id, []))

    def get_last_context_debug(self, session_id: str) -> dict[str, Any] | None:
        """Return the most recent context build for a session."""
        history = self._history.get(session_id, [])
        return history[-1] if history else None

    def get_last_full_context(self, session_id: str) -> dict[str, Any] | None:
        """Return the full context (system_prompt + messages) of the last build."""
        return self._last_full_context.get(session_id)

    def rehydrate_artifact(
        self,
        artifact: ArtifactSummary,
        full_content: str,
    ) -> ContextSection:
        """Selectively rehydrate an artifact into a full context section.

        Used when the agent explicitly needs the full content of a
        specific artifact (e.g. reviewing generated code).
        """
        content = (
            f"## Rehydrated Artifact: {artifact.title}\n"
            f"Type: {artifact.artifact_type} | Source: {artifact.source_tool}\n"
            f"---\n{full_content}\n---"
        )
        return ContextSection(
            section_type=ContextSectionType.WORKING_CONTEXT,
            content=content,
            priority=5,
            token_estimate=estimate_tokens(content),
        )

    def summarize_tool_output(
        self,
        tool_name: str,
        raw_output: str,
        session_id: str,
        task_id: str,
    ) -> ArtifactSummary:
        """Delegate to compactor — raw output never enters context."""
        return self.compactor.summarize_tool_output(
            tool_name=tool_name,
            raw_output=raw_output,
            session_id=session_id,
            task_id=task_id,
        )

    # ------------------------------------------------------------------
    # Section builders
    # ------------------------------------------------------------------

    def _build_immutable_rules(
        self, session: SessionState, agent: AgentProfile
    ) -> ContextSection:
        parts = [
            "# System Rules (Immutable)",
            "",
            agent.system_prompt,
            "",
        ]
        if session.immutable_rules:
            parts.append("## Session Rules")
            for rule in session.immutable_rules:
                parts.append(f"- {rule}")
        parts.extend(
            [
                "",
                "## Contrato de Contexto",
                "- Los resultados de herramientas se incluyen completos en la conversación.",
                "- Los steps anteriores pueden estar compactados como resúmenes.",
                "- Mantén las respuestas enfocadas en el paso actual.",
                "- Si ya tienes la información necesaria, genera tu respuesta final.",
                "- NO repitas llamadas a herramientas con los mismos argumentos.",
                "- Responde SIEMPRE en español.",
            ]
        )
        content = "\n".join(parts)
        return ContextSection(
            section_type=ContextSectionType.IMMUTABLE_RULES,
            content=content,
            priority=100,
            token_estimate=estimate_tokens(content),
        )

    def _build_project_profile(self, session: SessionState) -> ContextSection:
        if not session.project_profile:
            content = "# Project Profile\nNo project profile configured."
        else:
            lines = ["# Project Profile"]
            for key, value in session.project_profile.items():
                lines.append(f"- **{key}**: {value}")
            content = "\n".join(lines)
        return ContextSection(
            section_type=ContextSectionType.PROJECT_PROFILE,
            content=content,
            priority=80,
            token_estimate=estimate_tokens(content),
        )

    async def _build_knowledge_base(
        self, session: SessionState
    ) -> ContextSection | None:
        """Load relevant knowledge documents via semantic search.

        Uses embeddings to find the most relevant docs for the current
        task. Always includes a title index of ALL docs so the agent
        knows what exists and can request more.
        """
        if not self.memory:
            return None

        all_docs: list[MemoryDocument] = []
        all_docs.extend(await self.memory.list_documents(
            namespace="knowledge",
            memory_type=MemoryType.DOCUMENT,
        ))
        all_docs.extend(await self.memory.list_documents(
            namespace=f"knowledge:{session.session_id}",
            memory_type=MemoryType.DOCUMENT,
        ))

        if not all_docs:
            return None

        doc_map = {d.memory_id: d for d in all_docs}

        # Rank docs by semantic similarity
        query = self._build_search_query(session)
        ranked_ids: list[str] = []

        if query:
            ranked_ids = await self._semantic_rank(query)

        if not ranked_ids:
            # No embeddings or no task — sort by size (smallest first)
            ranked_ids = [
                d.memory_id
                for d in sorted(all_docs, key=lambda d: len(d.content))
            ]

        # Include ALL docs — 42K tokens fits well within model context (128K)
        max_kb_tokens = 50_000
        token_budget = max_kb_tokens
        full_docs: list[MemoryDocument] = []

        for doc_id in ranked_ids:
            doc = doc_map.get(doc_id)
            if not doc:
                continue
            doc_tokens = estimate_tokens(doc.content)
            if doc_tokens <= token_budget:
                full_docs.append(doc)
                token_budget -= doc_tokens

        # Build section — ALWAYS include title index of ALL docs
        included_ids = {d.memory_id for d in full_docs}
        not_included = [d for d in all_docs if d.memory_id not in included_ids]

        lines = [
            "# Knowledge Base",
            f"_{len(full_docs)} doc(s) loaded in full, "
            f"{len(not_included)} available on request_",
            "",
        ]

        for doc in full_docs:
            lines.append(f"## {doc.title}")
            lines.append(doc.content)
            lines.append("")

        if not_included:
            lines.append("## Other Available Docs")
            lines.append("_Ask for any of these if you need the full content:_")
            for doc in not_included:
                lines.append(f"- **{doc.title}** ({doc.memory_id}): {doc.summary[:150]}")
            lines.append("")

        content = "\n".join(lines)
        return ContextSection(
            section_type=ContextSectionType.KNOWLEDGE_BASE,
            content=content,
            priority=60,
            token_estimate=estimate_tokens(content),
        )

    async def _semantic_rank(self, query: str) -> list[str]:
        """Rank knowledge docs by cosine similarity to the query."""
        try:
            if not self._embed_service:
                self._embed_service = EmbeddingService()

            query_embedding = await self._embed_service.embed(query)
            results = await self.memory.search_by_similarity(
                query_embedding=query_embedding,
                namespace="knowledge",
                top_k=50,
            )
            return [doc_id for doc_id, _score in results]

        except Exception as e:
            logger.warning("Semantic search failed: %s — loading all docs", e)
            return []

    @staticmethod
    def _build_search_query(session: SessionState) -> str:
        """Build a natural language query from the current task."""
        if not session.current_task:
            return ""
        parts = [session.current_task.objective]
        step = session.current_task.current_step()
        if step:
            parts.append(step.description)
        if session.current_task.facts_extracted:
            parts.extend(session.current_task.facts_extracted[-5:])
        return " ".join(parts)

    def _build_task_history(self, session: SessionState) -> ContextSection:
        """Build a compact summary of past tasks in this session.

        Each completed task is ~50 tokens instead of hundreds.
        The agent retains awareness of what was done before.
        """
        lines = [
            "# Session History",
            f"_{len(session.task_history)} previous task(s) in this session_",
            "",
        ]

        for i, entry in enumerate(session.task_history):
            status = entry.get("status", "?")
            objective = entry.get("objective", "")[:100]
            summary = entry.get("summary", "")[:150]
            facts = entry.get("facts", [])

            lines.append(f"**Task {i + 1}** [{status}]: {objective}")
            if summary:
                lines.append(f"  Result: {summary}")
            if facts:
                lines.append(f"  Facts: {'; '.join(facts[:5])}")
            # Key structured data (recordNums, sectionIds, etc.)
            key_data = entry.get("key_data", {})
            if key_data:
                kd_parts = []
                for table, nums in key_data.get("tables", {}).items():
                    kd_parts.append(f"{table}: records {nums}")
                for page, num in key_data.get("pages", {}).items():
                    kd_parts.append(f"page '{page}' = record {num}")
                if key_data.get("sections"):
                    kd_parts.append(f"sections: {key_data['sections']}")
                if key_data.get("modules"):
                    kd_parts.append(f"modules: {key_data['modules']}")
                if kd_parts:
                    lines.append(f"  Key data: {'; '.join(kd_parts)}")
            review = entry.get("review", "")
            if review:
                lines.append(f"  Review: {review[:100]}")
            lines.append("")

        content = "\n".join(lines)
        return ContextSection(
            section_type=ContextSectionType.TASK_STATE,
            content=content,
            priority=55,  # Below knowledge (60), above artifacts (50)
            token_estimate=estimate_tokens(content),
        )

    def _build_task_state(self, task: TaskState) -> ContextSection:
        lines = [
            "# Current Task",
            f"**Objective**: {task.objective}",
            f"**Status**: {task.status}",
            f"**Step**: {task.current_step_index + 1}/{len(task.plan)}",
        ]

        current = task.current_step()
        if current:
            lines.extend(
                [
                    "",
                    "## Current Step",
                    f"- Description: {current.description}",
                    f"- Agent: {current.agent_role}",
                    f"- Status: {current.status}",
                ]
            )

        if task.facts_extracted:
            lines.append("")
            lines.append("## Established Facts")
            for fact in task.facts_extracted[-10:]:
                lines.append(f"- {fact}")

        if task.constraints:
            lines.append("")
            lines.append("## Constraints")
            for c in task.constraints:
                lines.append(f"- {c}")

        # Show compacted previous steps results
        compacted_steps = [s for s in task.plan if s.compacted and s.result_summary]
        if compacted_steps:
            lines.append("")
            lines.append("## Previous Steps (compacted)")
            for step in compacted_steps:
                lines.append(f"- [{step.agent_role}] {step.description}: {step.result_summary[:300]}")

        # Show plan overview (compact)
        if task.plan:
            lines.append("")
            lines.append("## Plan Overview")
            for i, step in enumerate(task.plan):
                marker = "→" if i == task.current_step_index else "·"
                status_label = step.status.value
                compacted_label = " (compacted)" if step.compacted else ""
                lines.append(
                    f"  {marker} Step {i + 1} [{status_label}{compacted_label}]: {step.description}"
                )

        content = "\n".join(lines)
        return ContextSection(
            section_type=ContextSectionType.TASK_STATE,
            content=content,
            priority=70,
            token_estimate=estimate_tokens(content),
        )

    def _build_artifact_memory(
        self, artifacts: list[ArtifactSummary]
    ) -> ContextSection:
        content = self.compactor.compact_artifact_summaries(
            artifacts, max_chars=settings.artifact_summary_max_chars
        )
        return ContextSection(
            section_type=ContextSectionType.ARTIFACT_MEMORY,
            content=content,
            priority=50,
            token_estimate=estimate_tokens(content),
        )

    # ------------------------------------------------------------------
    # Assembly
    # ------------------------------------------------------------------

    def _assemble_system_prompt(self, sections: list[ContextSection]) -> str:
        """Combine sections into a single system prompt string."""
        parts: list[str] = []
        order = [
            ContextSectionType.IMMUTABLE_RULES,
            ContextSectionType.PROJECT_PROFILE,
            ContextSectionType.KNOWLEDGE_BASE,
            ContextSectionType.TASK_STATE,
        ]
        section_map: dict[ContextSectionType, ContextSection] = {
            s.section_type: s for s in sections
        }
        for st in order:
            if st in section_map:
                parts.append(section_map[st].content)
        return "\n\n---\n\n".join(parts)

    def _build_messages(
        self,
        session: SessionState,
        conversation: list[dict[str, Any]] | None = None,
    ) -> list[dict[str, Any]]:
        """Build the messages array with real conversation history.

        Includes the user objective message followed by the full
        assistant/tool conversation — like professional agentic tools.
        """
        if session.current_task:
            step = session.current_task.current_step()
            if step:
                user_content = (
                    f"Execute this step: {step.description}\n"
                    f"Overall objective: {session.current_task.objective}"
                )
            else:
                user_content = session.current_task.objective
        else:
            user_content = "Awaiting task assignment."

        messages: list[dict[str, Any]] = []

        # Include previous task exchanges as compact conversation history
        if session.task_history:
            history_lines = ["[HISTORIAL DE CONVERSACIÓN ANTERIOR — NO ejecutar de nuevo, solo contexto]"]
            for entry in session.task_history[-10:]:
                objective = entry.get("objective", "")[:200]
                summary = entry.get("summary", "")
                key_data = entry.get("key_data", {})
                tools = entry.get("tools_used", [])

                history_lines.append(f"Usuario pidió: {objective}")
                if tools:
                    history_lines.append(f"  Tools usadas: {', '.join(tools[:5])}")
                if key_data:
                    kd_parts = []
                    for table, nums in key_data.get("tables", {}).items():
                        kd_parts.append(f"{table}: records {nums}")
                    if key_data.get("sections"):
                        kd_parts.append(f"sections: {key_data['sections'][:5]}")
                    if key_data.get("modules"):
                        kd_parts.append(f"modules: {key_data['modules'][:5]}")
                    if kd_parts:
                        history_lines.append(f"  Datos clave: {'; '.join(kd_parts)}")
                # Extract agent response from summary
                if " → Agent: " in summary:
                    agent_part = summary.split(" → Agent: ", 1)[1][:200]
                    history_lines.append(f"  Resultado: {agent_part}")
                history_lines.append("")

            messages.append({"role": "user", "content": "\n".join(history_lines)})
            messages.append({"role": "assistant", "content": "Entendido, tengo el contexto del historial. ¿En qué puedo ayudarte ahora?"})

        # Current user message
        messages.append({"role": "user", "content": user_content})

        # Append real conversation (assistant messages + tool results from current step)
        if conversation:
            messages.extend(conversation)

        return messages