"""secure_llm_call() — sanitize prompt before LLM, restore pseudonyms in response.

Usage::

    from upf import secure_llm_call

    response = secure_llm_call(
        llm_fn=client.messages.create,
        prompt="My name is John Smith and my email is john@example.com. Summarize my profile.",
        # extra kwargs are forwarded to llm_fn
        model="claude-3-haiku-20240307",
        max_tokens=256,
    )
    # response is the return value of llm_fn, with PII restored in any text fields

The function:
1. Sanitizes the prompt using PSEUDONYM mode so the LLM sees PERSON_1 / EMAIL_1 etc.
2. Calls ``llm_fn(prompt=sanitized_prompt, **kwargs)`` (or passes as positional arg if
   the function uses a different keyword — see ``prompt_kwarg`` parameter).
3. Extracts text from the response, restores the pseudonym→original mapping,
   and returns a ``SecureLLMResult`` with both the raw and restored text.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Optional

from upf.config import UPFConfig
from upf.pipeline.text_pipeline import TextPipeline
from upf.redaction.pseudonym import PseudonymSession
from upf.redaction.text import RedactionMode


@dataclass
class SecureLLMResult:
    """Result from :func:`secure_llm_call`.

    Attributes:
        sanitized_prompt: The prompt that was actually sent to the LLM.
        raw_response: The unmodified return value from ``llm_fn``.
        restored_text: Response text with pseudonyms replaced by original PII values.
        session: The :class:`PseudonymSession` used, in case you need extra mappings.
    """

    sanitized_prompt: str
    raw_response: Any
    restored_text: str
    session: PseudonymSession = field(default_factory=PseudonymSession)


def _extract_text(response: Any) -> str:
    """Best-effort extraction of text content from common LLM response shapes."""
    # Anthropic SDK: response.content[0].text
    if hasattr(response, "content"):
        content = response.content
        if isinstance(content, list) and content:
            first = content[0]
            if hasattr(first, "text"):
                return str(first.text)
        if isinstance(content, str):
            return content

    # OpenAI SDK: response.choices[0].message.content
    if hasattr(response, "choices"):
        choices = response.choices
        if choices:
            msg = choices[0]
            if hasattr(msg, "message") and hasattr(msg.message, "content"):
                return str(msg.message.content or "")
            if hasattr(msg, "text"):
                return str(msg.text)

    # Generic: str fallback
    return str(response)


def secure_llm_call(
    llm_fn: Callable[..., Any],
    prompt: str,
    *,
    config: Optional[UPFConfig] = None,
    session: Optional[PseudonymSession] = None,
    prompt_kwarg: str = "prompt",
    **kwargs: Any,
) -> SecureLLMResult:
    """Call *llm_fn* with a PII-sanitized prompt and restore pseudonyms in the response.

    Parameters
    ----------
    llm_fn:
        Any callable that accepts a text prompt and returns a response with text
        content (Anthropic, OpenAI, or custom callables all work).
    prompt:
        The original prompt that may contain PII.
    config:
        Optional :class:`~upf.config.UPFConfig`.  ``redaction_mode`` is forced to
        ``PSEUDONYM`` regardless of what you pass here, so that the session can restore
        values from the response.
    session:
        Provide an existing :class:`~upf.redaction.pseudonym.PseudonymSession` to
        share state across multiple calls (e.g., multi-turn conversations).  A new
        session is created automatically when ``None``.
    prompt_kwarg:
        Keyword argument name used to pass the prompt to *llm_fn*.  Defaults to
        ``"prompt"``.  Use ``"messages"`` for OpenAI chat-style APIs (you will also
        need to format the list yourself and pass via ``**kwargs``).
    **kwargs:
        Extra keyword arguments forwarded verbatim to *llm_fn*.

    Returns
    -------
    SecureLLMResult
    """
    # Force PSEUDONYM mode so we can restore after the call
    cfg = UPFConfig(
        **{
            k: v
            for k, v in (config.__dict__ if config else UPFConfig().__dict__).items()
            if k != "redaction_mode"
        }
    )
    cfg.redaction_mode = RedactionMode.PSEUDONYM

    pipeline = TextPipeline(cfg)
    s = session or PseudonymSession()

    sanitized_prompt = pipeline.sanitize(prompt, session=s)

    raw_response = llm_fn(**{prompt_kwarg: sanitized_prompt}, **kwargs)

    raw_text = _extract_text(raw_response)
    restored_text = s.restore(raw_text)

    return SecureLLMResult(
        sanitized_prompt=sanitized_prompt,
        raw_response=raw_response,
        restored_text=restored_text,
        session=s,
    )
