"""SDK in-process transport driver and shared turn results."""

from __future__ import annotations

import asyncio
import json
import shutil
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx

from omnigent.runtime.harnesses.process_manager import HarnessProcessManager
from tests.e2e._harness_probes import cli_unavailable_reason
from tests.harness_bench.profile import BenchProfile
from tests.harness_bench.runtime_env import bench_creds_skip_reason, resolve_bench_env


class ProvisioningError(RuntimeError):
    """Expected environment failure that should skip one harness."""


POLICY_ALLOW = "POLICY_ACTION_ALLOW"
POLICY_DENY = "POLICY_ACTION_DENY"

# Denying another phase does not prove tool-call enforcement.
PHASE_TOOL_CALL = "PHASE_TOOL_CALL"

_CONV_ID = "conv_bench"

_STREAM_PROMPT = (
    "Count from 1 to 30 in words, one number per line, and add a short note after each."
)
_LONG_PROMPT = (
    "Write a very detailed 600-word essay about the history of computing, in full paragraphs."
)
_BENCH_TOOL_NAME = "bench_tool"
_BENCH_DENY_REASON = "bench-policy-deny"
_BENCH_TOOL_SPEC = [
    {
        "type": "function",
        "function": {
            "name": _BENCH_TOOL_NAME,
            "description": "A bench probe tool. Call it when asked.",
            "parameters": {
                "type": "object",
                "properties": {"arg": {"type": "string"}},
                "required": ["arg"],
            },
        },
    }
]

# Infrastructure failures must not be reported as capability gaps.
_INFRA_ERROR_MARKERS: tuple[str, ...] = (
    "403",
    "401",
    "Forbidden",
    "Unauthorized",
    "Invalid Token",
    "invalid token",
    "unexpected status",
    "Connection",
    "connection",
    "Temporarily Unavailable",
    "502",
    "503",
    "504",
    "already processing",
    "could not fetch a gateway token",
    "provider auth command",
    "empty token",
    "Failed to resolve external API key auth",
    "are logged in",
    "AcpProcessExited",
    "ACP subprocess",
    "ACP session",
)


def _error_text(error: object) -> str:
    if isinstance(error, dict):
        return f"{error.get('message', '')} {error.get('code', '')}"
    return str(error or "")


def infra_failure_reason(result: TurnResult) -> str | None:
    """Return a skip reason when failure reflects infrastructure, not capability."""
    if not result.failed:
        return None
    text = _error_text(result.error)
    if not any(marker in text for marker in _INFRA_ERROR_MARKERS):
        return None
    for code in ("403", "401"):
        if code in text:
            return (
                f"auth rejected ({code} Invalid/Forbidden token); the harness "
                "credential is stale or shadowed by an ambient env var. Refresh "
                "the harness auth source (profile, API key, or token env var)"
            )
    if "already processing" in text:
        return "session busy from a prior turn (sequencing, not a capability gap)"
    if any(
        marker in text
        for marker in (
            "could not fetch a gateway token",
            "provider auth command",
            "empty token",
            "Failed to resolve external API key auth",
        )
    ):
        return (
            "gateway/provider token could not be provisioned for this transport "
            "(environment/auth gap, not a capability the harness lacks)"
        )
    if any(
        marker in text
        for marker in ("are logged in", "AcpProcessExited", "ACP subprocess", "ACP session")
    ):
        return (
            "vendor CLI not installed or not logged in (own-auth harness); "
            "the agent process exited before a turn could run"
        )
    if "unexpected status" in text:
        return "gateway returned an unexpected status (environment/auth issue)"
    return "environment/connectivity error reaching the gateway"


@dataclass
class TurnResult:
    """Probe-observable state from one turn."""

    events: list[dict[str, Any]] = field(default_factory=list)
    text: str = ""
    text_delta_count: int = 0
    reasoning_delta_count: int = 0
    tool_calls: list[dict[str, Any]] = field(default_factory=list)
    policy_actions: list[tuple[str, str]] = field(default_factory=list)
    tool_call_denied: bool = False
    completed: bool = False
    cancelled: bool = False
    failed: bool = False
    error: Any = None
    timed_out: bool = False
    total_tokens: int | None = None
    total_cost_usd: float | None = None
    elicitation_requested: bool = False
    tool_call_allowed: bool = False

    @property
    def reached_terminal(self) -> bool:
        return self.completed or self.cancelled or self.failed

    @property
    def event_types(self) -> list[str]:
        return [e.get("type", "") for e in self.events]


