"""
novyx_agent.agent — The core Agent class.

5 lines to a working agent with persistent memory:

    from novyx_agent import Agent, tool

    agent = Agent(name="Aria", api_key="nram_...")

    @agent.tool
    def search(query: str) -> str:
        return "result"

    print(agent.run("What did we discuss yesterday?"))
"""

from __future__ import annotations

import json
import logging
from typing import Any, Callable, Dict, List, Optional, Sequence, Union

from novyx import Novyx

from .providers import LLMProvider, get_provider
from .tools import tool as tool_decorator
from .types import (
    AgentConfig,
    MemoryConfig,
    Message,
    Role,
    RunResult,
    ToolCall,
    ToolDef,
)

logger = logging.getLogger("novyx_agent")

# Score thresholds for tiered memory formatting.
# Memories above STRONG are presented as core context.
# Memories between MODERATE and STRONG are supplementary.
# Memories below MODERATE are background — the LLM can ignore them.
_TIER_STRONG = 0.65
_TIER_MODERATE = 0.45


def _format_tiered_memories(memories: Sequence[Any]) -> str:
    """Format recalled memories with signal-strength tiers.

    Instead of a flat list (which LLMs overfit to equally), we label each
    memory as [strong], [moderate], or [background] based on its composite
    score.  This gives the model explicit permission to deprioritize low-
    signal memories — addressing the "one question from 2 months ago keeps
    coming up" problem identified by Karpathy and others.

    Only the top 2 background memories are included to avoid noise.
    """
    strong, moderate, background = [], [], []
    for m in memories:
        score = getattr(m, "score", 0.0) or 0.0
        line = m.observation
        if score >= _TIER_STRONG:
            strong.append(f"- [strong] {line}")
        elif score >= _TIER_MODERATE:
            moderate.append(f"- [moderate] {line}")
        else:
            background.append(f"- [background] {line}")

    # Cap background memories to avoid noise
    background = background[:2]

    parts = []
    if strong:
        parts.append("Core context (directly relevant):\n" + "\n".join(strong))
    if moderate:
        parts.append("Supplementary context:\n" + "\n".join(moderate))
    if background:
        parts.append("Background (may not be relevant — use only if asked):\n" + "\n".join(background))

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


