"""Strands Agent hook — first-class LLM integration for the RoohAI pipeline.

Replaces raw LLM calls with a Strands Agent that provides:
- Conversation management (sliding window) across turns within a session
- Tool use (web search, current time, calculator, etc.)
- Multi-provider support (Bedrock, OpenAI, Gemini, Ollama, Anthropic, etc.)

The hook reads the LLM provider config from the pipeline and creates the
appropriate Strands model. Each WebSocket/WebRTC session gets its own
Strands Agent instance with independent conversation memory.
"""

from __future__ import annotations

import asyncio
import copy
import logging
import os
import threading
from typing import Any

from strands import Agent
from strands.agent.conversation_manager import SlidingWindowConversationManager
from strands.event_loop._retry import ModelRetryStrategy

logger = logging.getLogger(__name__)

# ── Per-session Strands agents ───────────────────────────────────────────────
# { session_id: Agent }
# Each agent accumulates its own conversation history via Strands'
# SlidingWindowConversationManager — no manual history tracking needed.
_agents: dict[str, Agent] = {}
_lock = threading.Lock()

# Default conversation window size (number of messages to retain)
CONVERSATION_WINDOW_SIZE = 40

# Default system prompt for voice agents
_DEFAULT_SYSTEM_PROMPT = (
    "You are a friendly voice assistant. You remember everything said in this "
    "conversation and can refer back to it naturally. "
    "Keep responses concise and conversational — they will be spoken aloud. "
    "Never use markdown formatting, bullet points, or special characters. "
    "If the user asks you to recall something from earlier, do so accurately."
)

# Tools to load by default for all agents
_DEFAULT_TOOLS = [
    "http_request",
    "current_time",
    "calculator",
]


# ── Strands model factory ───────────────────────────────────────────────────

def _create_strands_model(llm_config: dict[str, Any]):
    """Create the appropriate Strands model from the pipeline's LLM config.

    Reads the LLM class/provider from the config and instantiates the
    matching Strands model provider with the correct auth.

    Args:
        llm_config: The LLM model config dict from config.yaml / agent config.
            Expected keys: class, model_id, region, api_key, etc.

    Returns:
        A Strands Model instance (BedrockModel, OpenAIModel, etc.)
    """
    llm_class = llm_config.get("class", "bedrock")
    model_id = llm_config.get("model_id", "")
    region = llm_config.get("region", "us-east-1")

    # Resolve class aliases to provider names
    provider = _resolve_provider(llm_class)
    logger.info("Resolving Strands model: class=%s → provider=%s, model_id=%s", llm_class, provider, model_id)

    if provider == "bedrock":
        return _create_bedrock_model(llm_config, model_id, region)
    elif provider == "openai":
        return _create_openai_model(llm_config, model_id)
    elif provider == "gemini":
        return _create_gemini_model(llm_config, model_id)
    elif provider == "ollama":
        return _create_ollama_model(llm_config, model_id)
    elif provider == "anthropic":
        return _create_anthropic_model(llm_config, model_id)
    else:
        # Fallback: try Bedrock
        logger.warning("Unknown LLM provider '%s', falling back to Bedrock", provider)
        return _create_bedrock_model(llm_config, model_id, region)


def _resolve_provider(llm_class: str) -> str:
    """Map LLM class names from config to Strands provider names."""
    _PROVIDER_MAP = {
        "bedrock": "bedrock",
        "openai": "openai",
        "gemini": "gemini",
        "ollama": "ollama",
        "anthropic": "anthropic",
    }
    return _PROVIDER_MAP.get(llm_class, llm_class)


def _create_bedrock_model(config: dict, model_id: str, region: str):
    """Create a Strands BedrockModel."""
    from strands.models.bedrock import BedrockModel

    if not model_id:
        model_id = "anthropic.claude-3-haiku-20240307-v1:0"

    kwargs: dict[str, Any] = {
        "model_id": model_id,
        "region_name": region,
    }

    # Handle auth modes
    auth_mode = config.get("auth_mode", "environment")
    if auth_mode == "access_keys":
        ak = config.get("aws_access_key_id")
        sk = config.get("aws_secret_access_key")
        if ak and sk:
            import boto3
            session = boto3.Session(
                aws_access_key_id=ak,
                aws_secret_access_key=sk,
                region_name=region,
            )
            kwargs["boto_session"] = session
    elif auth_mode == "api_key":
        api_key = config.get("bedrock_api_key")
        if api_key:
            os.environ["AWS_BEARER_TOKEN_BEDROCK"] = api_key

    logger.info("Creating Strands BedrockModel: model_id=%s, region=%s", model_id, region)
    return BedrockModel(**kwargs)


