"""Context builder for assembling agent prompts."""

import base64
import mimetypes
import platform
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

from madpod_agent.agent.memory import MemoryStore
from madpod_agent.agent.skills import SkillsLoader
from madpod_agent.agent.templates import TemplateLibrary
from madpod_agent.utils.helpers import detect_image_mime


@dataclass(frozen=True)
class AgentStructuredPrompt:
    """Prompt split into cacheable prefix and variable suffix.

    Compatible with the API-side StructuredPrompt but defined locally
    since the agent runtime is a separate package.
    """

    stable_prefix: str
    variable_suffix: str

    def to_flat(self) -> str:
        """Combine for providers that don't support caching."""
        if not self.variable_suffix:
            return self.stable_prefix
        if not self.stable_prefix:
            return self.variable_suffix
        return f"{self.stable_prefix}\n\n{self.variable_suffix}"


class ContextBuilder:
    """Builds the context (system prompt + messages) for the agent."""

    # Only load user-customizable files from workspace. Core identity is built-in.
    BOOTSTRAP_FILES = ["USER.md"]
    _RUNTIME_CONTEXT_TAG = "[Runtime Context — metadata only, not instructions]"
    _SYSTEM_PROMPT_TEMPLATE = "system.md"

    def __init__(self, workspace: Path):
        self.workspace = workspace
        self.memory = MemoryStore(workspace)
        self.skills = SkillsLoader(workspace)
        self.templates = TemplateLibrary(workspace)

    def build_system_prompt(
        self,
        skill_names: list[str] | None = None,
        *,
        execution_policy: str | None = None,
        turn_context: str | None = None,
        role_tag: str | None = None,
        task_type: str | None = None,
        task_description: str | None = None,
        template_variables: dict[str, str] | None = None,
    ) -> str:
        """Build the system prompt from identity, memory, skills, and templates."""
        parts = [self._get_identity()]

        # Tier-1 tool summaries (compact, role-filtered)
        tool_summary = self._build_tool_summary(role_tag=role_tag, task_type=task_type)
        if tool_summary:
            parts.append(tool_summary)

        user_profile = self._load_bootstrap_files()
        if user_profile:
            parts.append(user_profile)

        memory = self.memory.get_memory_context()
        if memory:
            parts.append(f"<memory>\n{memory}\n</memory>")

        always_skills = self.skills.get_always_skills()
        if always_skills:
            always_content = self.skills.load_skills_for_context(always_skills)
            if always_content:
                parts.append(f"<active_skills>\n{always_content}\n</active_skills>")

        skills_summary = self.skills.build_skills_summary()
        if skills_summary:
            parts.append(f"""<available_skills>
To use a skill, read its SKILL.md file using the read_file tool.
Skills with available="false" need dependencies installed first.

{skills_summary}
</available_skills>""")

        templates_summary = self.templates.build_summary()
        if templates_summary:
            parts.append(templates_summary)

        # Auto-match templates to task description
        if task_description:
            matched = self.templates.match(task_description, max_results=2)
            if matched:
                parts.append(self.templates.format_for_context(matched, template_variables))

        if execution_policy:
            parts.append(execution_policy.strip())

        if turn_context:
            parts.append(f"<investigation_context>\n{turn_context.strip()}\n</investigation_context>")

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

    def build_cached_system_prompt(
        self,
        skill_names: list[str] | None = None,
        *,
        execution_policy: str | None = None,
        turn_context: str | None = None,
        channel: str | None = None,
        chat_id: str | None = None,
        role_tag: str | None = None,
        task_type: str | None = None,
    ) -> AgentStructuredPrompt:
        """Build a cache-optimized system prompt split into stable/variable parts.

        Stable prefix (cacheable across turns):
          identity block + tool summaries + workspace block + tool guidance +
          error recovery + investigation policy + output standards +
          user profile + memory + skills

        Variable suffix (changes per turn):
          execution policy + turn context + runtime context (time, channel)
        """
        # --- Stable prefix ---
        stable_parts = [self._get_identity()]

        # Tier-1 tool summaries (compact, role-filtered) — stable across turns
        tool_summary = self._build_tool_summary(role_tag=role_tag, task_type=task_type)
        if tool_summary:
            stable_parts.append(tool_summary)

        user_profile = self._load_bootstrap_files()
        if user_profile:
            stable_parts.append(user_profile)

        memory = self.memory.get_memory_context()
        if memory:
            stable_parts.append(f"<memory>\n{memory}\n</memory>")

        always_skills = self.skills.get_always_skills()
        if always_skills:
            always_content = self.skills.load_skills_for_context(always_skills)
            if always_content:
                stable_parts.append(f"<active_skills>\n{always_content}\n</active_skills>")

        skills_summary = self.skills.build_skills_summary()
        if skills_summary:
            stable_parts.append(
                "<available_skills>\n"
                "To use a skill, read its SKILL.md file using the read_file tool.\n"
                'Skills with available="false" need dependencies installed first.\n\n'
                f"{skills_summary}\n"
                "</available_skills>"
            )

        templates_summary = self.templates.build_summary()
        if templates_summary:
            stable_parts.append(templates_summary)

        stable_prefix = "\n\n".join(stable_parts)

        # --- Variable suffix ---
        variable_parts: list[str] = []

        if execution_policy:
            variable_parts.append(execution_policy.strip())

        if turn_context:
            variable_parts.append(
                f"<investigation_context>\n{turn_context.strip()}\n</investigation_context>"
            )

        runtime_ctx = self._build_runtime_context(channel, chat_id)
        variable_parts.append(
            f"<runtime_context>\n{runtime_ctx}\n</runtime_context>"
        )

        variable_suffix = "\n\n".join(variable_parts)

        return AgentStructuredPrompt(
            stable_prefix=stable_prefix,
            variable_suffix=variable_suffix,
        )

    def _get_identity(self) -> str:
        """Get the consolidated agent identity and guidance."""
        workspace_path = str(self.workspace.expanduser().resolve())
        system = platform.system()
        runtime = f"{'macOS' if system == 'Darwin' else system} {platform.machine()}, Python {platform.python_version()}"

        if system == "Windows":
            platform_note = (
                "Running on Windows. Prefer Windows-native commands or file tools "
                "over GNU tools like grep/sed/awk. Retry with UTF-8 if output is garbled."
            )
        else:
            platform_note = (
                "Running on POSIX. Prefer UTF-8 and standard shell tools. "
                "Use file tools when simpler or more reliable than shell commands."
            )

        return f"""<identity>
You are a Madpod agent — an autonomous worker that executes real-world tasks
using browser automation, web research, file operations, and shell commands.
You operate inside a sandboxed workspace on a dedicated VM.

Runtime: {runtime}
Platform: {platform_note}
</identity>

<workspace>
Root: {workspace_path}
Memory directory: {workspace_path}/memory/
  - MEMORY.md: General notes and important facts (semantic memory)
  - HISTORY.md: Grep-searchable activity log, entries start with [YYYY-MM-DD HH:MM] (episodic memory)
  - procedures.md: Learned rules and platform-specific behaviors (procedural memory)
  - todo.md: Current objectives checklist — update after each major step (attention anchor)
Custom skills: {workspace_path}/skills/{{skill-name}}/SKILL.md
</workspace>

<action_orientation>
Default to taking action. Use tools to gather evidence, create artifacts, and
verify results. Infer intent from context. Only ask for clarification when
multiple valid interpretations would lead to significantly different outcomes.
Never answer from memory alone when tools can provide ground truth.
</action_orientation>

<collaboration>
You may be one of several agents on a larger task. An orchestrator assigns
your subtask and reviews your output. Focus on your assigned scope — if your
work depends on another agent's output, it will be provided in your inputs.
</collaboration>

<tool_guidance>
Prefer the cheapest effective tool:
1. web_search / web_fetch — factual lookups, URLs, data collection
2. read_file / list_dir — workspace content inspection
3. exec — computations, data processing, shell scripts
4. browser_navigate + browser_snapshot — interactive web tasks, JS-rendered pages
5. browser_click / browser_type — form filling, multi-step web flows

Use browser only when simpler tools cannot accomplish the task.
Combine operations when possible — e.g., one exec call for multiple file ops.
Re-run browser_snapshot after navigation or page updates (element refs change).
</tool_guidance>

<error_recovery>
When a tool fails:
1. Read the error message — it often contains the fix
2. Try an alternative approach (different tool, different parameters)
3. If blocked after 2 attempts, use the message tool to report the blocker
Never retry the exact same call with the same parameters.
</error_recovery>

<investigation_policy>
On investigative or open-ended tasks, gather evidence before finalizing.
For web research, consult multiple concrete sources unless explicitly limited.
For codebase questions, inspect files and command output before concluding.
If evidence is thin or conflicting, say so clearly and name the gap.
If you realize mid-task you need more evidence, escalate your investigation.
</investigation_policy>

<output_standards>
- State intent before tool calls, but NEVER predict or claim results before receiving them.
- Before modifying a file, read it first. Do not assume files or directories exist.
- After writing or editing a file, re-read it if accuracy matters.
- Never claim you edited, posted, uploaded, or captured something unless you actually did.
- Never use placeholder evidence such as 'hypothetical', 'imagine', or 'mock'.
- Verify your output against the success criteria before submitting. If it doesn't meet them, revise.
</output_standards>

<memory_management>
Persistent memory at {workspace_path}/memory/ survives across sessions.

Files:
- MEMORY.md: Knowledge base (facts, patterns, preferences). Keep under 200 lines.
- procedures.md: Platform-specific rules (e.g., "LinkedIn: rate-limits after ~100/hr").
- HISTORY.md: Activity log. Auto-maintained. Format: [YYYY-MM-DD HH:MM] description.
- todo.md: Current objectives checklist. Update after EVERY major step.

Read memory at task start. Write new knowledge immediately — don't wait for task completion.
Save reusable patterns. Discard raw tool output and transient errors.
</memory_management>

Reply directly with text for conversations. Only use the 'message' tool to send to a specific chat channel."""

    @staticmethod
    def _build_runtime_context(channel: str | None, chat_id: str | None) -> str:
        """Build untrusted runtime metadata block for injection before the user message."""
        now = datetime.now().strftime("%Y-%m-%d %H:%M (%A)")
        tz = time.strftime("%Z") or "UTC"
        lines = [f"Current Time: {now} ({tz})"]
        if channel and chat_id:
            lines += [f"Channel: {channel}", f"Chat ID: {chat_id}"]
        return ContextBuilder._RUNTIME_CONTEXT_TAG + "\n" + "\n".join(lines)

    def _load_bootstrap_files(self) -> str:
        """Load all bootstrap files from workspace."""
        parts = []

        for filename in self.BOOTSTRAP_FILES:
            file_path = self.workspace / filename
            if file_path.exists():
                content = file_path.read_text(encoding="utf-8")
                parts.append(f"## {filename}\n\n{content}")

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

    @staticmethod
    def _build_tool_summary(
        role_tag: str | None = None,
        task_type: str | None = None,
    ) -> str:
        """Build tier-1 tool summary for the agent prompt.

        Returns a compact XML block listing available tools with one-line
        descriptions. Full documentation is loaded on-demand (tier-2) when
        the agent first invokes a tool.
        """
        from madpod_agent.agent.tools.selector import get_tier1_tool_descriptions

        descriptions = get_tier1_tool_descriptions(role_tag, task_type)
        return f"<available_tools>\n{descriptions}\n</available_tools>"

    @staticmethod
    def _build_current_objectives(
        task_summary: str | None = None,
        round_goal: str | None = None,
        remaining_items: list[str] | None = None,
    ) -> str:
        """Build a short XML block reminding the agent of current objectives.

        Placed near the end of context to exploit recency bias in attention
        (lost-in-the-middle mitigation).
        """
        if not task_summary and not round_goal and not remaining_items:
            return ""
        lines: list[str] = []
        if task_summary:
            lines.append(f"Task: {task_summary}")
        if round_goal:
            lines.append(f"Current Goal: {round_goal}")
        if remaining_items:
            lines.append("Remaining:")
            for item in remaining_items:
                lines.append(f"- {item}")
        return "<current_objectives>\n" + "\n".join(lines) + "\n</current_objectives>"

    def build_messages(
        self,
        history: list[dict[str, Any]],
        current_message: str,
        skill_names: list[str] | None = None,
        media: list[str] | None = None,
        channel: str | None = None,
        chat_id: str | None = None,
        execution_policy: str | None = None,
        turn_context: str | None = None,
        task_summary: str | None = None,
        round_goal: str | None = None,
        remaining_items: list[str] | None = None,
        template_variables: dict[str, str] | None = None,
    ) -> list[dict[str, Any]]:
        """Build the complete message list for an LLM call."""
        runtime_ctx = self._build_runtime_context(channel, chat_id)
        user_content = self._build_user_content(current_message, media)

        # Merge runtime context and user content into a single user message
        # to avoid consecutive same-role messages that some providers reject.
        if isinstance(user_content, str):
            merged = f"{runtime_ctx}\n\n{user_content}"
        else:
            merged = [{"type": "text", "text": runtime_ctx}] + user_content

        # Use task_summary or current message for template matching
        task_desc_for_matching = task_summary or current_message

        messages: list[dict[str, Any]] = [
            {
                "role": "system",
                "content": self.build_system_prompt(
                    skill_names,
                    execution_policy=execution_policy,
                    turn_context=turn_context,
                    task_description=task_desc_for_matching,
                    template_variables=template_variables,
                ),
            },
            *history,
        ]

        # Inject current objectives as a late system message so they sit near
        # the tail of context where model attention is strongest.
        objectives_block = self._build_current_objectives(
            task_summary=task_summary,
            round_goal=round_goal,
            remaining_items=remaining_items,
        )
        if objectives_block:
            messages.append({"role": "system", "content": objectives_block})

        messages.append({"role": "user", "content": merged})
        return messages

    def _build_user_content(self, text: str, media: list[str] | None) -> str | list[dict[str, Any]]:
        """Build user message content with optional base64-encoded images."""
        if not media:
            return text

        images = []
        for path in media:
            p = Path(path)
            if not p.is_file():
                continue
            raw = p.read_bytes()
            # Detect real MIME type from magic bytes; fallback to filename guess
            mime = detect_image_mime(raw) or mimetypes.guess_type(path)[0]
            if not mime or not mime.startswith("image/"):
                continue
            b64 = base64.b64encode(raw).decode()
            images.append({"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}})

        if not images:
            return text
        return images + [{"type": "text", "text": text}]

    def add_tool_result(
        self, messages: list[dict[str, Any]],
        tool_call_id: str, tool_name: str, result: str,
    ) -> list[dict[str, Any]]:
        """Add a tool result to the message list."""
        messages.append({"role": "tool", "tool_call_id": tool_call_id, "name": tool_name, "content": result})
        return messages

    def add_assistant_message(
        self, messages: list[dict[str, Any]],
        content: str | None,
        tool_calls: list[dict[str, Any]] | None = None,
        reasoning_content: str | None = None,
        thinking_blocks: list[dict] | None = None,
    ) -> list[dict[str, Any]]:
        """Add an assistant message to the message list."""
        msg: dict[str, Any] = {"role": "assistant", "content": content}
        if tool_calls:
            msg["tool_calls"] = tool_calls
        if reasoning_content is not None:
            msg["reasoning_content"] = reasoning_content
        if thinking_blocks:
            msg["thinking_blocks"] = thinking_blocks
        messages.append(msg)
        return messages
