"""
sdk.py — Zero-config auto-instrumentation for OpenAI and Anthropic clients.

Usage:
    from qppg_service.sdk import instrument

    # Instrument the OpenAI client — every call is auto-scored
    import openai
    instrument(
        openai_client=openai,
        base_url="https://api.llm-guard-kit.com",  # or http://localhost:8000
        api_key="lgk_your_api_key",
        domain="my-app",
    )

    # Now use OpenAI normally — scoring happens transparently
    resp = openai.chat.completions.create(model="gpt-4o-mini", messages=[...])

    # Access the latest score from the patched client
    print(openai._lgk_last_score)  # {"risk_score": 0.23, "needs_review": False}

Alternatively, use the context manager:
    from qppg_service.sdk import GuardContext
    with GuardContext(api_key="lgk_...", domain="my-app", base_url="...") as guard:
        resp = openai.chat.completions.create(...)
    score = guard.last_score   # {"risk_score": ..., "needs_review": ...}

Environment variable shortcut:
    export LGK_API_KEY=lgk_...
    export LGK_BASE_URL=https://api.llm-guard-kit.com
    export LGK_DOMAIN=my-app
    # Then just call instrument() with no args.
"""

from __future__ import annotations

import os
import threading
from typing import Any

_SENTINEL = object()


def instrument(
    openai_client=None,
    anthropic_client=None,
    base_url: str = "",
    api_key: str = "",
    domain: str = "default",
    score_threshold: float = 0.65,
    on_high_risk=None,
) -> None:
    """
    Monkey-patch OpenAI / Anthropic client to auto-score every completion.

    Parameters
    ----------
    openai_client    : openai module or AsyncOpenAI/OpenAI instance (optional)
    anthropic_client : anthropic module or Anthropic instance (optional)
    base_url         : llm-guard-kit API base URL (env: LGK_BASE_URL)
    api_key          : llm-guard-kit API key (env: LGK_API_KEY)
    domain           : monitoring domain name (env: LGK_DOMAIN)
    score_threshold  : risk score above which on_high_risk is called (default 0.65)
    on_high_risk     : optional callable(score_dict) called when risk >= threshold
    """
    base_url  = base_url  or os.environ.get("LGK_BASE_URL", "http://localhost:8000")
    api_key   = api_key   or os.environ.get("LGK_API_KEY", "")
    domain    = domain    or os.environ.get("LGK_DOMAIN", "default")

    if openai_client is not None:
        _patch_openai(openai_client, base_url, api_key, domain, score_threshold, on_high_risk)
    if anthropic_client is not None:
        _patch_anthropic(anthropic_client, base_url, api_key, domain, score_threshold, on_high_risk)


# ── Internal helpers ─────────────────────────────────────────────────────────

def _score_async(base_url: str, api_key: str, domain: str,
                 question: str, steps: list, on_high_risk, threshold: float,
                 result_container: list) -> None:
    """Fire-and-forget scoring in a background thread."""
    import urllib.request
    import json
    payload = json.dumps({
        "question": question,
        "steps":    steps,
    }).encode()
    url = base_url.rstrip("/") + f"/api/{domain}/score"
    req = urllib.request.Request(
        url, data=payload,
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read())
            result_container.append(data)
            if on_high_risk and data.get("risk_score", 0) >= threshold:
                try:
                    on_high_risk(data)
                except Exception:
                    pass
    except Exception:
        pass  # never crash the user's code


def _fire_score(base_url: str, api_key: str, domain: str,
                question: str, steps: list,
                on_high_risk=None, threshold: float = 0.65,
                store_on=None, attr_name: str = "_lgk_last_score") -> None:
    """Non-blocking scoring: runs in background thread, stores result on `store_on`."""
    if not api_key:
        return
    container: list = []

    def _bg():
        _score_async(base_url, api_key, domain, question, steps, on_high_risk, threshold, container)
        if store_on is not None and container:
            try:
                setattr(store_on, attr_name, container[0])
            except Exception:
                pass

    t = threading.Thread(target=_bg, daemon=True)
    t.start()


# ── OpenAI patch ──────────────────────────────────────────────────────────────

