"""
MuBit Learn — Closed-loop agentic memory with one-liner setup.

Usage::

    import mubit.learn

    mubit.learn.init(api_key="mbt_...", agent_id="my-agent")

    # All OpenAI/Anthropic/LiteLLM/Google GenAI calls now auto-ingest
    # traces AND auto-inject relevant lessons before each call.

    @mubit.learn.run(agent_id="planner")
    def plan(task):
        client = openai.OpenAI()
        return client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": task}],
        ).choices[0].message.content

    # For raw HTTP or unsupported LLM libraries, use manual helpers:
    context = mubit.learn.get_lessons("user question")
    # ... make your LLM call ...
    mubit.learn.capture(
        messages=[{"role": "user", "content": "question"}],
        response="answer",
        model="gpt-4o",
    )
"""

import atexit
import functools
import inspect
import json
import logging
import uuid
from typing import Any, Callable, Dict, List, Optional, TypeVar

from mubit.auto._context import Span, set_current_span, get_current_span
from mubit.auto._items import build_items
from mubit.auto._worker import IngestWorker
from mubit.learn._client import LearnClient
from mubit.learn._config import LearnConfig
from mubit.learn._lesson_cache import LessonCache
from mubit.learn._run_manager import RunManager

logger = logging.getLogger("mubit.learn")

F = TypeVar("F", bound=Callable[..., Any])

# Module-level state
_active_config: Optional[LearnConfig] = None
_learn_client: Optional[LearnClient] = None
_lesson_cache: Optional[LessonCache] = None
_run_manager: Optional[RunManager] = None
_original_refs: dict = {}


def init(
    api_key: Optional[str] = None,
    endpoint: Optional[str] = None,
    agent_id: str = "auto",
    user_id: str = "",
    session_id: Optional[str] = None,
    *,
    inject_lessons: bool = True,
    injection_position: str = "system",
    max_token_budget: int = 2048,
    entry_types: Optional[List[str]] = None,
    context_sections: Optional[List[str]] = None,
    capture: str = "all",
    min_length: int = 0,
    lane: str = "",
    auto_extract: bool = True,
    extraction_mode: str = "heuristic",
    auto_reflect: bool = True,
    reflect_after_n_calls: Optional[int] = None,
    cache_ttl_seconds: float = 30.0,
    cache_max_entries: int = 100,
    fail_open: bool = True,
) -> "RunManager":
    """Initialize mubit.learn — one-liner closed-loop agentic memory.

    After calling init(), all OpenAI/Anthropic/LiteLLM client instances
    will automatically:
    - Retrieve relevant lessons and inject them before each LLM call
    - Ingest all inputs and outputs in the background
    - Trigger reflection when the run ends

    Args:
        api_key: MuBit API key (falls back to MUBIT_API_KEY env var).
        endpoint: MuBit endpoint (falls back to MUBIT_ENDPOINT env var).
        agent_id: Default agent ID for all captured traces.
        user_id: Default user ID.
        session_id: Session/run ID (None = auto-generate UUID).
        inject_lessons: Whether to inject lessons before LLM calls.
        injection_position: Where to inject ("system", "prepend", "last_system").
        max_token_budget: Token budget for context retrieval.
        entry_types: Memory entry types to retrieve (None = server default).
        context_sections: Context sections to include (None = server default).
        capture: Ingestion capture mode ("all", "input_only", "output_only").
        min_length: Minimum text length to capture.
        auto_reflect: Trigger reflection when run ends.
        reflect_after_n_calls: Trigger reflection every N LLM calls (None = off).
        cache_ttl_seconds: Lesson cache TTL in seconds.
        cache_max_entries: Max lesson cache entries.
        fail_open: If True, proceed without lessons on retrieval failure.

    Returns:
        The RunManager instance for this session.
    """
    global _active_config, _learn_client, _lesson_cache, _run_manager

    config = LearnConfig(
        api_key=api_key,
        endpoint=endpoint,
        agent_id=agent_id,
        user_id=user_id,
        session_id=session_id,
        inject_lessons=inject_lessons,
        injection_position=injection_position,
        max_token_budget=max_token_budget,
        entry_types=entry_types,
        context_sections=context_sections,
        capture=capture,
        min_length=min_length,
        lane=lane,
        auto_extract=auto_extract,
        extraction_mode=extraction_mode,
        auto_reflect=auto_reflect,
        reflect_after_n_calls=reflect_after_n_calls,
        cache_ttl_seconds=cache_ttl_seconds,
        cache_max_entries=cache_max_entries,
        fail_open=fail_open,
    ).resolve()

    _active_config = config
    _learn_client = LearnClient(config.api_key, config.endpoint)
    _lesson_cache = LessonCache(config.cache_ttl_seconds, config.cache_max_entries)
    _run_manager = RunManager(config, _learn_client)

    # Register shutdown hook
    atexit.register(_shutdown)

    # Monkey-patch LLM libraries
    _patch_openai(config, _lesson_cache, _learn_client, _run_manager)
    _patch_anthropic(config, _lesson_cache, _learn_client, _run_manager)
    _patch_litellm(config, _lesson_cache, _learn_client, _run_manager)
    _patch_google_genai(config, _lesson_cache, _learn_client, _run_manager)

    return _run_manager


