"""
Main guard() function — the primary entry point for scandar-guard.

Supports sync and async Anthropic + OpenAI clients.
Handles streaming by buffering the full response before inspection.
Zero data leaves your environment.
"""
from __future__ import annotations

import json
from contextlib import asynccontextmanager, contextmanager
from typing import Any, AsyncGenerator, Generator, Optional

from .anomaly import AnomalyTracker
from .asg_telemetry import AsgTelemetrySender
from .audit import AuditLogger
from .config import GuardConfig
from .findings import RuntimeFinding
from .inspector import MessageInspector
from .multi_turn import MultiTurnTracker
from .profile import AgentProfile
from .scorer import ThreatScorer
from .session import SessionManager
from .tool_inspector import ToolCallInspector


class ScandarBlockedError(Exception):
    """Raised in block mode when a critical finding is detected."""
    def __init__(self, finding: RuntimeFinding) -> None:
        self.finding = finding
        super().__init__(
            f"[Scandar Guard] Blocked: {finding.title} (severity: {finding.severity})"
        )


def guard(client: Any, config: Optional[GuardConfig] = None) -> Any:
    """
    Wrap an LLM client with Scandar runtime security monitoring.

    Supports: Anthropic, AsyncAnthropic, OpenAI, AsyncOpenAI.
    Returns a wrapped client with identical API — one-line drop-in replacement.

    Streaming note: Guard buffers streaming responses to enable full inspection
    before yielding. Pass ``stream=True`` as normal — the response is returned
    as a completed Message object rather than a stream. Use ``stream=False``
    (the default) when lowest-latency is critical.

    Example::

        # Sync
        from anthropic import Anthropic
        from scandar_guard import guard
        client = guard(Anthropic())

        # Async
        from anthropic import AsyncAnthropic
        client = guard(AsyncAnthropic())
        response = await client.messages.create(...)
    """
    cfg = config or GuardConfig()
    session = SessionManager(session_id=cfg.session_id)
    msg_inspector = MessageInspector()
    tool_inspector = ToolCallInspector()
    anomaly = AnomalyTracker(session_id=session.session_id)
    logger = AuditLogger(log_file=cfg.log_file)
    scorer = ThreatScorer() if cfg.threat_scoring else None

    # Multi-turn tracking
    multi_turn = MultiTurnTracker(session_id=session.session_id) if cfg.multi_turn_detection else None
    if multi_turn:
        msg_inspector.set_multi_turn_tracker(multi_turn)

    # Per-agent behavioral profile
    profile = AgentProfile(agent_id=cfg.agent_id, profile_dir=cfg.profile_dir) if cfg.agent_id else None

    # ASG fleet telemetry sender
    asg = (
        AsgTelemetrySender(
            api_key=cfg.asg_api_key or "",
            agent_id=cfg.agent_id or "unnamed-agent",
            endpoint=cfg.asg_endpoint,
            parent_agent_id=cfg.asg_parent_agent_id,
        )
        if cfg.asg_telemetry and cfg.asg_api_key and cfg.agent_id
        else None
    )

    # Send session start
    if asg:
        asg.send_session_start(session.session_id)

    # Accumulated findings for threat scoring
    _call_pattern_findings: list[RuntimeFinding] = []
    _call_encoding_findings: list[RuntimeFinding] = []
    _call_multi_turn_findings: list[RuntimeFinding] = []
    _call_anomaly_findings: list[RuntimeFinding] = []
    _call_profile_findings: list[RuntimeFinding] = []

    client_name = type(client).__name__
    is_async = client_name.startswith("Async")

    # ── Shared helpers ────────────────────────────────────────────────────────

    def _handle_findings(findings: list[RuntimeFinding], event_type: str, meta: dict) -> None:
        if not findings:
            return
        session.increment_findings(len(findings))
        logger.log_event(event_type, session.session_id, findings, meta)

        # Categorize for threat scoring
        for f in findings:
            if f.category == "ENCODED_INJECTION":
                _call_encoding_findings.append(f)
            elif f.category in ("SPLIT_INJECTION", "SOCIAL_ENGINEERING_ESCALATION"):
                _call_multi_turn_findings.append(f)
            elif f.category in ("NEW_TOOL_MID_SESSION", "DATA_EXFIL_SEQUENCE", "COMMAND_THEN_EXFIL",
                                "TOOL_CALL_VOLUME_SPIKE", "RECON_EXFIL", "CODE_RECON_EXEC",
                                "PRIVILEGE_ESCALATION", "USER_DATA_EXFIL"):
                _call_anomaly_findings.append(f)
            elif f.category in ("PROFILE_DEVIATION", "PROFILE_FREQUENCY_DEVIATION"):
                _call_profile_findings.append(f)
            else:
                _call_pattern_findings.append(f)

        if cfg.mode == "block":
            for f in findings:
                if f.severity in cfg.block_on:
                    raise ScandarBlockedError(f)

    def _inspect_system(params: dict) -> list[RuntimeFinding]:
        system = params.get("system", "")
        if not system:
            return []
        if isinstance(system, list):
            system = " ".join(b.get("text", "") for b in system if isinstance(b, dict))
        drift = session.record_system_prompt(str(system))
        return [drift] if drift else []

    def _inspect_messages(messages: list) -> list[RuntimeFinding]:
        """
        Inspect all messages in a conversation history.

        Handles:
        - Plain string content (all roles)
        - Anthropic content block arrays (text, tool_result)
        - OpenAI function/tool role messages

        The tool_result block handling is the most critical path:
        external content fed back to the model after a tool call is
        the #1 runtime prompt injection vector.
        """
        findings: list[RuntimeFinding] = []
        for msg in messages:
            if not isinstance(msg, dict):
                continue
            role = msg.get("role", "user")
            content = msg.get("content")

            if content is None:
                continue

            if isinstance(content, str):
                findings.extend(
                    msg_inspector.inspect_message(role, content, session.session_id).findings
                )

            elif isinstance(content, list):
                for block in content:
                    if not isinstance(block, dict):
                        continue
                    btype = block.get("type", "")

                    if btype == "text":
                        findings.extend(
                            msg_inspector.inspect_message(
                                role, block.get("text", ""), session.session_id
                            ).findings
                        )

                    elif btype == "tool_result":
                        # CRITICAL: external content returned by a tool, fed back
                        # to the model. Primary vector for runtime prompt injection.
                        tool_use_id = block.get("tool_use_id", "unknown")
                        tool_content = block.get("content", "")
                        if isinstance(tool_content, str):
                            findings.extend(
                                tool_inspector.inspect_tool_result(
                                    tool_use_id, tool_content, session.session_id
                                ).findings
                            )
                        elif isinstance(tool_content, list):
                            for tc in tool_content:
                                if isinstance(tc, dict) and tc.get("type") == "text":
                                    findings.extend(
                                        tool_inspector.inspect_tool_result(
                                            tool_use_id, tc.get("text", ""), session.session_id
                                        ).findings
                                    )

            # OpenAI tool/function result role
            if role in ("tool", "function") and isinstance(content, str):
                findings.extend(
                    tool_inspector.inspect_tool_result(
                        msg.get("name", "unknown"), content, session.session_id
                    ).findings
                )

        return findings

    def _inspect_response(response: Any) -> list[RuntimeFinding]:
        """Inspect the model's response for threats and tool calls."""
        findings: list[RuntimeFinding] = []

        # Anthropic Message object
        if hasattr(response, "content") and isinstance(response.content, list):
            for block in response.content:
                btype = getattr(block, "type", None)
                if btype == "text":
                    findings.extend(
                        msg_inspector.inspect_response(
                            getattr(block, "text", ""), session.session_id
                        ).findings
                    )
                elif btype == "tool_use":
                    name = getattr(block, "name", "unknown")
                    inp = getattr(block, "input", {})
                    if not isinstance(inp, dict):
                        try:
                            inp = json.loads(inp) if isinstance(inp, str) else {}
                        except Exception:
                            inp = {}
                    tc_result = tool_inspector.inspect_tool_call(name, inp, session.session_id)
                    findings.extend(tc_result.findings)
                    findings.extend(anomaly.record_tool_call(name))
                    # Profile deviation check
                    if profile:
                        dev = profile.check_tool_deviation(name, session.session_id)
                        if dev:
                            findings.append(dev)
                    # ASG tool call event
                    if asg:
                        top_finding = next((f for f in sorted(tc_result.findings, key=lambda f: ["critical","high","medium","low","info"].index(f.severity)) ), None)
                        asg.send_tool_call(
                            session.session_id, name,
                            finding_category=top_finding.category if top_finding else None,
                            finding_severity=top_finding.severity if top_finding else None,
                        )
                    session.increment_tool_calls()

        # OpenAI ChatCompletion object
        elif hasattr(response, "choices"):
            for choice in response.choices:
                msg = getattr(choice, "message", None)
                if msg is None:
                    continue
                text = getattr(msg, "content", None)
                if text:
                    findings.extend(
                        msg_inspector.inspect_response(str(text), session.session_id).findings
                    )
                for tc in getattr(msg, "tool_calls", None) or []:
                    fn = getattr(tc, "function", None)
                    if fn is None:
                        continue
                    name = getattr(fn, "name", "unknown")
                    try:
                        args = json.loads(getattr(fn, "arguments", "{}") or "{}")
                    except Exception:
                        args = {}
                    findings.extend(tool_inspector.inspect_tool_call(name, args, session.session_id).findings)
                    findings.extend(anomaly.record_tool_call(name))
                    if profile:
                        dev = profile.check_tool_deviation(name, session.session_id)
                        if dev:
                            findings.append(dev)
                    session.increment_tool_calls()

        return findings

    # ── Sync wrapper ──────────────────────────────────────────────────────────

    _is_openai = client_name in ("OpenAI", "AsyncOpenAI") or "openai" in client_name.lower()

    def _make_call_sync(*args: Any, **kwargs: Any) -> Any:
        """Execute the real API call, buffering streams."""
        kwargs.pop("stream", None)
        if _is_openai:
            return client.chat.completions.create(*args, **kwargs)
        return client.messages.create(*args, **kwargs)

    def _create_sync(*args: Any, **kwargs: Any) -> Any:
        # Reset per-call finding accumulators
        _call_pattern_findings.clear()
        _call_encoding_findings.clear()
        _call_multi_turn_findings.clear()
        _call_anomaly_findings.clear()
        _call_profile_findings.clear()

        session.increment_messages()
        messages = kwargs.get("messages", args[0] if args else [])
        if not isinstance(messages, list):
            messages = []

        pre = _inspect_system(kwargs)
        pre.extend(_inspect_messages(messages))
        _handle_findings(pre, "pre_call_inspection", {"message_count": len(messages)})

        response = _make_call_sync(*args, **kwargs)

        post = _inspect_response(response)
        _handle_findings(post, "post_call_inspection", {"client_type": client_name})

        # Composite threat scoring
        if scorer:
            threat = scorer.score(
                pattern_findings=_call_pattern_findings,
                encoding_findings=_call_encoding_findings,
                multi_turn_findings=_call_multi_turn_findings,
                anomaly_findings=_call_anomaly_findings,
                profile_findings=_call_profile_findings,
                scan_trust_score=cfg.scan_trust_score,
                block_on=cfg.block_on if cfg.mode == "block" else None,
            )
            if threat.findings_count > 0:
                logger.log_event(
                    "threat_score",
                    session.session_id,
                    [],
                    {"score": threat.score, "classification": threat.classification, "signals": threat.signals},
                )
            if cfg.mode == "block" and threat.should_block:
                all_f = _call_pattern_findings + _call_encoding_findings + _call_multi_turn_findings
                critical_f = [f for f in all_f if f.severity in cfg.block_on]
                if critical_f:
                    raise ScandarBlockedError(critical_f[0])

        return response

    # ── Async wrapper ─────────────────────────────────────────────────────────

    async def _make_call_async(*args: Any, **kwargs: Any) -> Any:
        """Execute the real async API call, buffering streams."""
        kwargs.pop("stream", None)
        if _is_openai:
            return await client.chat.completions.create(*args, **kwargs)
        return await client.messages.create(*args, **kwargs)

    async def _create_async(*args: Any, **kwargs: Any) -> Any:
        _call_pattern_findings.clear()
        _call_encoding_findings.clear()
        _call_multi_turn_findings.clear()
        _call_anomaly_findings.clear()
        _call_profile_findings.clear()

        session.increment_messages()
        messages = kwargs.get("messages", args[0] if args else [])
        if not isinstance(messages, list):
            messages = []

        pre = _inspect_system(kwargs)
        pre.extend(_inspect_messages(messages))
        _handle_findings(pre, "pre_call_inspection", {"message_count": len(messages)})

        response = await _make_call_async(*args, **kwargs)

        post = _inspect_response(response)
        _handle_findings(post, "post_call_inspection", {"client_type": client_name})

        if scorer:
            threat = scorer.score(
                pattern_findings=_call_pattern_findings,
                encoding_findings=_call_encoding_findings,
                multi_turn_findings=_call_multi_turn_findings,
                anomaly_findings=_call_anomaly_findings,
                profile_findings=_call_profile_findings,
                scan_trust_score=cfg.scan_trust_score,
                block_on=cfg.block_on if cfg.mode == "block" else None,
            )
            if threat.findings_count > 0:
                logger.log_event(
                    "threat_score",
                    session.session_id,
                    [],
                    {"score": threat.score, "classification": threat.classification, "signals": threat.signals},
                )

        return response

    # ── Proxy classes ─────────────────────────────────────────────────────────

    class _SyncMessagesProxy:
        def __init__(self, target: Any) -> None:
            self._target = target

        def __getattr__(self, name: str) -> Any:
            if name == "create":
                return _create_sync
            return getattr(self._target, name)

    class _AsyncMessagesProxy:
        def __init__(self, target: Any) -> None:
            self._target = target

        def __getattr__(self, name: str) -> Any:
            if name == "create":
                return _create_async
            return getattr(self._target, name)

    class _SyncChatProxy:
        def __init__(self, target: Any) -> None:
            self._target = target

        def __getattr__(self, name: str) -> Any:
            if name == "completions":
                return _SyncMessagesProxy(getattr(self._target, name))
            return getattr(self._target, name)

    class _AsyncChatProxy:
        def __init__(self, target: Any) -> None:
            self._target = target

        def __getattr__(self, name: str) -> Any:
            if name == "completions":
                return _AsyncMessagesProxy(getattr(self._target, name))
            return getattr(self._target, name)

    # ── Session finalization (sent once on close, not per-call) ──────────────

    def _finalize_session() -> None:
        """Send session_end telemetry to ASG. Called once when the session closes."""
        if not asg or not scorer:
            return
        all_findings = (
            _call_pattern_findings + _call_encoding_findings +
            _call_multi_turn_findings + _call_anomaly_findings +
            _call_profile_findings
        )
        threat = scorer.score(
            _call_pattern_findings, _call_encoding_findings,
            _call_multi_turn_findings, _call_anomaly_findings,
            _call_profile_findings, cfg.scan_trust_score,
        )
        findings_summary: dict[str, int] = {}
        for f in all_findings:
            findings_summary[f.category] = findings_summary.get(f.category, 0) + 1
        stats = session.get_stats()
        asg.send_session_end(
            session.session_id,
            message_count=stats.get("message_count", 0),
            tool_call_count=stats.get("tool_call_count", 0),
            threat_score=threat.score,
            classification=threat.classification,
            findings_summary=findings_summary,
        )

    class SyncGuardedClient:
        def __getattr__(self, name: str) -> Any:
            attr = getattr(client, name)
            if name == "messages":
                return _SyncMessagesProxy(attr)
            if name == "chat":
                return _SyncChatProxy(attr)
            return attr

        def close(self) -> None:
            """Finalize the session and send telemetry to ASG."""
            _finalize_session()

        def __repr__(self) -> str:
            return f"GuardedClient({client_name}, session={session.session_id})"

    class AsyncGuardedClient:
        def __getattr__(self, name: str) -> Any:
            attr = getattr(client, name)
            if name == "messages":
                return _AsyncMessagesProxy(attr)
            if name == "chat":
                return _AsyncChatProxy(attr)
            return attr

        def close(self) -> None:
            """Finalize the session and send telemetry to ASG."""
            _finalize_session()

        def __repr__(self) -> str:
            return f"GuardedClient({client_name}, session={session.session_id})"

    return AsyncGuardedClient() if is_async else SyncGuardedClient()


# ── Context managers ──────────────────────────────────────────────────────────

@contextmanager
def guard_session(
    client: Any,
    config: Optional[GuardConfig] = None,
) -> Generator[Any, None, None]:
    """Sync context manager. Sends session_end telemetry on exit."""
    guarded = guard(client, config)
    try:
        yield guarded
    finally:
        guarded.close()


@asynccontextmanager
async def async_guard_session(
    client: Any,
    config: Optional[GuardConfig] = None,
) -> AsyncGenerator[Any, None]:
    """Async context manager. Sends session_end telemetry on exit."""
    guarded = guard(client, config)
    try:
        yield guarded
    finally:
        guarded.close()
