"""
MuBit Learn Google GenAI Wrapper.

Extends auto-capture with pre-call lesson retrieval and injection
for google.genai.Client.models.generate_content().
"""

import inspect
import json
import logging
import time
import uuid
from typing import Any, Optional

from mubit.auto._context import get_current_span, is_capture_enabled
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._extraction import extract_structured_items
from mubit.learn._injection import inject_context_openai
from mubit.learn._lesson_cache import LessonCache
from mubit.learn._run_manager import RunManager

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


def _extract_query_from_contents(contents: Any) -> str:
    """Extract query text from google.genai contents parameter.

    Contents can be:
    - str: plain text prompt
    - list[str | Part | Content]: list of parts
    - Content: a single Content object
    """
    if isinstance(contents, str):
        return contents[:200]

    if isinstance(contents, list):
        texts = []
        for item in reversed(contents):
            if isinstance(item, str):
                texts.append(item)
                break
            # Content object with parts
            if hasattr(item, "parts"):
                for part in item.parts:
                    if hasattr(part, "text") and part.text:
                        texts.append(part.text)
                if texts:
                    break
            # Part object
            if hasattr(item, "text") and item.text:
                texts.append(item.text)
                break
        return " ".join(texts)[:200]

    # Single Content object
    if hasattr(contents, "parts"):
        parts = []
        for part in contents.parts:
            if hasattr(part, "text") and part.text:
                parts.append(part.text)
        return " ".join(parts)[:200]

    return str(contents)[:200]


def _extract_response_text(response: Any) -> str:
    """Extract text from GenerateContentResponse."""
    try:
        if hasattr(response, "text"):
            return response.text or ""
    except Exception:
        pass

    try:
        if hasattr(response, "candidates") and response.candidates:
            candidate = response.candidates[0]
            if hasattr(candidate, "content") and candidate.content:
                parts = candidate.content.parts or []
                texts = [p.text for p in parts if hasattr(p, "text") and p.text]
                return "\n".join(texts)
    except Exception:
        pass

    return ""


def _contents_to_messages(contents: Any) -> list:
    """Convert google.genai contents to OpenAI-style messages for ingestion."""
    messages = []
    if isinstance(contents, str):
        messages.append({"role": "user", "content": contents})
    elif isinstance(contents, list):
        for item in contents:
            if isinstance(item, str):
                messages.append({"role": "user", "content": item})
            elif hasattr(item, "role") and hasattr(item, "parts"):
                parts_text = []
                for part in (item.parts or []):
                    if hasattr(part, "text") and part.text:
                        parts_text.append(part.text)
                if parts_text:
                    role = "assistant" if item.role == "model" else (item.role or "user")
                    messages.append({"role": role, "content": "\n".join(parts_text)})
            elif hasattr(item, "text") and item.text:
                messages.append({"role": "user", "content": item.text})
    elif hasattr(contents, "parts"):
        parts_text = []
        for part in (contents.parts or []):
            if hasattr(part, "text") and part.text:
                parts_text.append(part.text)
        if parts_text:
            messages.append({"role": "user", "content": "\n".join(parts_text)})

    return messages


def _inject_context_genai(contents: Any, context_block: str) -> Any:
    """Inject context into google.genai contents.

    Prepends context as a system-like instruction in the contents.
    """
    if not context_block or not context_block.strip():
        return contents

    memory_text = (
        f"<memory_context>\n{context_block}\n</memory_context>\n\n"
        "Use the above memory context to inform your response.\n\n---\n\n"
    )

    if isinstance(contents, str):
        return memory_text + contents

    if isinstance(contents, list) and contents:
        first = contents[0]
        if isinstance(first, str):
            return [memory_text + first] + contents[1:]
        # For structured content, prepend a text part
        try:
            from google.genai import types as genai_types
            context_content = genai_types.Content(
                parts=[genai_types.Part(text=memory_text)],
                role="user",
            )
            return [context_content] + list(contents)
        except ImportError:
            return contents

    return contents


