"""Claude Code session bridge — runs prompts via `claude -p` or Anthropic SDK.

History: originally shelled out to `claude -p` exclusively (zero direct API
cost via the developer's claude.ai session). That works locally but fails
in the legal-ai Docker container, which does not ship the CLI. To keep the
same call sites working in production, the module now tries the CLI first
and falls back to the Anthropic SDK using ``ANTHROPIC_API_KEY`` when the
CLI binary is absent.

Both paths share a single shape: ``query()`` returns text, ``query_json()``
parses that text as JSON. Callers don't need to know which path executed.

Async history: originally synchronous (``subprocess.run``) with a 120 s
timeout. That broke for large legal documents — sync subprocess stalled
the asyncio loop, and 120 s was far too short for cold-cache Hebrew prompts
(case 8174-24 hit three timeouts in a row). Fixed by going async with a
30-minute ceiling.

Caching contract (SDK path): pass long, repeated instruction text via
``system=...``. The SDK path adds ``cache_control: ephemeral`` so back-to-back
chunk calls reuse the cached prefix at ~10% of read cost. The CLI path doesn't
expose API-level caching; with ``system`` set, we just prepend it to the
prompt — same observable behavior, no caching benefit.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil

from legal_mcp.config import parse_llm_json

logger = logging.getLogger(__name__)

# Default ceiling for any single LLM call, in seconds.
# 30 min covers any single-document call we make in practice (chunking
# handles the rest); the bound exists only to prevent runaway zombies.
DEFAULT_TIMEOUT = 1800
LONG_TIMEOUT = 3600  # opus block writing on full case context

# Anthropic SDK fallback config — used when `claude` CLI is not on PATH.
# Default to Sonnet 4.6: strong balance of Hebrew legal-text quality and
# cost for the per-chunk extraction workload. Override via env if needed.
DEFAULT_SDK_MODEL = os.environ.get("CLAUDE_SDK_MODEL", "claude-sonnet-4-6")
DEFAULT_SDK_MAX_TOKENS = int(os.environ.get("CLAUDE_SDK_MAX_TOKENS", "8192"))

_anthropic_client = None


def _has_cli() -> bool:
    return shutil.which("claude") is not None


def _get_anthropic_client():
    """Lazy-init the AsyncAnthropic client. Raises with a clear message if
    the package or API key is missing — better than letting the SDK 401 in
    the middle of a multi-chunk extraction.
    """
    global _anthropic_client
    if _anthropic_client is not None:
        return _anthropic_client
    try:
        import anthropic
    except ImportError as e:
        raise RuntimeError(
            "The 'anthropic' package is required when the Claude CLI is "
            "unavailable. Add it to mcp-server/pyproject.toml."
        ) from e
    if not os.environ.get("ANTHROPIC_API_KEY"):
        raise RuntimeError(
            "ANTHROPIC_API_KEY is not set; cannot fall back to Anthropic SDK."
        )
    _anthropic_client = anthropic.AsyncAnthropic()
    return _anthropic_client


async def _query_cli(
    prompt: str, system: str | None, timeout: int, max_turns: int,
) -> str:
    """Run the prompt via the local `claude` CLI subprocess.

    Uses the developer's claude.ai session — zero direct API cost. With
    ``system`` set, we just prepend it to the prompt; the CLI doesn't
    expose API-level caching anyway.
    """
    full_prompt = f"{system}\n\n{prompt}" if system else prompt

    cmd = [
        "claude", "-p",
        "--output-format", "json",
        "--max-turns", str(max_turns),
    ]

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    try:
        stdout_b, stderr_b = await asyncio.wait_for(
            proc.communicate(input=full_prompt.encode("utf-8")),
            timeout=timeout,
        )
    except asyncio.TimeoutError:
        # wait_for cancellation alone leaves the child running.
        try:
            proc.kill()
            await proc.wait()
        except ProcessLookupError:
            pass
        raise RuntimeError(f"Claude CLI timed out after {timeout}s")

    if proc.returncode != 0:
        stderr = stderr_b.decode("utf-8", errors="replace").strip()[:500] or "unknown error"
        raise RuntimeError(f"Claude CLI failed (exit {proc.returncode}): {stderr}")

    stdout = stdout_b.decode("utf-8", errors="replace").strip()
    if not stdout:
        raise RuntimeError("Claude CLI returned empty response")

    # claude -p --output-format json returns {"type":"result","result":"..."}
    try:
        data = json.loads(stdout)
        if isinstance(data, dict) and "result" in data:
            return data["result"]
        return stdout
    except json.JSONDecodeError:
        return stdout


async def _query_sdk(prompt: str, system: str | None, timeout: int) -> str:
    """Run the prompt via the Anthropic SDK with 5-min ephemeral caching
    on the system message. Streams the response to dodge HTTP read timeouts
    on long Hebrew JSON outputs.
    """
    import anthropic

    client = _get_anthropic_client()

    kwargs: dict = {
        "model": DEFAULT_SDK_MODEL,
        "max_tokens": DEFAULT_SDK_MAX_TOKENS,
        "messages": [{"role": "user", "content": prompt}],
    }
    if system:
        # cache_control: ephemeral → 5-min TTL. The same system text repeats
        # across every chunk in an extraction run, so we get
        # cache_read_input_tokens on every call after the first.
        kwargs["system"] = [{
            "type": "text",
            "text": system,
            "cache_control": {"type": "ephemeral"},
        }]

    try:
        async with client.messages.stream(**kwargs) as stream:
            message = await asyncio.wait_for(
                stream.get_final_message(),
                timeout=timeout,
            )
    except asyncio.TimeoutError:
        raise RuntimeError(f"Anthropic SDK call timed out after {timeout}s")
    except anthropic.APIError as e:
        raise RuntimeError(f"Anthropic SDK call failed: {e}") from e

    text_parts: list[str] = []
    for block in message.content:
        if getattr(block, "type", None) == "text":
            text_parts.append(block.text)
    out = "".join(text_parts).strip()
    if not out:
        raise RuntimeError("Anthropic SDK returned no text content")

    usage = getattr(message, "usage", None)
    if usage is not None:
        logger.debug(
            "claude_session SDK usage: input=%s cache_read=%s cache_write=%s output=%s",
            getattr(usage, "input_tokens", None),
            getattr(usage, "cache_read_input_tokens", None),
            getattr(usage, "cache_creation_input_tokens", None),
            getattr(usage, "output_tokens", None),
        )
    return out


async def query(
    prompt: str,
    timeout: int = DEFAULT_TIMEOUT,
    max_turns: int = 1,
    *,
    system: str | None = None,
) -> str:
    """Send a prompt to Claude and return the text response.

    Tries the Claude CLI first (zero API cost, uses claude.ai session).
    Falls back to the Anthropic SDK with ANTHROPIC_API_KEY when the CLI is
    absent — this is the production-Docker path.

    Pass repeating instruction text via ``system=`` so the SDK path can
    cache it (5-min ephemeral). Pass only the per-chunk content via
    ``prompt``. The CLI path concatenates them; the SDK path keeps them
    separate so caching works.

    Args:
        prompt: The user-facing prompt text.
        timeout: Max seconds before the call is aborted.
        max_turns: CLI-only — max conversation turns (1 = single response).
        system: Optional system message. With the SDK path, gets cached
            with 5-min ephemeral TTL when set.

    Raises:
        RuntimeError: if both paths fail or time out. The message includes
            which path raised so the caller can distinguish CLI from SDK.
    """
    if _has_cli():
        return await _query_cli(prompt, system, timeout, max_turns)
    return await _query_sdk(prompt, system, timeout)


async def query_json(
    prompt: str,
    timeout: int = DEFAULT_TIMEOUT,
    *,
    system: str | None = None,
) -> dict | list | None:
    """Send a prompt and parse the response as JSON.

    Uses parse_llm_json for robust parsing (handles markdown wrapping, truncation).
    """
    raw = await query(prompt, timeout=timeout, system=system)
    return parse_llm_json(raw)