def fill_snapshot_cost(result: TurnResult, snapshot: dict[str, Any]) -> None:
    """Copy observed usage and cost from a session snapshot."""
    tokens = snapshot.get("last_total_tokens")
    if isinstance(tokens, int):
        result.total_tokens = tokens
    cost = snapshot.get("total_cost_usd")
    if isinstance(cost, (int, float)):
        result.total_cost_usd = float(cost)


class SdkInprocDriver:
    """Drive turns through a harness wrap subprocess."""

    transport = "sdk-inproc"

    def __init__(self, profile: BenchProfile, *, databricks_profile: str | None) -> None:
        self._profile = profile
        self._databricks_profile = databricks_profile
        self._pm: HarnessProcessManager | None = None
        self._client: httpx.AsyncClient | None = None
        self._tmp_parent: Path | None = None

    @staticmethod
    def unavailable(profile: BenchProfile, *, databricks_profile: str | None) -> str | None:
        """Return why this driver cannot run the profile, if applicable."""
        if profile.transport != SdkInprocDriver.transport:
            return (
                f"transport {profile.transport!r} not supported by the "
                f"{SdkInprocDriver.transport!r} driver"
            )
        creds_skip = bench_creds_skip_reason(databricks_profile)
        if creds_skip is not None:
            return creds_skip
        if profile.cli_binary is not None:
            reason = cli_unavailable_reason(profile.cli_binary)
            if reason is not None:
                return reason
        return None

    async def __aenter__(self) -> SdkInprocDriver:
        self._tmp_parent = Path("/tmp") / f"omni-bench-{uuid.uuid4().hex[:8]}"
        self._tmp_parent.mkdir(mode=0o700)
        self._pm = HarnessProcessManager(tmp_parent=self._tmp_parent)
        await self._pm.start()
        p = self._profile
        resolved = resolve_bench_env(self._databricks_profile)
        wrap_env = {
            f"{p.env_prefix}GATEWAY": "true",
            f"{p.env_prefix}MODEL": p.model,
        }
        if resolved.db_profile:
            wrap_env[f"{p.env_prefix}DATABRICKS_PROFILE"] = resolved.db_profile
        self._client = await self._pm.get_client(_CONV_ID, p.harness, env=wrap_env)
        return self

    async def __aexit__(self, *exc: object) -> None:
        if self._pm is not None:
            await self._pm.shutdown()
        if self._tmp_parent is not None:
            shutil.rmtree(self._tmp_parent, ignore_errors=True)

    async def run_turn(
        self,
        prompt: str,
        *,
        tools: list[dict[str, Any]] | None = None,
        deny_phases: frozenset[str] = frozenset(),
        policy_reason: str | None = None,
        auto_tool_output: str | None = None,
        interrupt_on_first_delta: bool = False,
        timeout: float = 120.0,
    ) -> TurnResult:
        """Start one turn and drain its event stream."""
        assert self._client is not None, "driver used outside its async context"
        body: dict[str, Any] = {
            "type": "message",
            "role": "user",
            "model": f"{self._profile.harness}-bench-agent",
            "content": [{"type": "input_text", "text": prompt}],
        }
        if tools is not None:
            body["tools"] = tools

        result = TurnResult()
        try:
            await asyncio.wait_for(
                self._drive(
                    body,
                    result,
                    deny_phases,
                    policy_reason,
                    auto_tool_output,
                    interrupt_on_first_delta,
                ),
                timeout=timeout,
            )
        except (asyncio.TimeoutError, httpx.ReadTimeout):
            result.timed_out = True
        return result

    async def run_basic_turn(self, marker: str) -> TurnResult:
        return await self.run_turn(
            f"Reply with exactly the literal string {marker} and nothing else."
        )

    async def run_streaming_turn(self) -> TurnResult:
        return await self.run_turn(_STREAM_PROMPT)

    async def run_tool_turn(self, *, deny: bool) -> TurnResult:
        """Provoke a tool call and optionally deny its policy evaluation."""
        if deny:
            return await self.run_turn(
                f"Call the {_BENCH_TOOL_NAME} tool with arg='go'. It is required.",
                tools=_BENCH_TOOL_SPEC,
                deny_phases=frozenset({PHASE_TOOL_CALL}),
                policy_reason=_BENCH_DENY_REASON,
                timeout=150.0,
            )
        return await self.run_turn(
            f"You must call the {_BENCH_TOOL_NAME} tool with arg='go', "
            "then reply with the tool's output verbatim.",
            tools=_BENCH_TOOL_SPEC,
            auto_tool_output="bench-tool-ok",
            timeout=150.0,
        )

    async def run_policy_turn(self, *, action: str) -> TurnResult:
        """Return unmeasured because wrap-direct cannot observe ALLOW or ASK."""
        return TurnResult()

    async def run_interrupt_turn(self) -> TurnResult:
        return await self.run_turn(_LONG_PROMPT, interrupt_on_first_delta=True, timeout=120.0)

    async def _drive(
        self,
        body: dict[str, Any],
        result: TurnResult,
        deny_phases: frozenset[str],
        policy_reason: str | None,
        auto_tool_output: str | None,
        interrupt_on_first_delta: bool,
    ) -> None:
        """POST the turn and consume the SSE stream into *result* in place."""
        client = self._client
        assert client is not None
        interrupted = False
        async with client.stream("POST", f"/v1/sessions/{_CONV_ID}/events", json=body) as response:
            response.raise_for_status()
            buffer = ""
            async for chunk in response.aiter_text():
                buffer += chunk
                while "\n\n" in buffer:
                    frame, _, buffer = buffer.partition("\n\n")
                    event = _decode_frame(frame)
                    if event is None:
                        continue
                    result.events.append(event)
                    etype = event.get("type", "")

                    if etype == "response.output_text.delta":
                        result.text += event.get("delta", "")
                        result.text_delta_count += 1
                        if interrupt_on_first_delta and not interrupted:
                            interrupted = True
                            await self._post({"type": "interrupt"})
                    elif etype in _REASONING_DELTA_TYPES:
                        result.reasoning_delta_count += 1
                    elif etype == "response.output_item.done":
                        # Action-required calls park until a tool result arrives.
                        item = event.get("item") or {}
                        if (
                            item.get("type") == "function_call"
                            and item.get("status") == "action_required"
                        ):
                            call_id = item.get("call_id", "")
                            result.tool_calls.append(item)
                            if auto_tool_output is not None:
                                await self._post(
                                    {
                                        "type": "tool_result",
                                        "call_id": call_id,
                                        "output": auto_tool_output,
                                    }
                                )
                    elif etype == "policy_evaluation.requested":
                        phase = str(event.get("phase", ""))
                        action = POLICY_DENY if phase in deny_phases else POLICY_ALLOW
                        verdict: dict[str, Any] = {
                            "type": "policy_verdict",
                            "evaluation_id": event["evaluation_id"],
                            "action": action,
                        }
                        if action == POLICY_DENY and policy_reason is not None:
                            verdict["reason"] = policy_reason
                        # A raced or rejected verdict was not delivered.
                        if await self._post(verdict):
                            result.policy_actions.append((phase, action))
                            if action == POLICY_DENY and phase == PHASE_TOOL_CALL:
                                result.tool_call_denied = True
                    elif etype == "response.completed":
                        result.completed = True
                        usage = (event.get("response") or {}).get("usage") or {}
                        tok = usage.get("total_tokens")
                        if isinstance(tok, int):
                            result.total_tokens = tok
                        cost = usage.get("cost_usd")
                        if isinstance(cost, (int, float)):
                            result.total_cost_usd = float(cost)
                    elif etype == "response.cancelled":
                        result.cancelled = True
                    elif etype == "response.failed":
                        result.failed = True
                        result.error = event.get("error") or event.get("response", {}).get("error")

    async def _post(self, payload: dict[str, Any]) -> bool:
        """Post a downward event, tolerating races with turn completion."""
        assert self._client is not None
        try:
            resp = await self._client.post(f"/v1/sessions/{_CONV_ID}/events", json=payload)
        except httpx.HTTPError:
            return False
        return not resp.is_error


# Harness wraps use both reasoning-delta spellings.
_REASONING_DELTA_TYPES: frozenset[str] = frozenset(
    {"response.reasoning.delta", "response.reasoning_summary_text.delta"}
)


def _decode_frame(frame: str) -> dict[str, Any] | None:
    data_line = next(
        (line for line in frame.splitlines() if line.startswith("data:")),
        None,
    )
    if data_line is None:
        return None
    try:
        decoded = json.loads(data_line[len("data:") :].strip())
    except json.JSONDecodeError:
        return None
    return decoded if isinstance(decoded, dict) else None