def wrap_google_genai_learn(
    client: Any,
    *,
    learn_config: LearnConfig,
    lesson_cache: LessonCache,
    learn_client: LearnClient,
    run_manager: RunManager,
    session_id: str = "",
    agent_id: str = "auto",
    user_id: str = "",
    capture: str = "all",
    min_length: int = 0,
    mubit_api_key: Optional[str] = None,
    mubit_endpoint: Optional[str] = None,
) -> Any:
    """Wrap a google.genai Client with learn (inject + ingest) capabilities."""
    if not hasattr(client, "models"):
        return client

    if getattr(client, "_mubit_learn_wrapped", None) is True:
        return client

    models = client.models

    # Set up ingest worker for auto-capture
    worker = IngestWorker(api_key=mubit_api_key, endpoint=mubit_endpoint)
    worker.start()

    # Wrap generate_content
    if hasattr(models, "generate_content"):
        original_generate = models.generate_content

        def _enrich_kwargs(kwargs):
            if not learn_config.inject_lessons:
                return kwargs

            contents = kwargs.get("contents")
            if contents is None and len(kwargs) > 0:
                # contents might be a positional arg handled by caller
                return kwargs

            query = _extract_query_from_contents(contents)
            if not query:
                return kwargs

            context_block = lesson_cache.get(run_manager.session_id, query)
            if context_block is None:
                try:
                    context_block = learn_client.get_context(
                        session_id=run_manager.session_id,
                        query=query,
                        max_token_budget=learn_config.max_token_budget,
                        entry_types=learn_config.entry_types,
                        sections=learn_config.context_sections,
                    )
                    lesson_cache.set(run_manager.session_id, query, context_block)
                except Exception as e:
                    if learn_config.fail_open:
                        logger.debug("mubit.learn context fetch failed: %s", e)
                        context_block = ""
                    else:
                        raise

            if context_block:
                kwargs = dict(kwargs)
                kwargs["contents"] = _inject_context_genai(
                    kwargs.get("contents"), context_block
                )

            return kwargs

        def _capture_interaction(kwargs, response):
            """Capture the interaction for ingestion."""
            try:
                contents = kwargs.get("contents")
                messages = _contents_to_messages(contents)
                assistant_text = _extract_response_text(response)
                model_name = kwargs.get("model", "unknown")

                items = build_items(
                    messages=messages,
                    assistant_text=assistant_text,
                    model=str(model_name),
                    latency_ms=0,
                    capture=capture,
                    min_length=min_length,
                    user_id=learn_config.user_id,
                )
                if items:
                    span = get_current_span()
                    sid = getattr(span, "session_id", None) or run_manager.session_id
                    aid = getattr(span, "agent_id", None) or learn_config.agent_id
                    worker.enqueue(sid, aid, items)

                # Auto-extraction
                if learn_config.auto_extract and assistant_text:
                    extracted = extract_structured_items(
                        messages=messages,
                        assistant_text=assistant_text,
                        model=str(model_name),
                        user_id=learn_config.user_id,
                    )
                    if extracted:
                        span = get_current_span()
                        sid = getattr(span, "session_id", None) or run_manager.session_id
                        aid = getattr(span, "agent_id", None) or learn_config.agent_id
                        worker.enqueue(sid, aid, extracted)
            except Exception as e:
                logger.debug("mubit.learn genai capture failed: %s", e)

        is_async = inspect.iscoroutinefunction(original_generate)

        if is_async:
            async def learn_generate_content(*args, **kwargs):
                # Handle positional 'contents' arg
                if args and "contents" not in kwargs:
                    # First positional arg after 'model' is 'contents'
                    if len(args) >= 2:
                        kwargs["model"] = args[0]
                        kwargs["contents"] = args[1]
                        args = args[2:]
                    elif len(args) == 1:
                        kwargs["contents"] = args[0]
                        args = ()

                kwargs = _enrich_kwargs(kwargs)
                result = await original_generate(*args, **kwargs)
                run_manager.increment()
                _capture_interaction(kwargs, result)
                return result

            models.generate_content = learn_generate_content
        else:
            def learn_generate_content(*args, **kwargs):
                if args and "contents" not in kwargs:
                    if len(args) >= 2:
                        kwargs["model"] = args[0]
                        kwargs["contents"] = args[1]
                        args = args[2:]
                    elif len(args) == 1:
                        kwargs["contents"] = args[0]
                        args = ()

                kwargs = _enrich_kwargs(kwargs)
                result = original_generate(*args, **kwargs)
                run_manager.increment()
                _capture_interaction(kwargs, result)
                return result

            models.generate_content = learn_generate_content

    # Wrap generate_content_stream similarly
    if hasattr(models, "generate_content_stream"):
        original_stream = models.generate_content_stream

        is_stream_async = inspect.iscoroutinefunction(original_stream)

        if is_stream_async:
            async def learn_generate_stream(*args, **kwargs):
                if args and "contents" not in kwargs:
                    if len(args) >= 2:
                        kwargs["model"] = args[0]
                        kwargs["contents"] = args[1]
                        args = args[2:]
                    elif len(args) == 1:
                        kwargs["contents"] = args[0]
                        args = ()

                kwargs = _enrich_kwargs(kwargs)
                result = await original_stream(*args, **kwargs)
                run_manager.increment()
                return result

            models.generate_content_stream = learn_generate_stream
        else:
            def learn_generate_stream(*args, **kwargs):
                if args and "contents" not in kwargs:
                    if len(args) >= 2:
                        kwargs["model"] = args[0]
                        kwargs["contents"] = args[1]
                        args = args[2:]
                    elif len(args) == 1:
                        kwargs["contents"] = args[0]
                        args = ()

                kwargs = _enrich_kwargs(kwargs)
                result = original_stream(*args, **kwargs)
                run_manager.increment()
                return result

            models.generate_content_stream = learn_generate_stream

    client._mubit_learn_wrapped = True
    client._mubit_worker = worker
    return client
