"""
Ollama SDK instrumentation.

Patches Ollama SDK methods to capture:
- Chat (sync module-level and async client)
- Generate (sync module-level and async client)
- Streaming responses for both chat and generate

Ollama uses module-level functions (ollama.chat, ollama.generate) and an
AsyncClient class (ollama.AsyncClient.chat, ollama.AsyncClient.generate).
Responses are plain dicts, not Pydantic models.

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

import logging
import time
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, TypeVar

import wrapt
from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

# Track if already instrumented
_instrumented = False


def instrument_ollama(module: Any) -> None:
    """
    Apply instrumentation to Ollama module.

    This patches:
    - ollama.chat (sync, module-level function)
    - ollama.generate (sync, module-level function)
    - ollama.AsyncClient.chat (async class method)
    - ollama.AsyncClient.generate (async class method)
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch sync module-level functions
        wrapt.wrap_function_wrapper(module, "chat", _wrap_chat)
        wrapt.wrap_function_wrapper(module, "generate", _wrap_generate)

        # Patch async client methods
        wrapt.wrap_function_wrapper(module, "AsyncClient.chat", _wrap_async_chat)
        wrapt.wrap_function_wrapper(
            module, "AsyncClient.generate", _wrap_async_generate
        )

        _instrumented = True
        logger.debug("Instrumented Ollama SDK")

    except Exception as e:
        logger.warning(f"Failed to instrument Ollama SDK: {e}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _get_tracer() -> Optional[Any]:
    """Get the Risicare tracer if available."""
    try:
        from risicare.tracer import get_tracer

        return get_tracer()
    except ImportError:
        return None


def _get_config() -> Optional[Any]:
    """Get the Risicare client config if available."""
    try:
        from risicare.client import get_client

        client = get_client()
        return client.config if client else None
    except ImportError:
        return None


def _should_trace_content() -> bool:
    """Check if content should be traced."""
    config = _get_config()
    return config.trace_content if config else False


def _extract_model_info(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """Extract model information from request kwargs."""
    return {
        "gen_ai.system": "ollama",
        "gen_ai.request.model": kwargs.get("model", "unknown"),
        "gen_ai.request.stream": kwargs.get("stream", False),
    }


def _get_model(args: tuple, kwargs: Dict[str, Any]) -> str:
    """Extract model name from args/kwargs."""
    return kwargs.get("model", args[0] if args else "unknown")


# ---------------------------------------------------------------------------
# Input capture
# ---------------------------------------------------------------------------


def _capture_chat_messages(
    span: Any, messages: list, trace_content: bool
) -> None:
    """Capture input messages to span.

    ``messages`` is a list of dicts, e.g.
    ``[{"role": "user", "content": "Hi"}]``.
    """
    if not trace_content:
        span.set_attribute("gen_ai.prompt.count", len(messages))
        return

    for i, msg in enumerate(messages[:10]):  # Limit to 10 messages
        span.set_attribute(f"gen_ai.prompt.{i}.role", msg.get("role", ""))
        content = msg.get("content", "")
        if isinstance(content, str) and content:
            span.set_attribute(
                f"gen_ai.prompt.{i}.content",
                content[:10000] if len(content) > 10000 else content,
            )


# ---------------------------------------------------------------------------
# Response capture
# ---------------------------------------------------------------------------


def _capture_chat_response(
    span: Any,
    response: Dict[str, Any],
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture chat response data into span.

    Ollama chat responses are plain dicts::

        {
            "model": "llama3.1",
            "message": {"role": "assistant", "content": "Hello!"},
            "done": True,
            "prompt_eval_count": 10,
            "eval_count": 20,
            "total_duration": 1234567890,
        }
    """
    span.set_attribute("gen_ai.latency_ms", latency_ms)
    span.set_attribute("gen_ai.response.model", response.get("model", "unknown"))

    prompt_tokens = response.get("prompt_eval_count", 0)
    completion_tokens = response.get("eval_count", 0)
    span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
    span.set_attribute("gen_ai.usage.completion_tokens", completion_tokens)
    span.set_attribute(
        "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
    )

    if trace_content:
        message = response.get("message", {})
        content = message.get("content", "")
        if content:
            span.set_attribute(
                "gen_ai.completion.content",
                content[:10000] if len(content) > 10000 else content,
            )


def _capture_generate_response(
    span: Any,
    response: Dict[str, Any],
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture generate response data into span.

    Ollama generate responses are plain dicts::

        {
            "model": "llama3.1",
            "response": "Hello!",
            "done": True,
            "prompt_eval_count": 10,
            "eval_count": 20,
        }
    """
    span.set_attribute("gen_ai.latency_ms", latency_ms)
    span.set_attribute("gen_ai.response.model", response.get("model", "unknown"))

    prompt_tokens = response.get("prompt_eval_count", 0)
    completion_tokens = response.get("eval_count", 0)
    span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
    span.set_attribute("gen_ai.usage.completion_tokens", completion_tokens)
    span.set_attribute(
        "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
    )

    if trace_content:
        content = response.get("response", "")
        if content:
            span.set_attribute(
                "gen_ai.completion.content",
                content[:10000] if len(content) > 10000 else content,
            )


# ---------------------------------------------------------------------------
# Streaming wrappers — Chat
# ---------------------------------------------------------------------------


def _wrap_stream_chat_sync(
    stream: Iterator[Dict[str, Any]],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Dict[str, Any]]:
    """Wrap a sync streaming chat response.

    Streaming chunks::

        {"message": {"content": "Hel"}, "done": False}
        {"message": {"content": "lo"},  "done": False}
        {"message": {"content": "!"},   "done": True,
         "prompt_eval_count": 10, "eval_count": 3}
    """
    chunk_count = 0
    content_buffer: list[str] = []
    content_length = 0
    token_data: Dict[str, int] = {}

    try:
        for chunk in stream:
            chunk_count += 1

            # Accumulate content (bounded)
            if trace_content and content_length < 10000:
                msg_content = chunk.get("message", {}).get("content", "")
                if msg_content:
                    content_buffer.append(msg_content)
                    content_length += len(msg_content)

            # Final chunk carries token counts
            if chunk.get("done", False):
                token_data["prompt_tokens"] = chunk.get("prompt_eval_count", 0)
                token_data["completion_tokens"] = chunk.get("eval_count", 0)

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if token_data:
            prompt_tokens = token_data.get("prompt_tokens", 0)
            completion_tokens = token_data.get("completion_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
            span.set_attribute(
                "gen_ai.usage.completion_tokens", completion_tokens
            )
            span.set_attribute(
                "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
            )

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


async def _wrap_stream_chat_async(
    stream: AsyncIterator[Dict[str, Any]],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Dict[str, Any]]:
    """Wrap an async streaming chat response."""
    chunk_count = 0
    content_buffer: list[str] = []
    content_length = 0
    token_data: Dict[str, int] = {}

    try:
        async for chunk in stream:
            chunk_count += 1

            if trace_content and content_length < 10000:
                msg_content = chunk.get("message", {}).get("content", "")
                if msg_content:
                    content_buffer.append(msg_content)
                    content_length += len(msg_content)

            if chunk.get("done", False):
                token_data["prompt_tokens"] = chunk.get("prompt_eval_count", 0)
                token_data["completion_tokens"] = chunk.get("eval_count", 0)

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if token_data:
            prompt_tokens = token_data.get("prompt_tokens", 0)
            completion_tokens = token_data.get("completion_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
            span.set_attribute(
                "gen_ai.usage.completion_tokens", completion_tokens
            )
            span.set_attribute(
                "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
            )

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


# ---------------------------------------------------------------------------
# Streaming wrappers — Generate
# ---------------------------------------------------------------------------


def _wrap_stream_generate_sync(
    stream: Iterator[Dict[str, Any]],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Dict[str, Any]]:
    """Wrap a sync streaming generate response.

    Streaming chunks::

        {"response": "Hel", "done": False}
        {"response": "lo!", "done": True,
         "prompt_eval_count": 10, "eval_count": 3}
    """
    chunk_count = 0
    content_buffer: list[str] = []
    content_length = 0
    token_data: Dict[str, int] = {}

    try:
        for chunk in stream:
            chunk_count += 1

            if trace_content and content_length < 10000:
                text = chunk.get("response", "")
                if text:
                    content_buffer.append(text)
                    content_length += len(text)

            if chunk.get("done", False):
                token_data["prompt_tokens"] = chunk.get("prompt_eval_count", 0)
                token_data["completion_tokens"] = chunk.get("eval_count", 0)

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if token_data:
            prompt_tokens = token_data.get("prompt_tokens", 0)
            completion_tokens = token_data.get("completion_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
            span.set_attribute(
                "gen_ai.usage.completion_tokens", completion_tokens
            )
            span.set_attribute(
                "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
            )

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


async def _wrap_stream_generate_async(
    stream: AsyncIterator[Dict[str, Any]],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Dict[str, Any]]:
    """Wrap an async streaming generate response."""
    chunk_count = 0
    content_buffer: list[str] = []
    content_length = 0
    token_data: Dict[str, int] = {}

    try:
        async for chunk in stream:
            chunk_count += 1

            if trace_content and content_length < 10000:
                text = chunk.get("response", "")
                if text:
                    content_buffer.append(text)
                    content_length += len(text)

            if chunk.get("done", False):
                token_data["prompt_tokens"] = chunk.get("prompt_eval_count", 0)
                token_data["completion_tokens"] = chunk.get("eval_count", 0)

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if token_data:
            prompt_tokens = token_data.get("prompt_tokens", 0)
            completion_tokens = token_data.get("completion_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
            span.set_attribute(
                "gen_ai.usage.completion_tokens", completion_tokens
            )
            span.set_attribute(
                "gen_ai.usage.total_tokens", prompt_tokens + completion_tokens
            )

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


# ---------------------------------------------------------------------------
# Main wrappers — Chat
# ---------------------------------------------------------------------------


def _wrap_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ollama.chat (sync, module-level function)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = _get_model(args, kwargs)
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"ollama.chat/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            messages = kwargs.get("messages", [])
            _capture_chat_messages(span, messages, trace_content)

            try:
                start_time = time.perf_counter()
                response = wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_chat_response(span, response, latency_ms, trace_content)
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"ollama.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    messages = kwargs.get("messages", [])
    _capture_chat_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_stream_chat_sync(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


async def _wrap_async_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ollama.AsyncClient.chat (async)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = _get_model(args, kwargs)
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"ollama.chat/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            messages = kwargs.get("messages", [])
            _capture_chat_messages(span, messages, trace_content)

            try:
                start_time = time.perf_counter()
                response = await wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_chat_response(span, response, latency_ms, trace_content)
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"ollama.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    messages = kwargs.get("messages", [])
    _capture_chat_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_stream_chat_async(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


# ---------------------------------------------------------------------------
# Main wrappers — Generate
# ---------------------------------------------------------------------------


def _wrap_generate(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ollama.generate (sync, module-level function)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = _get_model(args, kwargs)
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"ollama.generate/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            prompt = kwargs.get("prompt", args[1] if len(args) > 1 else "")
            if trace_content and prompt:
                span.set_attribute(
                    "gen_ai.prompt.content",
                    prompt[:10000] if len(prompt) > 10000 else prompt,
                )

            try:
                start_time = time.perf_counter()
                response = wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_generate_response(
                    span, response, latency_ms, trace_content
                )
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"ollama.generate/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    prompt = kwargs.get("prompt", args[1] if len(args) > 1 else "")
    if trace_content and prompt:
        span.set_attribute(
            "gen_ai.prompt.content",
            prompt[:10000] if len(prompt) > 10000 else prompt,
        )

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_stream_generate_sync(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


async def _wrap_async_generate(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ollama.AsyncClient.generate (async)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = _get_model(args, kwargs)
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"ollama.generate/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            prompt = kwargs.get("prompt", args[1] if len(args) > 1 else "")
            if trace_content and prompt:
                span.set_attribute(
                    "gen_ai.prompt.content",
                    prompt[:10000] if len(prompt) > 10000 else prompt,
                )

            try:
                start_time = time.perf_counter()
                response = await wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_generate_response(
                    span, response, latency_ms, trace_content
                )
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"ollama.generate/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    prompt = kwargs.get("prompt", args[1] if len(args) > 1 else "")
    if trace_content and prompt:
        span.set_attribute(
            "gen_ai.prompt.content",
            prompt[:10000] if len(prompt) > 10000 else prompt,
        )

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_stream_generate_async(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise
