"""
hollow.agent — multi-step AI agent loop powered by Hollow.

Model-agnostic: pass any ModelAdapter (Anthropic, OpenAI, or your own).
Defaults to AnthropicAdapter when no adapter is provided.
"""

from __future__ import annotations

import json
import os
import re
from typing import Callable, Optional, Protocol, Union, runtime_checkable

from hollow.client import HollowClient
from hollow.types import Action, AgentStep

_DEFAULT_MAX_STEPS = 15
_DEFAULT_START_URL = "https://www.startpage.com"

_SYSTEM_PROMPT = """\
You are an AI agent with access to Hollow, \
a serverless web intelligence layer.
You receive GDG Spatial maps — structured representations of web pages \
with actionable element IDs like [1], [2], [3].

On each step, reason then act:

EVALUATE: what did the last action accomplish?
MEMORY: what do I know so far?
GOAL: what is my immediate next step?
ACTION: exactly one JSON action object

Actions:
{ "type": "navigate", "url": "https://..." }
{ "type": "click", "elementId": 3 }
{ "type": "fill", "elementId": 4, "value": "text" }
{ "type": "scroll", "direction": "down" }
{ "type": "done", "result": "final answer here" }

When the task is complete respond with done.
Use startpage.com to search when you need to find a URL.\
"""


# ── ModelAdapter protocol ───────────────────────────────────────────────────────

@runtime_checkable
class ModelAdapter(Protocol):
    """
    Implement this protocol to plug any LLM into the Hollow agent loop.

    ``messages`` is a list of dicts with ``role`` and ``content`` keys.
    Roles: ``"system"``, ``"user"``, ``"assistant"``.

    Example::

        class MyAdapter:
            def complete(self, messages: list[dict]) -> str:
                return my_model.chat(messages)

        run_agent(client, task="...", model=MyAdapter())
    """

    def complete(self, messages: list[dict]) -> str: ...


# ── Built-in adapters ──────────────────────────────────────────────────────────

class AnthropicAdapter:
    """
    Default adapter. Uses the Anthropic API via the ``anthropic`` package.

    Reads ``ANTHROPIC_API_KEY`` from the environment unless ``api_key`` is given.

    Args:
        api_key: Anthropic API key. Falls back to ``ANTHROPIC_API_KEY`` env var.
        model:   Anthropic model ID. Default: ``claude-opus-4-6``.

    Example::

        from hollow import AnthropicAdapter, run_agent

        result = run_agent(client, task="...",
            model=AnthropicAdapter(model="claude-opus-4-6"))
    """

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "claude-opus-4-6",
    ) -> None:
        self._api_key = api_key
        self._model = model

    def complete(self, messages: list[dict]) -> str:
        try:
            import anthropic as _anthropic
        except ImportError as exc:
            raise ImportError(
                "AnthropicAdapter requires the 'anthropic' package. "
                "Install it with: pip install anthropic"
            ) from exc

        api_key = self._api_key or os.environ.get("ANTHROPIC_API_KEY")
        if not api_key:
            raise EnvironmentError(
                "AnthropicAdapter requires ANTHROPIC_API_KEY in the environment "
                "or the api_key constructor argument."
            )

        client = _anthropic.Anthropic(api_key=api_key)

        # Anthropic takes system as a top-level param, not a message
        system_parts = [m["content"] for m in messages if m["role"] == "system"]
        turns = [m for m in messages if m["role"] in ("user", "assistant")]
        system = "\n\n".join(system_parts) if system_parts else None

        kwargs: dict = dict(
            model=self._model,
            max_tokens=1024,
            messages=turns,
        )
        if system:
            kwargs["system"] = system

        response = client.messages.create(**kwargs)
        return next((b.text for b in response.content if b.type == "text"), "")


class OpenAIAdapter:
    """
    Adapter for OpenAI and any OpenAI-compatible API (Groq, Together, Ollama, etc.).

    Requires the ``openai`` package: ``pip install hollow-sdk[openai]``
    Reads ``OPENAI_API_KEY`` from the environment unless ``api_key`` is given.

    Args:
        api_key:  OpenAI API key (or provider key). Falls back to ``OPENAI_API_KEY``.
        model:    Model ID. Default: ``gpt-4o``.
        base_url: Base URL for OpenAI-compatible providers.

    Example::

        from hollow import OpenAIAdapter, run_agent

        # OpenAI
        result = run_agent(client, task="...",
            model=OpenAIAdapter(model="gpt-4o"))

        # Groq
        result = run_agent(client, task="...",
            model=OpenAIAdapter(
                api_key=os.environ["GROQ_API_KEY"],
                base_url="https://api.groq.com/openai/v1",
                model="llama-3.3-70b-versatile",
            ))

        # Local Ollama
        result = run_agent(client, task="...",
            model=OpenAIAdapter(
                api_key="ollama",
                base_url="http://localhost:11434/v1",
                model="llama3.2",
            ))
    """

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "gpt-4o",
        base_url: str | None = None,
    ) -> None:
        self._api_key = api_key
        self._model = model
        self._base_url = base_url

    def complete(self, messages: list[dict]) -> str:
        try:
            from openai import OpenAI as _OpenAI
        except ImportError as exc:
            raise ImportError(
                "OpenAIAdapter requires the 'openai' package. "
                "Install it with: pip install 'hollow-sdk[openai]'"
            ) from exc

        api_key = self._api_key or os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise EnvironmentError(
                "OpenAIAdapter requires OPENAI_API_KEY in the environment "
                "or the api_key constructor argument."
            )

        kwargs: dict = dict(api_key=api_key)
        if self._base_url:
            kwargs["base_url"] = self._base_url

        openai_client = _OpenAI(**kwargs)

        # OpenAI supports system messages natively
        response = openai_client.chat.completions.create(
            model=self._model,
            max_tokens=1024,
            messages=messages,  # type: ignore[arg-type]
        )
        return response.choices[0].message.content or ""