def run(
    agent_id: Optional[str] = None,
    session_id: Optional[str] = None,
    *,
    auto_reflect: Optional[bool] = None,
    inject_lessons: Optional[bool] = None,
    max_token_budget: Optional[int] = None,
) -> Callable[[F], F]:
    """Decorator for a scoped learn run.

    Creates a span with a managed session_id and triggers reflection on exit.
    Requires init() to have been called first.

    Usage::

        @mubit.learn.run(agent_id="planner")
        def my_func():
            ...  # LLM calls here auto-enriched
    """

    def decorator(fn: F) -> F:
        is_async = inspect.iscoroutinefunction(fn)

        @functools.wraps(fn)
        def sync_wrapper(*args, **kwargs):
            rm = _enter_run(agent_id, session_id, auto_reflect, inject_lessons, max_token_budget)
            try:
                return fn(*args, **kwargs)
            finally:
                rm.end()

        @functools.wraps(fn)
        async def async_wrapper(*args, **kwargs):
            rm = _enter_run(agent_id, session_id, auto_reflect, inject_lessons, max_token_budget)
            try:
                return await fn(*args, **kwargs)
            finally:
                rm.end()

        return async_wrapper if is_async else sync_wrapper

    return decorator


def start_run(
    agent_id: Optional[str] = None,
    session_id: Optional[str] = None,
    *,
    auto_reflect: Optional[bool] = None,
) -> "RunManager":
    """Start an explicit learn run. Call run.end() when done.

    Requires init() to have been called first.

    Usage::

        run = mubit.learn.start_run(agent_id="researcher")
        # ... do work ...
        run.end()  # triggers flush + reflect
    """
    return _enter_run(agent_id, session_id, auto_reflect)


def uninstrument() -> None:
    """Restore original LLM library behavior and stop the learn module."""
    global _active_config, _learn_client, _lesson_cache, _run_manager

    # Restore OpenAI
    try:
        import openai
        if "openai.OpenAI.__init__" in _original_refs:
            openai.OpenAI.__init__ = _original_refs.pop("openai.OpenAI.__init__")
            if hasattr(openai.OpenAI, "_mubit_learn_instrumented"):
                del openai.OpenAI._mubit_learn_instrumented
        if "openai.AsyncOpenAI.__init__" in _original_refs:
            openai.AsyncOpenAI.__init__ = _original_refs.pop("openai.AsyncOpenAI.__init__")
            if hasattr(openai.AsyncOpenAI, "_mubit_learn_instrumented"):
                del openai.AsyncOpenAI._mubit_learn_instrumented
    except ImportError:
        pass

    # Restore Anthropic
    try:
        import anthropic
        if "anthropic.Anthropic.__init__" in _original_refs:
            anthropic.Anthropic.__init__ = _original_refs.pop("anthropic.Anthropic.__init__")
            if hasattr(anthropic.Anthropic, "_mubit_learn_instrumented"):
                del anthropic.Anthropic._mubit_learn_instrumented
        if "anthropic.AsyncAnthropic.__init__" in _original_refs:
            anthropic.AsyncAnthropic.__init__ = _original_refs.pop("anthropic.AsyncAnthropic.__init__")
            if hasattr(anthropic.AsyncAnthropic, "_mubit_learn_instrumented"):
                del anthropic.AsyncAnthropic._mubit_learn_instrumented
    except ImportError:
        pass

    # Restore Google GenAI
    try:
        import google.genai
        if "google.genai.Client.__init__" in _original_refs:
            google.genai.Client.__init__ = _original_refs.pop("google.genai.Client.__init__")
            if hasattr(google.genai.Client, "_mubit_learn_instrumented"):
                del google.genai.Client._mubit_learn_instrumented
    except ImportError:
        pass

    _original_refs.clear()
    _active_config = None
    _learn_client = None
    _lesson_cache = None
    _run_manager = None