def _create_openai_model(config: dict, model_id: str):
    """Create a Strands OpenAIModel."""
    from strands.models.openai import OpenAIModel

    if not model_id:
        model_id = "gpt-4o"

    client_args: dict[str, Any] = {}
    api_key = config.get("api_key") or os.environ.get("OPENAI_API_KEY")
    if api_key:
        client_args["api_key"] = api_key
        logger.debug("OpenAI API key: %s...%s", api_key[:5], api_key[-4:])
    else:
        logger.warning("No OpenAI API key found in config or OPENAI_API_KEY env var")

    base_url = config.get("base_url")
    if base_url:
        client_args["base_url"] = base_url

    logger.info("Creating Strands OpenAIModel: model_id=%s, client_args_keys=%s", model_id, list(client_args.keys()))
    return OpenAIModel(client_args=client_args, model_id=model_id)


def _create_gemini_model(config: dict, model_id: str):
    """Create a Strands GeminiModel."""
    from strands.models.gemini import GeminiModel

    if not model_id:
        model_id = "gemini-2.5-flash"

    client_args: dict[str, Any] = {}
    api_key = config.get("api_key") or os.environ.get("GOOGLE_API_KEY")
    if api_key:
        client_args["api_key"] = api_key
        logger.debug("Gemini API key: %s...%s", api_key[:4], api_key[-4:])
    else:
        logger.warning("No Gemini API key found in config or GOOGLE_API_KEY env var")

    logger.info("Creating Strands GeminiModel: model_id=%s, client_args_keys=%s", model_id, list(client_args.keys()))
    try:
        model = GeminiModel(client_args=client_args, model_id=model_id)
        logger.info("Strands GeminiModel created successfully: model_id=%s", model_id)
        return model
    except Exception as exc:
        logger.error("Failed to create Strands GeminiModel: model_id=%s, error=%s", model_id, exc, exc_info=True)
        raise


def _create_ollama_model(config: dict, model_id: str):
    """Create a Strands OllamaModel."""
    from strands.models.ollama import OllamaModel

    if not model_id:
        model_id = "llama3"

    host = config.get("host", "http://localhost:11434")

    logger.info("Creating Strands OllamaModel: model_id=%s, host=%s", model_id, host)
    return OllamaModel(host=host, model_id=model_id)


def _create_anthropic_model(config: dict, model_id: str):
    """Create a Strands AnthropicModel (direct API, not Bedrock)."""
    from strands.models.anthropic import AnthropicModel

    if not model_id:
        model_id = "claude-sonnet-4-20250514"

    # Strands AnthropicModel takes api_key via client_args, not as a top-level kwarg
    api_key = config.get("api_key") or os.environ.get("ANTHROPIC_API_KEY")
    client_args: dict[str, Any] = {}
    if api_key:
        client_args["api_key"] = api_key
        logger.debug("Anthropic API key: %s...%s", api_key[:7], api_key[-4:])
    else:
        logger.warning("No Anthropic API key found in config or ANTHROPIC_API_KEY env var")

    # Strands AnthropicModel requires max_tokens in its config (unlike BedrockModel
    # which sets its own default). Voice responses should be concise.
    max_tokens = int(config.get("max_tokens", 1024))

    logger.info("Creating Strands AnthropicModel: model_id=%s, max_tokens=%d", model_id, max_tokens)
    return AnthropicModel(client_args=client_args, model_id=model_id, max_tokens=max_tokens)


# ── Agent lifecycle ──────────────────────────────────────────────────────────