class Agent:
    """
    AI agent with Novyx memory baked in.

    Every run automatically:
    1. Recalls relevant memories before the LLM sees the prompt
    2. Executes tool calls in a loop until the model responds
    3. Remembers key outputs for future recall
    4. Records everything in the Novyx audit trail

    Plus rollback, eval, and knowledge graph — all one method call away.
    """

    def __init__(
        self,
        *,
        name: str = "Agent",
        instructions: str = "You are a helpful assistant.",
        api_key: str,
        model: str = "gpt-4o-mini",
        provider: str = "openai",
        temperature: float = 0.7,
        max_tokens: int = 4096,
        max_tool_rounds: int = 10,
        memory: Optional[Union[MemoryConfig, Dict[str, Any]]] = None,
        tools: Optional[Sequence[Callable]] = None,
        novyx_url: Optional[str] = None,
        llm_api_key: Optional[str] = None,
        llm_base_url: Optional[str] = None,
    ) -> None:
        """
        Create an agent.

        Args:
            name: Agent name (used as agent_id in Novyx)
            instructions: System prompt
            api_key: Novyx API key (nram_...)
            model: LLM model name
            provider: "openai", "anthropic", or "litellm"
            temperature: LLM temperature
            max_tokens: Max tokens per LLM call
            max_tool_rounds: Max tool-call loops per run
            memory: Memory configuration (auto-recall, auto-remember, etc.)
            tools: List of @tool-decorated functions
            novyx_url: Custom Novyx API URL
            llm_api_key: LLM provider API key (uses env var if not set)
            llm_base_url: Custom LLM base URL (OpenAI-compatible)
        """
        if isinstance(memory, dict):
            memory = MemoryConfig(**memory)

        self._config = AgentConfig(
            name=name,
            instructions=instructions,
            model=model,
            provider=provider,
            temperature=temperature,
            max_tokens=max_tokens,
            max_tool_rounds=max_tool_rounds,
            memory=memory or MemoryConfig(),
        )

        # Novyx memory client
        nx_kwargs: Dict[str, Any] = {"api_key": api_key, "agent_id": name, "source": "novyx-agent"}
        if novyx_url:
            nx_kwargs["api_url"] = novyx_url
        self._nx = Novyx(**nx_kwargs)

        # LLM provider
        provider_kwargs: Dict[str, Any] = {}
        if llm_api_key:
            provider_kwargs["api_key"] = llm_api_key
        if llm_base_url and provider == "openai":
            provider_kwargs["base_url"] = llm_base_url
        self._llm = get_provider(provider, **provider_kwargs)

        # Tools registry
        self._tools: Dict[str, ToolDef] = {}
        if tools:
            for t in tools:
                self._register_tool(t)

        # Conversation history (within a session)
        self._history: List[Message] = []

    # ------------------------------------------------------------------
    # Tool registration
    # ------------------------------------------------------------------

    def tool(self, fn: Optional[Callable] = None, **kwargs: Any) -> Any:
        """
        Register a tool via decorator.

            @agent.tool
            def search(query: str) -> str:
                '''Search the web.'''
                return "..."
        """

        def _wrap(f: Callable) -> Callable:
            if not hasattr(f, "_tool_def"):
                f = tool_decorator(f, **kwargs)
            self._register_tool(f)
            return f

        if fn is not None:
            return _wrap(fn)
        return _wrap

    def _register_tool(self, fn: Callable) -> None:
        """Register a decorated function as a tool."""
        if not hasattr(fn, "_tool_def"):
            raise ValueError(
                f"Function '{fn.__name__}' is not decorated with @tool. "
                "Use @tool or @agent.tool."
            )
        td: ToolDef = fn._tool_def  # type: ignore[attr-defined]
        self._tools[td.name] = td

    # ------------------------------------------------------------------
    # Core run loop
    # ------------------------------------------------------------------

    def run(self, prompt: str, *, context: Optional[str] = None) -> RunResult:
        """
        Run the agent on a prompt.

        1. Auto-recall relevant memories
        2. Build messages (system + memory context + history + user prompt)
        3. LLM loop: call model → execute tools → repeat until text response
        4. Auto-remember the exchange
        5. Return RunResult

        Args:
            prompt: User message
            context: Optional extra context to inject

        Returns:
            RunResult with output, tool_calls, memory stats
        """
        mem_cfg = self._config.memory
        tool_calls_log: List[ToolCall] = []
        memories_recalled = 0
        memories_stored = 0
        total_usage: Dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0}

        # --- 1. Auto-recall ---
        memory_context = ""
        if mem_cfg.auto_recall:
            try:
                recall_kwargs: Dict[str, Any] = {
                    "limit": mem_cfg.recall_limit,
                    "threshold": mem_cfg.recall_threshold,
                }
                if mem_cfg.recall_tags:
                    recall_kwargs["tags"] = mem_cfg.recall_tags
                results = self._nx.recall(prompt, **recall_kwargs)
                if results.memories:
                    memories_recalled = len(results.memories)
                    memory_context = _format_tiered_memories(results.memories)
                    logger.debug("Recalled %d memories", memories_recalled)
            except Exception as e:
                logger.warning("Memory recall failed: %s", e)

        # --- 2. Build messages ---
        messages: List[Message] = []

        # System prompt
        system_parts = [self._config.instructions]
        if memory_context:
            system_parts.append(memory_context)
        if context:
            system_parts.append(f"Additional context:\n{context}")
        messages.append(Message(role=Role.SYSTEM, content="\n\n".join(system_parts)))

        # Conversation history
        messages.extend(self._history)

        # User prompt
        user_msg = Message(role=Role.USER, content=prompt)
        messages.append(user_msg)
        self._history.append(user_msg)

        # --- 3. LLM tool loop ---
        tool_defs = list(self._tools.values())
        final_content = ""

        for _round in range(self._config.max_tool_rounds):
            content, calls, usage = self._llm.chat(
                messages,
                tool_defs,
                model=self._config.model,
                temperature=self._config.temperature,
                max_tokens=self._config.max_tokens,
            )
            total_usage["prompt_tokens"] += usage.get("prompt_tokens", 0)
            total_usage["completion_tokens"] += usage.get("completion_tokens", 0)

            if not calls:
                final_content = content
                break

            # Record assistant message with tool calls
            assistant_msg = Message(
                role=Role.ASSISTANT,
                content=content,
                tool_calls=[
                    {"id": c["id"], "name": c["name"], "arguments": c["arguments"]}
                    for c in calls
                ],
            )
            messages.append(assistant_msg)

            # Execute each tool call
            for call in calls:
                tool_name = call["name"]
                tool_args = call["arguments"]
                td = self._tools.get(tool_name)

                if td and td.fn:
                    try:
                        result = td.fn(**tool_args)
                        result_str = str(result) if not isinstance(result, str) else result
                    except Exception as e:
                        result_str = f"Error: {e}"
                        logger.error("Tool '%s' raised: %s", tool_name, e)
                else:
                    result_str = f"Error: Unknown tool '{tool_name}'"

                tool_calls_log.append(
                    ToolCall(tool=tool_name, args=tool_args, result=result_str)
                )

                messages.append(
                    Message(
                        role=Role.TOOL,
                        content=result_str,
                        tool_call_id=call["id"],
                        name=tool_name,
                    )
                )
        else:
            # Hit max rounds — use last content
            final_content = final_content or "I reached the maximum number of tool rounds."

        # Record assistant response in history
        self._history.append(Message(role=Role.ASSISTANT, content=final_content))

        # --- 4. Auto-remember ---
        if mem_cfg.auto_remember and final_content:
            try:
                remember_kwargs: Dict[str, Any] = {
                    "importance": mem_cfg.remember_importance,
                }
                if mem_cfg.remember_tags:
                    remember_kwargs["tags"] = mem_cfg.remember_tags
                observation = f"[{self._config.name}] Q: {prompt[:200]} → A: {final_content[:500]}"
                self._nx.remember(observation, **remember_kwargs)
                memories_stored = 1
                logger.debug("Stored memory for this exchange")
            except Exception as e:
                logger.warning("Memory store failed: %s", e)

        return RunResult(
            output=final_content,
            tool_calls=tool_calls_log,
            memories_recalled=memories_recalled,
            memories_stored=memories_stored,
            model=self._config.model,
            usage=total_usage,
        )

    # ------------------------------------------------------------------
    # Memory operations — direct passthrough to Novyx
    # ------------------------------------------------------------------

    def remember(self, observation: str, **kwargs: Any) -> Any:
        """Store a memory explicitly."""
        return self._nx.remember(observation, **kwargs)

    def recall(self, query: str, **kwargs: Any) -> Any:
        """Search memories semantically."""
        return self._nx.recall(query, **kwargs)

    def forget(self, memory_id: str) -> bool:
        """Delete a specific memory."""
        return self._nx.forget(memory_id)

    def memories(self, **kwargs: Any) -> Any:
        """List all memories."""
        return self._nx.memories(**kwargs)

    # ------------------------------------------------------------------
    # Rollback
    # ------------------------------------------------------------------

    def rollback(self, target: str) -> Any:
        """Roll back memory state to a point in time (e.g., '2 hours ago')."""
        return self._nx.rollback(target)

    def rollback_preview(self, target: str) -> Any:
        """Preview what a rollback would do without executing it."""
        return self._nx.rollback_preview(target)

    # ------------------------------------------------------------------
    # Audit
    # ------------------------------------------------------------------

    def audit(self, **kwargs: Any) -> Any:
        """Get cryptographic audit trail."""
        return self._nx.audit(**kwargs)

    # ------------------------------------------------------------------
    # Eval
    # ------------------------------------------------------------------

    def eval(self, **kwargs: Any) -> Any:
        """Run memory health evaluation (recall, drift, conflicts, staleness)."""
        return self._nx.eval_run(**kwargs)

    def eval_gate(self, min_score: float) -> Any:
        """Gate: pass/fail based on minimum eval score."""
        return self._nx.eval_gate(min_score)

    # ------------------------------------------------------------------
    # Knowledge graph
    # ------------------------------------------------------------------

    def add_triple(self, subject: str, predicate: str, obj: str, **kwargs: Any) -> Any:
        """Add a knowledge graph triple."""
        return self._nx.add_triple(subject, predicate, obj, **kwargs)

    def query_triples(self, **kwargs: Any) -> Any:
        """Query the knowledge graph."""
        return self._nx.query_triples(**kwargs)

    # ------------------------------------------------------------------
    # Session management
    # ------------------------------------------------------------------

    def clear_history(self) -> None:
        """Clear conversation history (does NOT affect Novyx memories)."""
        self._history.clear()

    @property
    def history(self) -> List[Message]:
        """Current conversation history."""
        return list(self._history)

    @property
    def novyx(self) -> Novyx:
        """Direct access to the underlying Novyx client for advanced operations."""
        return self._nx

    @property
    def config(self) -> AgentConfig:
        """Agent configuration."""
        return self._config

    # ------------------------------------------------------------------
    # Repr
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"Agent(name={self._config.name!r}, model={self._config.model!r}, "
            f"tools={list(self._tools.keys())}, "
            f"memory={'on' if self._config.memory.auto_recall else 'off'})"
        )
