# SPDX-License-Identifier: MIT
"""Run and verify the Langfuse integration without retaining raw traces."""

from __future__ import annotations

import argparse
from importlib.metadata import version
import json
import os
from pathlib import Path
import subprocess
import sys
import time
from typing import Dict, List, Mapping, Optional

from langfuse_adapter import LangfuseTraceSink
from workflow import FixtureTools, SupportRequest, WorkflowResult, run_support_workflow

RAW_EMAIL = "roman.test@example.com"
RAW_ACCOUNT = "acct_DEMO123"
FIXTURE_RELEASE = "6401879"
FIXTURE_ENVIRONMENT = "futurecraft-evidence"
SAFE_METADATA_KEYS = {
    "account_id",
    "attempt",
    "document",
    "email",
    "environment",
    "error_type",
    "fixture",
    "fixture_response",
    "logical_observation_id",
    "logical_parent_id",
    "needs_human_review",
    "question",
    "refund_window",
    "subscription",
}


def fixture_request() -> SupportRequest:
    return SupportRequest(
        question=f"Can {RAW_EMAIL} refund {RAW_ACCOUNT}?",
        email=RAW_EMAIL,
        account_id=RAW_ACCOUNT,
    )


def model_dict(value: object) -> Dict[str, object]:
    dump = getattr(value, "model_dump", None)
    if not callable(dump):
        raise TypeError(f"unsupported Langfuse response: {type(value).__name__}")
    result = dump(mode="json", by_alias=True)
    if not isinstance(result, dict):
        raise TypeError("Langfuse response did not serialize to an object")
    return result


def wait_for_trace(
    client: object, trace_id: str, minimum_observations: int
) -> Dict[str, object]:
    api = getattr(client, "api")
    last_error: Optional[Exception] = None
    for _ in range(20):
        try:
            trace = model_dict(api.trace.get(trace_id))
            observations = trace.get("observations")
            if (
                isinstance(observations, list)
                and len(observations) >= minimum_observations
            ):
                return trace
        except Exception as error:
            last_error = error
        time.sleep(1)
    raise RuntimeError(f"trace {trace_id} was not readable") from last_error


def assert_no_sensitive_values(
    payload: str, *, include_public_key: bool = True
) -> None:
    banned = {
        "synthetic-email": RAW_EMAIL,
        "synthetic-account": RAW_ACCOUNT,
        "langfuse-secret-key": os.environ.get("LANGFUSE_SECRET_KEY"),
        "langfuse-public-key": os.environ.get("LANGFUSE_PUBLIC_KEY"),
    }
    for label, value in banned.items():
        if not value:
            continue
        if label == "langfuse-public-key" and not include_public_key:
            continue
        if value in payload:
            raise AssertionError(
                f"trace export contains {label} at {_sensitive_path(payload, value)}"
            )


def _sensitive_path(payload: str, value: str) -> str:
    try:
        parsed = json.loads(payload)
    except json.JSONDecodeError:
        return "unparsed-payload"

    def walk(item: object, path: str) -> Optional[str]:
        if isinstance(item, dict):
            for key, child in item.items():
                found = walk(child, f"{path}.{key}")
                if found:
                    return found
        elif isinstance(item, list):
            for index, child in enumerate(item):
                found = walk(child, f"{path}[{index}]")
                if found:
                    return found
        elif value in str(item):
            return path
        return None

    return walk(parsed, "$") or "unknown"


def safe_observations(trace: Mapping[str, object]) -> List[Dict[str, object]]:
    raw_observations = trace.get("observations")
    if not isinstance(raw_observations, list):
        raise AssertionError("trace has no observation list")

    observations: List[Dict[str, object]] = []
    for raw in raw_observations:
        if not isinstance(raw, dict):
            raise AssertionError("observation is not an object")
        if raw.get("input") is not None or raw.get("output") is not None:
            raise AssertionError("adapter exported an input or output payload")
        metadata = raw.get("metadata")
        if not isinstance(metadata, dict):
            raise AssertionError("observation has no metadata object")
        safe_metadata = {
            key: metadata[key] for key in SAFE_METADATA_KEYS if key in metadata
        }
        observations.append(
            {
                "id": raw.get("id"),
                "type": raw.get("type"),
                "name": raw.get("name"),
                "parentObservationId": raw.get("parentObservationId"),
                "level": raw.get("level"),
                "statusMessage": raw.get("statusMessage"),
                "startTime": raw.get("startTime"),
                "endTime": raw.get("endTime"),
                "version": raw.get("version"),
                "environment": raw.get("environment"),
                "metadata": safe_metadata,
            }
        )
    return sorted(observations, key=lambda item: str(item["startTime"]))