def get_lessons(
    query: str,
    *,
    session_id: Optional[str] = None,
    max_token_budget: Optional[int] = None,
    entry_types: Optional[List[str]] = None,
    sections: Optional[List[str]] = None,
) -> str:
    """Fetch relevant lessons as a context string. Works without any LLM library.

    This is the manual equivalent of the automatic lesson injection. Use it
    when making raw HTTP calls to LLM APIs or using unsupported libraries.

    Requires init() to have been called first.

    Args:
        query: The user question or topic to retrieve lessons for.
        session_id: Override session ID (defaults to current run's session).
        max_token_budget: Override token budget for context retrieval.
        entry_types: Override entry types to retrieve.
        sections: Override context sections to include.

    Returns:
        Pre-formatted context string, or empty string on failure.

    Raises:
        RuntimeError: If init() has not been called.
    """
    if _learn_client is None or _active_config is None or _run_manager is None:
        raise RuntimeError("mubit.learn.init() must be called before get_lessons()")

    sid = session_id or _run_manager.session_id
    budget = max_token_budget or _active_config.max_token_budget
    etypes = entry_types or _active_config.entry_types
    sects = sections or _active_config.context_sections

    # Check cache first
    if _lesson_cache is not None:
        cached = _lesson_cache.get(sid, query)
        if cached is not None:
            return cached

    try:
        context_block = _learn_client.get_context(
            session_id=sid,
            query=query,
            max_token_budget=budget,
            entry_types=etypes,
            sections=sects,
        )
        if _lesson_cache is not None:
            _lesson_cache.set(sid, query, context_block)
        return context_block
    except Exception as e:
        if _active_config.fail_open:
            logger.debug("mubit.learn.get_lessons() failed: %s", e)
            return ""
        raise


