"""Method registry: get_method(name, **params) returns a fresh agent.

Methods with heavy optional dependencies are imported lazily.
"""
from __future__ import annotations

from .base import BaseAgent, MemoryBuildResult, MemoryContext
from .long_context import LongContextAgent
from .bm25_rag import BM25RAGAgent
from .dense_rag import DenseRAGAgent

__all__ = ["BaseAgent", "MemoryBuildResult", "MemoryContext", "LongContextAgent", "BM25RAGAgent",
           "DenseRAGAgent", "SUPPORTED_METHODS", "DEFAULT_METHODS", "INTERNAL_LLM_METHODS",
           "EMBEDDING_DIM_METHODS", "EMBEDDING_MODEL_METHODS", "EMBEDDING_BATCH_METHODS",
           "EMBEDDING_TOKENIZER_METHODS", "TOP_K_PARAMS", "embedding_params", "get_method"]

SUPPORTED_METHODS = ("long_context", "bm25_rag", "dense_rag",
                     "hipporag", "remem", "memos", "letta", "mirix",
                     "amem", "lightmem", "mem0")
DEFAULT_METHODS = ("long_context", "bm25_rag", "dense_rag")
INTERNAL_LLM_METHODS = ("memos", "letta", "mirix", "amem", "lightmem",
                        "hipporag", "remem", "mem0")
EMBEDDING_DIM_METHODS = ("hipporag", "remem", "letta", "mirix", "lightmem", "mem0", "memos")
EMBEDDING_MODEL_METHODS = EMBEDDING_DIM_METHODS + ("dense_rag", "amem")
# letta and lightmem embed one text per request inside their engines
EMBEDDING_BATCH_METHODS = ("dense_rag", "hipporag", "remem", "mirix", "amem", "mem0", "memos")
# methods that count embedding inputs with embedding.tokenizer (truncation, usage estimate)
EMBEDDING_TOKENIZER_METHODS = ("amem", "mem0", "mirix")
# the parameter that holds the number of retrieved items (top_k / --top-k)
TOP_K_PARAMS = {"bm25_rag": "top_k", "dense_rag": "top_k", "hipporag": "qa_top_k", "remem": "qa_top_k",
                "letta": "retrieve_num", "mirix": "retrieve_num", "amem": "retrieve_num",
                "lightmem": "retrieve_num", "mem0": "retrieve_num", "memos": "retrieve_num"}


def embedding_params(name: str, embedding: dict) -> dict:
    """The embedding settings of the backbone config that a method takes as parameters."""
    out = {}
    if name in EMBEDDING_MODEL_METHODS:
        out["embedding_model"] = embedding["model"]
    if name in EMBEDDING_DIM_METHODS:
        out["embedding_dim"] = embedding["dim"]
    if name in EMBEDDING_BATCH_METHODS:
        out["embedding_batch_size"] = embedding.get("batch_size")
    return out


def get_method(name: str, **params) -> BaseAgent:
    if name == "long_context":
        return LongContextAgent(**params)
    if name == "bm25_rag":
        return BM25RAGAgent(**params)
    if name == "dense_rag":
        from ..llm_client import make_embedder
        embed_model = params.pop("embedding_model", None)
        batch_size = params.pop("embedding_batch_size", None)
        return DenseRAGAgent(embedder=make_embedder(model=embed_model, batch_size=batch_size), **params)
    if name == "hipporag":
        from .hipporag_agent import HippoRAGAgent
        return HippoRAGAgent(**params)
    if name == "remem":
        from .remem_agent import ReMemAgent
        return ReMemAgent(**params)
    if name == "memos":
        from .memos_agent import MemOSAgent
        return MemOSAgent(**params)
    if name == "letta":
        from .letta_agent import LettaAgent
        return LettaAgent(**params)
    if name == "mirix":
        from .mirix_agent import MIRIXAgent
        return MIRIXAgent(**params)
    if name == "amem":
        from .amem_agent import AMemAgent
        return AMemAgent(**params)
    if name == "lightmem":
        from .lightmem_agent import LightMemAgent
        return LightMemAgent(**params)
    if name == "mem0":
        from .mem0_agent import Mem0Agent
        return Mem0Agent(**params)
    raise ValueError(f"unknown method '{name}'; supported: {', '.join(SUPPORTED_METHODS)}")