"""Local-only SDK payload probes. All credentials and content are synthetic."""
import gzip
import json
import os
import warnings
import threading
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from importlib.metadata import version

# Isolate probes from inherited telemetry destinations and defaults.
for key in list(os.environ):
    if key.startswith(("LANGSMITH_", "LANGCHAIN_", "LANGFUSE_", "OTEL_")):
        del os.environ[key]

from langsmith import Client, traceable, tracing_context
from langsmith.anonymizer import create_secret_anonymizer
from langfuse import Langfuse
from langfuse.types import MaskOtelSpansResult, OtelSpanPatch
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
from google.protobuf.json_format import MessageToDict

captures = []

class Receiver(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        body = json.dumps({"version": "local-test", "batch_ingest_config": {"use_multipart_endpoint": False}}).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def receive(self):
        body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
        if self.headers.get("Content-Encoding") == "gzip":
            body = gzip.decompress(body)
        if "protobuf" in self.headers.get("Content-Type", ""):
            msg = ExportTraceServiceRequest()
            msg.ParseFromString(body)
            payload = MessageToDict(msg, preserving_proto_field_name=True)
        else:
            payload = json.loads(body)
        captures.append({"path": self.path, "payload": payload})
        self.send_response(200)
        self.send_header("Content-Length", "2")
        self.end_headers()
        self.wfile.write(b"{}")

    do_POST = receive
    do_PATCH = receive

server = ThreadingHTTPServer(("127.0.0.1", 0), Receiver)
threading.Thread(target=server.serve_forever, daemon=True).start()
base = f"http://127.0.0.1:{server.server_port}"
results = []

def check(label, condition):
    assert condition, label
    results.append({"check": label, "pass": True})

def smith(**options):
    return Client(api_url=base, api_key="dummy-trace-key", auto_batch_tracing=False,
                  omit_traced_runtime_info=True, **options)

raw = "T28_FAKE_INPUT_7f3a"
client = smith(hide_inputs=True, hide_outputs=True)
run_id = uuid.uuid4()
client.create_run("safe-name", inputs={"prompt": raw}, run_type="chain", id=run_id,
                  extra={"metadata": {"note": "T28_FAKE_META_7f3a"}})
client.update_run(run_id, outputs={"answer": "T28_FAKE_OUTPUT_7f3a"})
p = json.dumps(captures)
check("LS hide_io removes input/output", raw not in p and "T28_FAKE_OUTPUT_7f3a" not in p)
check("LS hide_io leaves metadata", "T28_FAKE_META_7f3a" in p)

captures.clear()
client = smith(hide_inputs=True, hide_outputs=True, hide_metadata=True)
client.create_run("T28_FAKE_NAME_7f3a", inputs={"p": raw}, run_type="chain",
                  tags=["T28_FAKE_TAG_7f3a"], error="T28_FAKE_ERROR_7f3a",
                  extra={"metadata": {"note": "T28_FAKE_META_7f3a"}})
p = json.dumps(captures)
check("LS hide_metadata removes metadata", "T28_FAKE_META_7f3a" not in p)
check("LS hide flags leave name/tag/error", all(x in p for x in ["T28_FAKE_NAME_7f3a", "T28_FAKE_TAG_7f3a", "T28_FAKE_ERROR_7f3a"]))

captures.clear()
client = smith(anonymizer=create_secret_anonymizer())
secret = "ghp_" + "A" * 36
unknown = "T28_FAKE_UNRECOGNIZED_7f3a"
client.create_run("safe-name", inputs={"known": secret, "opaque": unknown}, run_type="chain",
                  error="request failed " + secret,
                  extra={"metadata": {"credential": secret}})
p = json.dumps(captures)
check("LS preset covers known secret and error", secret not in p and "[SECRET_DETECTED]" in p)
check("LS preset leaves unknown format", unknown in p)

captures.clear()
client = smith(hide_inputs=True, hide_outputs=True)
@traceable(client=client, process_inputs=lambda _: {"kept": "T28_FAKE_PROCESS_IN"},
           process_outputs=lambda _: {"kept": "T28_FAKE_PROCESS_OUT"})
def traced(value):
    return value
with tracing_context(enabled=True):
    check("LS application result unchanged", traced(raw) == raw)
client.flush()
p = json.dumps(captures)
check("LS client hide flags still apply after function processors", "T28_FAKE_PROCESS_IN" not in p and "T28_FAKE_PROCESS_OUT" not in p)

captures.clear()
client = smith(hide_inputs=lambda _: {"custom": "custom-hide"},
               anonymizer=lambda _: {"chosen": "anonymizer"})
client.create_run("priority-test", inputs={"p": raw}, run_type="chain")
p = json.dumps(captures)
check("LS Python anonymizer overrides callable hide_inputs", "anonymizer" in p and "custom-hide" not in p)

captures.clear()
client = smith()
@traceable(client=client)
def context_probe(value):
    return value
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    with tracing_context(enabled=True, hide_inputs=True, hide_outputs=True):
        context_probe(raw)
client.flush()
check("LS Python tracing_context hide kwargs unsupported in tested release",
      any("Unrecognized keyword" in str(w.message) for w in caught) and raw in json.dumps(captures))

def scrub(value):
    if isinstance(value, str):
        return value.replace("T28_FAKE_ATTR_7f3a", "[REDACTED]")
    if isinstance(value, dict):
        return {k: scrub(v) for k, v in value.items()}
    if isinstance(value, list):
        return [scrub(v) for v in value]
    if isinstance(value, tuple):
        return tuple(scrub(v) for v in value)
    return value

def export_mask(*, params):
    patches = {}
    for identifier, span in params.spans.items():
        changed = {k: scrub(v) for k, v in span.attributes.items() if scrub(v) != v}
        if changed:
            patches[identifier] = OtelSpanPatch(set_attributes=changed)
    return MaskOtelSpansResult(span_patches=patches)

def fuse_probe(label, **options):
    captures.clear()
    provider = TracerProvider()
    # A sibling processor/exporter is deliberately unmasked as a negative control.
    provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint=base + "/other")))
    lf = Langfuse(public_key="pk-dummy-" + label, secret_key="sk-dummy", base_url=base,
                  tracer_provider=provider, should_export_span=lambda _: True,
                  span_exporter=OTLPSpanExporter(endpoint=base + "/fuse"), **options)
    with lf.start_as_current_observation(name="sdk-observation", input="T28_FAKE_ATTR_7f3a",
                                        metadata={"note": "T28_FAKE_ATTR_7f3a"}):
        pass
    with provider.get_tracer("third-party-test").start_as_current_span("safe-third-party") as span:
        span.set_attribute("gen_ai.prompt.0.content", "T28_FAKE_ATTR_7f3a")
        span.set_attribute("custom.array", ["T28_FAKE_ATTR_7f3a", "safe"])
        span.add_event("safe-event", {"message": "T28_FAKE_EVENT_7f3a"})
    lf.flush()
    provider.force_flush()
    fuse = json.dumps([x for x in captures if x["path"] == "/fuse"])
    other = json.dumps([x for x in captures if x["path"] == "/other"])
    check(label + " positive capture", bool(fuse != "[]" and other != "[]"))
    if "mask_otel_spans" in options:
        check("LF export mask covers SDK/third-party/array attrs", "T28_FAKE_ATTR_7f3a" not in fuse)
        check("LF export mask leaves event attrs", "T28_FAKE_EVENT_7f3a" in fuse)
        check("LF sibling exporter remains unmasked", "T28_FAKE_ATTR_7f3a" in other)
    else:
        # SDK input and metadata are transformed, raw third-party attr is still sent.
        payloads = [x["payload"] for x in captures if x["path"] == "/fuse"]
        spans = [s for p in payloads for r in p["resource_spans"] for q in r["scope_spans"] for s in q["spans"]]
        sdk = json.dumps([s for s in spans if s["name"] == "sdk-observation"])
        third = json.dumps([s for s in spans if s["name"] == "safe-third-party"])
        check("LF legacy mask covers SDK fields", "T28_FAKE_ATTR_7f3a" not in sdk)
        check("LF legacy mask misses third-party attrs", "T28_FAKE_ATTR_7f3a" in third)
    provider.shutdown()
    lf.shutdown()

fuse_probe("legacy", mask=lambda *, data, **_: scrub(data))
fuse_probe("export", mask_otel_spans=export_mask)
server.shutdown()
print(json.dumps({"versions": {p: version(p) for p in ["langsmith", "langfuse", "opentelemetry-sdk"]},
                  "results": results, "passed": len(results)}, ensure_ascii=False, indent=2))