def capture(
    messages: List[Dict[str, Any]],
    response: str,
    *,
    model: str = "",
    session_id: Optional[str] = None,
    agent_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> None:
    """Ingest an LLM interaction. Works without any LLM library.

    This is the manual equivalent of the automatic trace capture. Use it
    when making raw HTTP calls to LLM APIs or using unsupported libraries.

    Requires init() to have been called first.

    Args:
        messages: OpenAI-format messages list (dicts with 'role' and 'content').
        response: The assistant's response text.
        model: The LLM model name (for metadata).
        session_id: Override session ID (defaults to current run's session).
        agent_id: Override agent ID.
        user_id: Override user ID.

    Raises:
        RuntimeError: If init() has not been called.
    """
    if _active_config is None or _run_manager is None:
        raise RuntimeError("mubit.learn.init() must be called before capture()")

    sid = session_id or _run_manager.session_id
    aid = agent_id or _active_config.agent_id
    uid = user_id or _active_config.user_id

    items = build_items(
        messages=messages,
        assistant_text=response,
        model=model or "unknown",
        latency_ms=0,
        capture=_active_config.capture,
        min_length=_active_config.min_length,
        user_id=uid,
    )

    if not items:
        return

    # Use a one-off worker or find an existing one
    worker = IngestWorker(
        api_key=_active_config.api_key,
        endpoint=_active_config.endpoint,
    )
    worker.start()
    worker.enqueue(sid, aid, items)

    # Also run auto-extraction if enabled
    if _active_config.auto_extract and response:
        from mubit.learn._extraction import extract_structured_items
        extracted = extract_structured_items(
            messages=messages,
            assistant_text=response,
            model=model or "unknown",
            user_id=uid,
        )
        if extracted:
            worker.enqueue(sid, aid, extracted)

    # Increment call count for reflection tracking
    _run_manager.increment()


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _enter_run(
    agent_id: Optional[str],
    session_id: Optional[str],
    auto_reflect: Optional[bool] = None,
    inject_lessons: Optional[bool] = None,
    max_token_budget: Optional[int] = None,
) -> RunManager:
    """Create a scoped RunManager and set up a Span."""
    if _active_config is None or _learn_client is None:
        raise RuntimeError("mubit.learn.init() must be called before starting a run")

    # Build a per-run config with overrides
    cfg = LearnConfig(
        api_key=_active_config.api_key,
        endpoint=_active_config.endpoint,
        agent_id=agent_id or _active_config.agent_id,
        user_id=_active_config.user_id,
        session_id=session_id,
        inject_lessons=inject_lessons if inject_lessons is not None else _active_config.inject_lessons,
        injection_position=_active_config.injection_position,
        max_token_budget=max_token_budget if max_token_budget is not None else _active_config.max_token_budget,
        entry_types=_active_config.entry_types,
        context_sections=_active_config.context_sections,
        capture=_active_config.capture,
        min_length=_active_config.min_length,
        auto_reflect=auto_reflect if auto_reflect is not None else _active_config.auto_reflect,
        reflect_after_n_calls=_active_config.reflect_after_n_calls,
        cache_ttl_seconds=_active_config.cache_ttl_seconds,
        cache_max_entries=_active_config.cache_max_entries,
        fail_open=_active_config.fail_open,
    )

    rm = RunManager(cfg, _learn_client)

    # Push a span so auto-capture wrappers pick up the session_id
    parent = get_current_span()
    span = Span(
        parent_id=parent.trace_id if parent else None,
        name=f"learn-run-{rm.session_id}",
        agent_id=cfg.agent_id,
        session_id=rm.session_id,
        user_id=cfg.user_id,
    )
    set_current_span(span)

    return rm


def _shutdown() -> None:
    """Atexit handler: end the global run if still active."""
    if _run_manager and not _run_manager._ended:
        _run_manager.end()


def _patch_openai(config, cache, client, run_manager):
    """Monkey-patch OpenAI classes to use learn wrappers."""
    try:
        import openai
        from mubit.learn._openai_learn import wrap_openai_learn

        if not getattr(openai.OpenAI, "_mubit_learn_instrumented", False):
            # Uninstrument auto module first if it was active
            if getattr(openai.OpenAI, "_mubit_instrumented", False):
                from mubit.auto._instrument import uninstrument as auto_uninstrument
                auto_uninstrument()

            _original_refs["openai.OpenAI.__init__"] = openai.OpenAI.__init__

            def patched_init(self, *args, **kwargs):
                _original_refs["openai.OpenAI.__init__"](self, *args, **kwargs)
                wrap_openai_learn(
                    self,
                    learn_config=config,
                    lesson_cache=cache,
                    learn_client=client,
                    run_manager=run_manager,
                    session_id=run_manager.session_id,
                    agent_id=config.agent_id,
                    user_id=config.user_id,
                    capture=config.capture,
                    min_length=config.min_length,
                    mubit_api_key=config.api_key,
                    mubit_endpoint=config.endpoint,
                )

            openai.OpenAI.__init__ = patched_init
            openai.OpenAI._mubit_learn_instrumented = True

        if not getattr(openai.AsyncOpenAI, "_mubit_learn_instrumented", False):
            _original_refs["openai.AsyncOpenAI.__init__"] = openai.AsyncOpenAI.__init__

            def patched_async_init(self, *args, **kwargs):
                _original_refs["openai.AsyncOpenAI.__init__"](self, *args, **kwargs)
                wrap_openai_learn(
                    self,
                    learn_config=config,
                    lesson_cache=cache,
                    learn_client=client,
                    run_manager=run_manager,
                    session_id=run_manager.session_id,
                    agent_id=config.agent_id,
                    user_id=config.user_id,
                    capture=config.capture,
                    min_length=config.min_length,
                    mubit_api_key=config.api_key,
                    mubit_endpoint=config.endpoint,
                )

            openai.AsyncOpenAI.__init__ = patched_async_init
            openai.AsyncOpenAI._mubit_learn_instrumented = True

    except ImportError:
        pass


def _patch_anthropic(config, cache, client, run_manager):
    """Monkey-patch Anthropic classes to use learn wrappers."""
    try:
        import anthropic
        from mubit.learn._anthropic_learn import wrap_anthropic_learn

        if not getattr(anthropic.Anthropic, "_mubit_learn_instrumented", False):
            if getattr(anthropic.Anthropic, "_mubit_instrumented", False):
                from mubit.auto._instrument import uninstrument as auto_uninstrument
                auto_uninstrument()

            _original_refs["anthropic.Anthropic.__init__"] = anthropic.Anthropic.__init__

            def patched_init(self, *args, **kwargs):
                _original_refs["anthropic.Anthropic.__init__"](self, *args, **kwargs)
                wrap_anthropic_learn(
                    self,
                    learn_config=config,
                    lesson_cache=cache,
                    learn_client=client,
                    run_manager=run_manager,
                    session_id=run_manager.session_id,
                    agent_id=config.agent_id,
                    user_id=config.user_id,
                    capture=config.capture,
                    min_length=config.min_length,
                    mubit_api_key=config.api_key,
                    mubit_endpoint=config.endpoint,
                )

            anthropic.Anthropic.__init__ = patched_init
            anthropic.Anthropic._mubit_learn_instrumented = True

        if not getattr(anthropic.AsyncAnthropic, "_mubit_learn_instrumented", False):
            _original_refs["anthropic.AsyncAnthropic.__init__"] = anthropic.AsyncAnthropic.__init__

            def patched_async_init(self, *args, **kwargs):
                _original_refs["anthropic.AsyncAnthropic.__init__"](self, *args, **kwargs)
                wrap_anthropic_learn(
                    self,
                    learn_config=config,
                    lesson_cache=cache,
                    learn_client=client,
                    run_manager=run_manager,
                    session_id=run_manager.session_id,
                    agent_id=config.agent_id,
                    user_id=config.user_id,
                    capture=config.capture,
                    min_length=config.min_length,
                    mubit_api_key=config.api_key,
                    mubit_endpoint=config.endpoint,
                )

            anthropic.AsyncAnthropic.__init__ = patched_async_init
            anthropic.AsyncAnthropic._mubit_learn_instrumented = True

    except ImportError:
        pass


def _patch_litellm(config, cache, client, run_manager):
    """Add learn callback to litellm."""
    try:
        import litellm
        from mubit.auto._litellm import MubitLiteLLMLogger
        from mubit.learn._litellm_learn import MubitLearnLiteLLMCallback

        # Add auto-capture logger if not already present
        if not getattr(litellm, "_mubit_instrumented", False):
            auto_cb = MubitLiteLLMLogger(
                session_id=run_manager.session_id,
                agent_id=config.agent_id,
                user_id=config.user_id,
                capture=config.capture,
                min_length=config.min_length,
                mubit_api_key=config.api_key,
                mubit_endpoint=config.endpoint,
            )
            if hasattr(litellm, "callbacks"):
                litellm.callbacks.append(auto_cb)
                litellm._mubit_instrumented = True

        # Add learn callback
        if not getattr(litellm, "_mubit_learn_instrumented", False):
            learn_cb = MubitLearnLiteLLMCallback(config, cache, client, run_manager)
            if hasattr(litellm, "callbacks"):
                litellm.callbacks.append(learn_cb)
                litellm._mubit_learn_instrumented = True

    except ImportError:
        pass


def _patch_google_genai(config, cache, client, run_manager):
    """Monkey-patch Google GenAI Client to use learn wrappers."""
    try:
        import google.genai
        from mubit.learn._google_genai import wrap_google_genai_learn

        if not getattr(google.genai.Client, "_mubit_learn_instrumented", False):
            _original_refs["google.genai.Client.__init__"] = google.genai.Client.__init__

            def patched_init(self, *args, **kwargs):
                _original_refs["google.genai.Client.__init__"](self, *args, **kwargs)
                wrap_google_genai_learn(
                    self,
                    learn_config=config,
                    lesson_cache=cache,
                    learn_client=client,
                    run_manager=run_manager,
                    session_id=run_manager.session_id,
                    agent_id=config.agent_id,
                    user_id=config.user_id,
                    capture=config.capture,
                    min_length=config.min_length,
                    mubit_api_key=config.api_key,
                    mubit_endpoint=config.endpoint,
                )

            google.genai.Client.__init__ = patched_init
            google.genai.Client._mubit_learn_instrumented = True

    except ImportError:
        pass