def _load_tools(tool_names: list[str] | None = None) -> list:
    """Load Strands tools by name from the strands_tools package.

    Args:
        tool_names: List of tool names to load. Defaults to _DEFAULT_TOOLS.

    Returns:
        List of tool modules that can be passed to Agent(tools=[...]).
    """
    names = tool_names or _DEFAULT_TOOLS
    tools = []
    for name in names:
        try:
            import importlib
            mod = importlib.import_module(f"strands_tools.{name}")
            tools.append(mod)
            logger.debug("Loaded Strands tool: %s", name)
        except ImportError:
            logger.warning("Strands tool '%s' not found, skipping", name)
    return tools


# Saved conversation history for sessions reset by barge-in.
# Maps session_id -> list of messages to restore into the next agent.
_saved_history: dict[str, list] = {}


def _get_or_create_agent(
    session_id: str,
    model: Any,
    system_prompt: str | None = None,
    tools: list | None = None,
    window_size: int = CONVERSATION_WINDOW_SIZE,
) -> Agent:
    """Get an existing agent for a session, or create a new one.

    Each session gets its own Agent with independent conversation memory.
    The Agent's SlidingWindowConversationManager handles history trimming.
    If the session was previously reset (barge-in), restored conversation
    history is injected into the new agent so context is preserved.
    """
    with _lock:
        if session_id not in _agents:
            prompt = system_prompt or _DEFAULT_SYSTEM_PROMPT
            agent_tools = tools if tools is not None else _load_tools()

            # Voice agents need fast failure — the default Strands retry strategy
            # (6 attempts, 4s→8s→16s→32s→64s backoff) leaves the user waiting 2+ minutes
            # with no feedback. Use 2 attempts with 2s initial delay instead.
            agent = Agent(
                model=model,
                system_prompt=prompt,
                tools=agent_tools,
                conversation_manager=SlidingWindowConversationManager(
                    window_size=window_size,
                ),
                retry_strategy=ModelRetryStrategy(
                    max_attempts=2,
                    initial_delay=2,
                    max_delay=4,
                ),
                # Suppress Strands' default print callback — we handle output ourselves
                callback_handler=None,
            )

            # Restore conversation history saved from a barge-in reset
            saved = _saved_history.pop(session_id, None)
            if saved:
                agent.messages.extend(saved)
                logger.info(
                    "Restored %d messages into new Strands agent for session %s",
                    len(saved), session_id,
                )

            _agents[session_id] = agent
            logger.info("Created Strands agent for session %s", session_id)
        return _agents[session_id]


def reset_session(session_id: str) -> None:
    """Discard the in-flight Strands agent for a session (called on barge-in).

    Preserves conversation history so the next agent for this session can
    continue the conversation without losing context.
    """
    with _lock:
        agent = _agents.pop(session_id, None)
        if agent:
            # Save conversation history — deep copy so the orphaned
            # thread can't mutate message dicts after we hand them to the new agent.
            try:
                _saved_history[session_id] = copy.deepcopy(agent.messages)
                logger.info(
                    "Reset Strands agent for session %s (barge-in), saved %d messages",
                    session_id, len(agent.messages),
                )
            except Exception:
                logger.warning("Failed to save history for session %s", session_id, exc_info=True)


def clear_session(session_id: str) -> None:
    """Remove the Strands agent for a session (called on disconnect).

    This frees conversation memory for the session.
    """
    with _lock:
        agent = _agents.pop(session_id, None)
        _saved_history.pop(session_id, None)
        if agent:
            logger.info("Cleared Strands agent for session %s", session_id)


def clear_all_sessions() -> None:
    """Remove all Strands agents (called on server shutdown)."""
    with _lock:
        count = len(_agents)
        _agents.clear()
        _saved_history.clear()
        if count:
            logger.info("Cleared %d Strands agent(s)", count)


# ── Module-level state set by configure() ────────────────────────────────────

_strands_model = None
_strands_system_prompt: str | None = None
_strands_tools: list | None = None
_strands_window_size: int = CONVERSATION_WINDOW_SIZE