def _patch_openai(client, base_url: str, api_key: str, domain: str,
                  threshold: float, on_high_risk) -> None:
    """Wrap openai.chat.completions.create (sync) and .acreate (async)."""
    # Support both the openai module itself and an OpenAI() instance
    try:
        chat_completions = client.chat.completions
    except AttributeError:
        return  # can't patch

    # Sync wrapper
    original_create = chat_completions.create

    def _wrapped_create(*args, **kwargs):
        resp = original_create(*args, **kwargs)
        try:
            messages = kwargs.get("messages", [])
            question = next((m.get("content", "") for m in messages if m.get("role") == "user"), "")
            choices  = getattr(resp, "choices", [])
            output   = choices[0].message.content if choices else ""
            steps    = [{"thought": output, "action_type": "Finish",
                         "action_arg": output, "observation": ""}]
            _fire_score(base_url, api_key, domain, question, steps,
                        on_high_risk, threshold, store_on=client)
        except Exception:
            pass
        return resp

    chat_completions.create = _wrapped_create
    client._lgk_last_score  = None
    client._lgk_instrumented = True


# ── Anthropic patch ───────────────────────────────────────────────────────────

def _patch_anthropic(client, base_url: str, api_key: str, domain: str,
                     threshold: float, on_high_risk) -> None:
    """Wrap anthropic.messages.create (sync)."""
    try:
        messages_obj = client.messages
    except AttributeError:
        return

    original_create = messages_obj.create

    def _wrapped_create(*args, **kwargs):
        resp = original_create(*args, **kwargs)
        try:
            msgs     = kwargs.get("messages", [])
            question = next((m.get("content", "") for m in msgs if m.get("role") == "user"), "")
            content  = getattr(resp, "content", [])
            output   = content[0].text if content else ""
            steps    = [{"thought": output, "action_type": "Finish",
                         "action_arg": output, "observation": ""}]
            _fire_score(base_url, api_key, domain, question, steps,
                        on_high_risk, threshold, store_on=client)
        except Exception:
            pass
        return resp

    messages_obj.create = _wrapped_create
    client._lgk_last_score   = None
    client._lgk_instrumented = True


# ── GuardContext ─────────────────────────────────────────────────────────────

class GuardContext:
    """
    Context manager that instruments clients for the duration of a block.

    Example:
        with GuardContext(api_key="lgk_...", domain="my-app") as guard:
            resp = openai.chat.completions.create(...)
        print(guard.last_score)
    """

    def __init__(
        self,
        openai_client=None,
        anthropic_client=None,
        base_url: str = "",
        api_key: str = "",
        domain: str = "default",
        score_threshold: float = 0.65,
    ):
        self._oc     = openai_client
        self._ac     = anthropic_client
        self.base_url = base_url or os.environ.get("LGK_BASE_URL", "http://localhost:8000")
        self.api_key  = api_key  or os.environ.get("LGK_API_KEY", "")
        self.domain   = domain   or os.environ.get("LGK_DOMAIN", "default")
        self.threshold = score_threshold
        self.last_score: dict | None = None
        self._scores: list = []

    def _on_score(self, score_dict: dict) -> None:
        self._scores.append(score_dict)
        self.last_score = score_dict

    def __enter__(self):
        instrument(
            openai_client=self._oc,
            anthropic_client=self._ac,
            base_url=self.base_url,
            api_key=self.api_key,
            domain=self.domain,
            score_threshold=self.threshold,
            on_high_risk=self._on_score,
        )
        return self

    def __exit__(self, *args):
        # Restore: nothing to undo since we fire-and-forget; user re-instruments as needed
        pass

    @property
    def scores(self) -> list[dict]:
        return list(self._scores)


# ── Convenience: generate SDK snippet for the UI ─────────────────────────────

def generate_snippet(api_key: str, domain: str = "default",
                     base_url: str = "https://api.llm-guard-kit.com",
                     language: str = "python") -> str:
    """Return a copy-paste SDK snippet for the given API key and domain."""
    if language == "python":
        return f"""# llm-guard-kit auto-instrumentation
from qppg_service.sdk import instrument
import openai

instrument(
    openai_client=openai,
    api_key="{api_key}",
    domain="{domain}",
    base_url="{base_url}",
)
# All openai.chat.completions.create() calls are now auto-scored.
"""
    elif language == "javascript":
        return f"""// llm-guard-kit proxy — route via our gateway, zero code changes
const openai = new OpenAI({{
  apiKey: process.env.OPENAI_API_KEY,
  baseURL: "{base_url}/proxy/openai/v1",
  defaultHeaders: {{ Authorization: "Bearer {api_key}" }},
}});
// All calls are now intercepted, scored, and logged.
"""
    return f"# SDK not available for language: {language}"
