"""LLM call wrapper with automatic PII redaction and audit logging.

Intercepts LLM calls, detects PII in inputs, optionally redacts before
sending to the provider, logs everything to the audit trail, and tracks
data flows for compliance documentation.
"""

import time
from typing import Any, Callable, Optional

from .audit import AuditLog
from . import pii


class ShieldedCall:
    """Result of a shielded LLM call."""

    def __init__(
        self,
        response: Any,
        pii_detected: dict[str, int],
        input_was_redacted: bool,
        audit_entry_hash: str,
    ):
        self.response = response
        self.pii_detected = pii_detected
        self.input_was_redacted = input_was_redacted
        self.audit_entry_hash = audit_entry_hash

    def __str__(self):
        return str(self.response)


def wrap_openai_call(
    client: Any,
    messages: list[dict],
    model: str = "gpt-4",
    audit_log: Optional[AuditLog] = None,
    purpose: str = "",
    user_id: Optional[str] = None,
    redact_input: bool = False,
    **kwargs,
) -> ShieldedCall:
    """Wrap an OpenAI chat completion call with compliance features.

    Args:
        client: An OpenAI client instance.
        messages: Chat messages in OpenAI format.
        model: Model name.
        audit_log: AuditLog instance for recording.
        purpose: Processing purpose (for audit trail).
        user_id: User identifier (for data subject requests).
        redact_input: If True, redact PII before sending to OpenAI.
        **kwargs: Additional arguments passed to the API call.

    Returns:
        ShieldedCall with the response and compliance metadata.
    """
    # Combine message content for PII scanning
    input_text = "\n".join(
        m.get("content", "") for m in messages if m.get("content")
    )
    pii_found = pii.summary(input_text)

    # Optionally redact before sending
    if redact_input and pii_found:
        messages = [
            {**m, "content": pii.redact(m.get("content", ""))}
            if m.get("content") else m
            for m in messages
        ]

    # Make the actual API call
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        **kwargs,
    )

    output_text = ""
    tokens_in = 0
    tokens_out = 0
    if response.choices:
        output_text = response.choices[0].message.content or ""
    if response.usage:
        tokens_in = response.usage.prompt_tokens
        tokens_out = response.usage.completion_tokens

    # Log to audit trail
    entry_hash = ""
    if audit_log:
        entry = audit_log.record(
            event_type="llm_call",
            provider="openai",
            model=model,
            purpose=purpose,
            user_id=user_id,
            input_text=input_text,
            output_text=output_text,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata={"redacted": redact_input},
        )
        entry_hash = entry.entry_hash

    return ShieldedCall(
        response=response,
        pii_detected=pii_found,
        input_was_redacted=redact_input and bool(pii_found),
        audit_entry_hash=entry_hash,
    )


def wrap_anthropic_call(
    client: Any,
    messages: list[dict],
    model: str = "claude-sonnet-4-20250514",
    audit_log: Optional[AuditLog] = None,
    purpose: str = "",
    user_id: Optional[str] = None,
    redact_input: bool = False,
    max_tokens: int = 1024,
    **kwargs,
) -> ShieldedCall:
    """Wrap an Anthropic message call with compliance features.

    Args:
        client: An Anthropic client instance.
        messages: Chat messages in Anthropic format.
        model: Model name.
        audit_log: AuditLog instance for recording.
        purpose: Processing purpose (for audit trail).
        user_id: User identifier (for data subject requests).
        redact_input: If True, redact PII before sending.
        max_tokens: Max tokens for response.
        **kwargs: Additional arguments passed to the API call.

    Returns:
        ShieldedCall with the response and compliance metadata.
    """
    input_text = "\n".join(
        m.get("content", "") for m in messages if m.get("content")
    )
    pii_found = pii.summary(input_text)

    if redact_input and pii_found:
        messages = [
            {**m, "content": pii.redact(m.get("content", ""))}
            if m.get("content") else m
            for m in messages
        ]

    response = client.messages.create(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
        **kwargs,
    )

    output_text = ""
    tokens_in = 0
    tokens_out = 0
    if response.content:
        output_text = response.content[0].text if response.content else ""
    if response.usage:
        tokens_in = response.usage.input_tokens
        tokens_out = response.usage.output_tokens

    entry_hash = ""
    if audit_log:
        entry = audit_log.record(
            event_type="llm_call",
            provider="anthropic",
            model=model,
            purpose=purpose,
            user_id=user_id,
            input_text=input_text,
            output_text=output_text,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata={"redacted": redact_input},
        )
        entry_hash = entry.entry_hash

    return ShieldedCall(
        response=response,
        pii_detected=pii_found,
        input_was_redacted=redact_input and bool(pii_found),
        audit_entry_hash=entry_hash,
    )


def wrap_generic_call(
    fn: Callable,
    input_text: str,
    provider: str = "custom",
    model: str = "custom",
    audit_log: Optional[AuditLog] = None,
    purpose: str = "",
    user_id: Optional[str] = None,
    redact_input: bool = False,
    **kwargs,
) -> ShieldedCall:
    """Wrap any LLM call function with compliance features.

    For providers without a specific wrapper. Pass any callable that
    takes text input and returns text output.

    Args:
        fn: Callable that takes (text, **kwargs) and returns response text.
        input_text: The input text.
        provider: Provider name for audit trail.
        model: Model name for audit trail.
        audit_log: AuditLog instance.
        purpose: Processing purpose.
        user_id: User identifier.
        redact_input: If True, redact PII before calling fn.
        **kwargs: Passed to fn.

    Returns:
        ShieldedCall with the response and compliance metadata.
    """
    pii_found = pii.summary(input_text)

    call_input = pii.redact(input_text) if (redact_input and pii_found) else input_text
    response = fn(call_input, **kwargs)

    output_text = str(response) if response else ""

    entry_hash = ""
    if audit_log:
        entry = audit_log.record(
            event_type="llm_call",
            provider=provider,
            model=model,
            purpose=purpose,
            user_id=user_id,
            input_text=input_text,
            output_text=output_text,
            tokens_in=0,
            tokens_out=0,
            metadata={"redacted": redact_input},
        )
        entry_hash = entry.entry_hash

    return ShieldedCall(
        response=response,
        pii_detected=pii_found,
        input_was_redacted=redact_input and bool(pii_found),
        audit_entry_hash=entry_hash,
    )