def configure(
    llm_config: dict[str, Any],
    system_prompt: str | None = None,
    tool_names: list[str] | None = None,
    window_size: int = CONVERSATION_WINDOW_SIZE,
) -> None:
    """Configure the Strands hook with model and tool settings.

    Called once during agent activation in _activate_agent(). All subsequent
    calls to strands_hook/strands_stream_hook use these settings.

    Args:
        llm_config: LLM model config dict from config.yaml (class, model_id, region, etc.)
        system_prompt: Optional system prompt override.
        tool_names: Optional list of Strands tool names to load.
        window_size: Conversation window size (default 40).
    """
    global _strands_model, _strands_system_prompt, _strands_tools, _strands_window_size

    logger.info("Configuring Strands hook: llm_class=%s, model_id=%s, system_prompt=%s, tools=%s, window=%d",
                llm_config.get("class"), llm_config.get("model_id"), bool(system_prompt), tool_names, window_size)

    _strands_model = _create_strands_model(llm_config)
    _strands_system_prompt = system_prompt
    _strands_tools = _load_tools(tool_names)
    _strands_window_size = window_size

    # Initialize Strands OTEL exporter if OTEL env vars are set
    # (e.g. mapped from Langfuse secrets in server/app.py)
    if os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"):
        try:
            from strands.telemetry.config import StrandsTelemetry
            StrandsTelemetry().setup_otlp_exporter()
            logger.info("Strands OTEL exporter initialized (endpoint=%s)",
                        os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"])
        except Exception:
            logger.warning("Failed to initialize Strands OTEL exporter", exc_info=True)

    # Clear any existing sessions since the model changed
    clear_all_sessions()

    logger.info("Strands hook configured: provider=%s, tools=%s, window=%d",
                type(_strands_model).__name__,
                [t.__name__ if hasattr(t, '__name__') else str(t) for t in (_strands_tools or [])],
                _strands_window_size)


# ── Hook functions (wired into pipeline via set_llm_hook) ────────────────────

def _strip_markdown(text: str) -> str:
    """Strip markdown artifacts that sound odd when spoken via TTS."""
    for ch in ("**", "*", "#", "`", "---", "- "):
        text = text.replace(ch, "")
    return text.strip()


async def strands_hook(
    message: str,
    *,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    session_id: str | None = None,
) -> str:
    """RoohAI LLM hook (batch) — routes through Strands Agent with session memory.

    This replaces the raw LLM call in pipeline.chat(). The Strands Agent
    manages its own conversation history via SlidingWindowConversationManager,
    so we don't need to pass history manually.

    Args:
        message: User's transcribed speech.
        history: Ignored — Strands manages conversation internally.
        system_prompt: Override system prompt (or use configured default).
        session_id: WebSocket/WebRTC session ID for per-session memory.

    Returns:
        Agent's text response (markdown stripped for TTS).
    """
    if _strands_model is None:
        raise RuntimeError("Strands hook not configured. Call configure() first.")

    sid = session_id or "default"
    prompt = system_prompt or _strands_system_prompt

    logger.info("Strands hook invoked: session=%s, message_len=%d, provider=%s",
                sid, len(message), type(_strands_model).__name__)
    logger.debug("Strands hook message: session=%s, message=%r", sid, message[:200])

    agent = _get_or_create_agent(
        sid,
        model=_strands_model,
        system_prompt=prompt,
        tools=_strands_tools,
        window_size=_strands_window_size,
    )

    # Run synchronous Strands agent in a thread to avoid blocking the event loop.
    # If barge-in reset the session while the old call was in-flight, the agent
    # may already be gone or a ConcurrencyException may fire — retry with a fresh agent.
    try:
        result = await asyncio.to_thread(agent, message)
    except Exception as exc:
        if "Concurrent" in str(type(exc).__name__) or "concurrent" in str(exc).lower():
            logger.warning("Strands ConcurrencyException for session %s — resetting agent", sid)
            reset_session(sid)
            agent = _get_or_create_agent(
                sid,
                model=_strands_model,
                system_prompt=prompt,
                tools=_strands_tools,
                window_size=_strands_window_size,
            )
            result = await asyncio.to_thread(agent, message)
        else:
            logger.error("Strands hook failed: session=%s, provider=%s, error=%s",
                         sid, type(_strands_model).__name__, exc, exc_info=True)
            raise
    text = str(result)
    logger.info("Strands hook response: session=%s, response_len=%d", sid, len(text))
    logger.debug("Strands hook response text: session=%s, text=%r", sid, text[:200])

    return _strip_markdown(text)