# ── Action parsing ──────────────────────────────────────────────────────────────

def _parse_action(text: str) -> dict:
    """Extract the last JSON object with a 'type' field from the response."""
    matches = list(re.finditer(r'\{[^{}]*"type"\s*:\s*"[^"]+"[^{}]*\}', text))
    if not matches:
        raise ValueError(
            f'No ACTION JSON found in response:\n{text[:300]}'
        )
    raw = matches[-1].group(0)
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        raise ValueError(f"Could not parse ACTION JSON: {raw}") from exc


def _dict_to_action(d: dict) -> Action:
    """Convert a parsed action dict to an Action dataclass."""
    return Action(
        type=d["type"],
        element_id=d.get("elementId"),
        url=d.get("url"),
        value=d.get("value"),
        direction=d.get("direction"),
    )


# ── Agent loop ──────────────────────────────────────────────────────────────────

# Accept either a ModelAdapter or a plain callable (for inline lambdas)
_AdapterArg = Union[ModelAdapter, Callable[[list[dict]], str]]


def run_agent(
    client: HollowClient,
    task: str,
    model: Optional[_AdapterArg] = None,
    max_steps: int = _DEFAULT_MAX_STEPS,
    start_url: str = _DEFAULT_START_URL,
    on_step: Optional[Callable[[AgentStep], None]] = None,
) -> str:
    """Run a multi-step AI agent that uses Hollow as its serverless browser.

    Args:
        client:    A HollowClient instance.
        task:      The natural-language task for the agent.
        model:     Model adapter to use. Defaults to AnthropicAdapter
                   (reads ANTHROPIC_API_KEY from the environment).
                   Pass any object with a ``complete(messages) -> str`` method,
                   or a plain callable ``messages -> str``.
        max_steps: Maximum steps before raising (default 15).
        start_url: The URL to load on step 1 (default startpage.com).
        on_step:   Optional callback called after each step.

    Returns:
        The agent's final answer string.

    Raises:
        RuntimeError: If max_steps is reached without a done action.

    Example::

        from hollow import HollowClient, run_agent, OpenAIAdapter

        client = HollowClient()

        # Default Anthropic (reads ANTHROPIC_API_KEY)
        result = run_agent(client, task="What are the top 3 stories on HN?")

        # OpenAI
        result = run_agent(client, task="...",
            model=OpenAIAdapter(model="gpt-4o"))

        # Any callable
        result = run_agent(client, task="...",
            model=lambda msgs: my_model.chat(msgs))
    """
    # Resolve adapter: default to Anthropic if not provided
    if model is None:
        adapter: Callable[[list[dict]], str] = AnthropicAdapter().complete
    elif callable(model) and not isinstance(model, ModelAdapter):
        # Plain callable (e.g. a lambda)
        adapter = model  # type: ignore[assignment]
    else:
        adapter = model.complete  # type: ignore[union-attr]

    # ── Load the first page ────────────────────────────────────────────────────
    percept = client.perceive(start_url)
    session_id = percept.session_id
    gdg_map = percept.gdg_map
    confidence = percept.confidence
    tier = percept.tier

    # Accumulate user/assistant turns; system is injected fresh at each call
    history: list[dict] = []

    # ── Step loop ──────────────────────────────────────────────────────────────
    for step in range(1, max_steps + 1):
        user_content = (
            f"Task: {task}\n\nCurrent page ({start_url}):\n\n{gdg_map}"
            if step == 1
            else f"Action executed. Current page (session: {session_id}):\n\n{gdg_map}"
        )

        history.append({"role": "user", "content": user_content})

        # Build full message list with system prepended
        full_messages = [{"role": "system", "content": _SYSTEM_PROMPT}, *history]

        assistant_text = adapter(full_messages)
        history.append({"role": "assistant", "content": assistant_text})

        # Parse action from response
        try:
            action_dict = _parse_action(assistant_text)
        except ValueError as exc:
            raise ValueError(f"Step {step}: {exc}") from exc

        if action_dict.get("type") == "done":
            return action_dict.get("result", "")

        action = _dict_to_action(action_dict)

        # Fire callback
        if on_step is not None:
            on_step(AgentStep(
                step=step,
                gdg_map=gdg_map,
                action=action,
                confidence=confidence,
                tier=tier,
            ))

        # Execute action via Hollow
        if action.type == "navigate" and action.url:
            percept = client.perceive(action.url, session_id=session_id)
        else:
            percept = client.act(session_id, action)

        session_id = percept.session_id
        gdg_map = percept.gdg_map
        confidence = percept.confidence
        tier = percept.tier

    raise RuntimeError(
        f"Agent reached max steps ({max_steps}) without completing the task. "
        f"Last page: confidence={confidence} tier={tier}"
    )