def assert_shape(scenario: str, observations: List[Dict[str, object]]) -> None:
    by_logical_id = {
        item["metadata"]["logical_observation_id"]: item for item in observations
    }
    expected = {
        "task-1",
        "retrieval-1",
        "plan-1",
        "subscription-1",
        "refund-1",
        "answer-1",
    }
    if scenario == "retry":
        expected.add("subscription-2")
    if set(by_logical_id) != expected:
        raise AssertionError(
            f"unexpected {scenario} observation set: {sorted(by_logical_id)}"
        )

    root = by_logical_id["task-1"]
    if root["parentObservationId"] is not None:
        raise AssertionError("root observation has a parent")
    if any(
        item["parentObservationId"] != root["id"]
        for logical_id, item in by_logical_id.items()
        if logical_id != "task-1"
    ):
        raise AssertionError("child observation has incorrect parentage")
    if any(not item["endTime"] for item in observations):
        raise AssertionError("an observation has no end time")

    if scenario == "retry":
        if by_logical_id["subscription-1"]["level"] != "ERROR":
            raise AssertionError("first retry attempt is not an error")
        if by_logical_id["subscription-2"]["level"] == "ERROR":
            raise AssertionError("successful retry remains an error")
    if scenario == "terminal":
        if by_logical_id["refund-1"]["level"] != "ERROR":
            raise AssertionError("terminal tool error is missing")
        if by_logical_id["task-1"]["level"] != "WARNING":
            raise AssertionError("terminal scenario root is not a warning")
        for logical_id in ("retrieval-1", "plan-1", "answer-1"):
            if by_logical_id[logical_id]["level"] == "ERROR":
                raise AssertionError("successful sibling was marked as an error")


def run_scenario(client: object, scenario: str) -> Dict[str, object]:
    tools = FixtureTools(
        subscription_failures_remaining=1 if scenario == "retry" else 0,
        refund_terminal_error=scenario == "terminal",
    )
    sink = LangfuseTraceSink(client)
    result = run_support_workflow(fixture_request(), tools, sink)
    sink.flush()

    root = sink.observations["task-1"]
    minimum_observations = 7 if scenario == "retry" else 6
    trace = wait_for_trace(client, root.trace_id, minimum_observations)
    serialized = json.dumps(trace, sort_keys=True)
    assert_no_sensitive_values(serialized, include_public_key=False)
    observations = safe_observations(trace)
    assert_shape(scenario, observations)
    return {
        "scenario": scenario,
        "traceId": root.trace_id,
        "result": {
            "answer": result.answer,
            "needsHumanReview": result.needs_human_review,
        },
        "observations": observations,
    }


def run_unavailable_child() -> None:
    from langfuse import Langfuse

    client = Langfuse(
        public_key="pk-lf-fixture",
        secret_key="sk-lf-fixture",
        base_url="http://127.0.0.1:9",
        timeout=1,
        flush_at=1,
        flush_interval=0.1,
        environment=FIXTURE_ENVIRONMENT,
        release=FIXTURE_RELEASE,
    )
    result = run_support_workflow(
        fixture_request(),
        FixtureTools(),
        LangfuseTraceSink(client),
    )
    print(
        json.dumps(
            {
                "answer": result.answer,
                "needsHumanReview": result.needs_human_review,
            },
            sort_keys=True,
        )
    )


def verify_unavailable_exporter() -> Dict[str, object]:
    environment = dict(os.environ)
    for key in ("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_BASE_URL"):
        environment.pop(key, None)
    completed = subprocess.run(
        [sys.executable, __file__, "--unavailable-child"],
        check=True,
        capture_output=True,
        text=True,
        env=environment,
        timeout=15,
    )
    result_line = completed.stdout.strip().splitlines()[-1]
    result = json.loads(result_line)
    expected = WorkflowResult(
        answer="Subscription: pro. Refund window: 14 days.",
        needs_human_review=False,
    )
    if result != {
        "answer": expected.answer,
        "needsHumanReview": expected.needs_human_review,
    }:
        raise AssertionError("unavailable exporter changed the application result")
    return {"applicationResultPreserved": True, "endpoint": "127.0.0.1:9"}


def run(output_path: Path) -> None:
    from langfuse import Langfuse

    base_url = os.environ.get("LANGFUSE_BASE_URL") or os.environ.get(
        "LANGFUSE_BASEURL"
    )
    if not base_url:
        raise RuntimeError("LANGFUSE_BASE_URL or LANGFUSE_BASEURL is required")
    client = Langfuse(
        public_key=os.environ.get("LANGFUSE_PUBLIC_KEY"),
        secret_key=os.environ.get("LANGFUSE_SECRET_KEY"),
        base_url=base_url,
        environment=FIXTURE_ENVIRONMENT,
        release=FIXTURE_RELEASE,
        flush_at=1,
    )
    health = client.api.health.health()
    scenarios = [
        run_scenario(client, scenario) for scenario in ("happy", "retry", "terminal")
    ]
    output = {
        "generatedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "fixtureRelease": FIXTURE_RELEASE,
        "environment": FIXTURE_ENVIRONMENT,
        "versions": {
            "langfusePython": version("langfuse"),
            "langfuseServer": health.version,
            "python": sys.version.split()[0],
        },
        "scenarios": scenarios,
        "unavailableExporter": verify_unavailable_exporter(),
    }
    serialized = json.dumps(output, indent=2, sort_keys=True) + "\n"
    assert_no_sensitive_values(serialized)
    output_path.write_text(serialized, encoding="utf-8")
    print(f"Verified {len(scenarios)} remote scenarios: {output_path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("/tmp/futurecraft-langfuse-evidence.json"),
    )
    parser.add_argument("--unavailable-child", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    if arguments.unavailable_child:
        run_unavailable_child()
    else:
        run(arguments.output)