"""
lambdagent.conversation — ConversationLam: conversation-aware Lambda abstraction.

Wraps any LLMProvider with conversation history management.
Each apply() appends the new input as a user message, calls the provider
with the full (or windowed) history, and records the response.

This is what eliminates hallucination: the LLM sees its complete
conversation history, not a lossy compressed state string.

L03: Added ChatMessage-based apply_typed() path for richer response metadata.
ConversationLam now supports both dict-based (legacy) and ChatMessage-based
(unified) conversation flows.

Usage:
    provider = ClaudeCodeProvider(config)
    lam = ConversationLam("agent", provider, system_prompt="...")

    r1 = lam.apply("read the README")     # creates conversation
    r2 = lam.apply("[tool result] ...")     # continues with full memory
    r3 = lam.apply("[tool result] ...")     # still remembers r1, r2
"""
from __future__ import annotations

import logging
import time
from typing import Any, Callable, List, Optional

from lambdagent.core import Term, Context
from lambdagent.providers.base import LLMProvider, ProviderError, ChatMessage, ChatResponse

logger = logging.getLogger("lambdagent.conversation")

# 可重试 provider 错误（如 claude-code resume 首字节卡死 / 瞬时 API 抖动）的总尝试次数。
# 不重试则一次卡死就把这步退化成 [X_ERROR] 字符串、run 失败。provider 在 stall 时已
# 清 session_id → 重试时退化成 fresh first-turn(小 prompt，通常很快)→ 恢复而非死。
_PROVIDER_RETRY_ATTEMPTS = 5


def _provider_call_with_retry(fn, what: str):
    """调用 provider，对 retryable ProviderError 重试。非 retryable 立即上抛；
    重试耗尽后把最后一次错误上抛，由调用方退化成 [X_ERROR] 字符串。"""
    last: Optional[ProviderError] = None
    for attempt in range(1, _PROVIDER_RETRY_ATTEMPTS + 1):
        try:
            return fn()
        except ProviderError as e:
            last = e
            if not getattr(e, "retryable", False) or attempt >= _PROVIDER_RETRY_ATTEMPTS:
                raise
            logger.warning(
                "provider retryable error on %s (attempt %d/%d) — 重试: %s",
                what, attempt, _PROVIDER_RETRY_ATTEMPTS, str(e)[:200],
            )
            time.sleep(1.0)
    assert last is not None
    raise last


class ConversationLam(Term):
    """
    Lambda abstraction with conversation persistence.

    Lambda semantics preserved:
        ConversationLam(provider, prompt) = lambda x. provider(history + x)
        apply() = beta-reduction with memory

    The key difference from stateless Lam:
        Lam:             each apply() is independent
        ConversationLam: each apply() builds on all previous calls

    L03: Now supports both legacy dict-based and ChatMessage-based flows.
    """

    def __init__(
        self,
        name: str,
        provider: LLMProvider,
        system_prompt: str,
        max_history_tokens: int = 80000,
        keep_recent_turns: int = 20,
        output_parser: Callable[[str], Any] | None = None,
        # L03: New unified interface params
        model: str = "",
        temperature: float = 0.0,
        max_tokens: int = 4096,
        # Hard character ceiling for the *rendered* request (system + history +
        # new input). Providers like DashScope/qwen reject inputs over a fixed
        # CHARACTER count (qwen: 258048 → "Range of input length should be
        # [1, 258048]"). The token-based budget above is not enough on its own:
        # for CJK text 1 token ≈ 1 char, so an 80k-"token" history estimated as
        # 320k chars sails past the char limit. We keep a conservative margin
        # below 258048 for the system prompt + the next user turn.
        max_input_chars: int = 200000,
    ):
        super().__init__(name)
        self.provider = provider
        self.system_prompt = system_prompt
        self.max_history_tokens = max_history_tokens
        self.keep_recent_turns = keep_recent_turns
        self.max_input_chars = max_input_chars
        self.output_parser = output_parser or (lambda x: x)
        # L03: Store model/temperature/max_tokens for chat_typed calls
        self._model = model
        self._temperature = temperature
        self._max_tokens = max_tokens

        # Conversation history (system message always first)
        self.messages: List[dict] = [
            {"role": "system", "content": system_prompt}
        ]

        # L03: ChatMessage-based history (parallel to dict-based for typed path)
        self._typed_messages: List[ChatMessage] = []

    # Expose model name for react_step logging
    @property
    def model(self) -> str:
        return self._model or self.provider.model_name

    # Expose _session_id for react_step session detection
    @property
    def _session_id(self):
        return getattr(self.provider, '_session_id', None)

    def apply(self, input: Any, ctx: Context | None = None) -> Any:
        """Beta-reduction with conversation memory."""
        ctx = ctx or Context()
        t0 = time.time()

        # Append user message
        self.messages.append({"role": "user", "content": str(input)})

        # Manage context window before sending
        managed = self._manage_context()

        # Call provider
        tokens_used = 0
        try:
            response = _provider_call_with_retry(
                lambda: self.provider.chat(managed), "chat")
        except ProviderError as e:
            response = f"[{e.provider.upper()}_ERROR] {e}"
            ctx.record_provider_error(e.provider, str(e))  # AUDIT ②: 结构化上报

        # Record assistant response
        self.messages.append({"role": "assistant", "content": response})

        duration = (time.time() - t0) * 1000
        result = self.output_parser(response)
        ctx.log(self._name, self._trace_id, input, result, duration, self.model, tokens_used)
        return result

    # ── L03: ChatMessage-based typed path ──

    def apply_typed(self, input: Any, ctx: Context | None = None) -> ChatResponse:
        """
        Apply with ChatMessage/ChatResponse for richer metadata.

        Uses provider.chat_typed() instead of provider.chat().
        Returns ChatResponse with token counts, model info, etc.
        """
        ctx = ctx or Context()
        t0 = time.time()

        # Build messages
        messages = [ChatMessage(role="system", content=self.system_prompt)]
        messages.extend(self._typed_messages)
        messages.append(ChatMessage(role="user", content=str(input)))

        # Compact if over budget
        messages = self._compact_typed(messages)

        # Call provider via typed interface
        try:
            response = _provider_call_with_retry(
                lambda: self.provider.chat_typed(
                    messages=messages,
                    model=self._model,
                    temperature=self._temperature,
                    max_tokens=self._max_tokens,
                ), "chat_typed")
        except ProviderError as e:
            response = ChatResponse(
                text=f"[{e.provider.upper()}_ERROR] {e}",
                model=self.model,
            )
            ctx.record_provider_error(e.provider, str(e))  # AUDIT ②: 结构化上报

        duration = (time.time() - t0) * 1000

        # Update typed history
        self._typed_messages.append(ChatMessage(role="user", content=str(input)))
        self._typed_messages.append(ChatMessage(role="assistant", content=response.text))

        # Also sync to dict-based history for compatibility
        self.messages.append({"role": "user", "content": str(input)})
        self.messages.append({"role": "assistant", "content": response.text})

        # Log to context
        ctx.log(self._name, self._trace_id, str(input)[:200],
                response.text[:200], duration, self.model,
                response.input_tokens + response.output_tokens)

        return response

    def _compact_typed(self, messages: List[ChatMessage]) -> List[ChatMessage]:
        """Compact ChatMessage history under both token and char budgets."""
        total_chars = sum(len(m.content) for m in messages)
        estimated_tokens = total_chars // 4

        # Hard character ceiling (qwen rejects >258048 chars). Drop oldest
        # history first, then middle-truncate a remaining oversized message.
        if total_chars > self.max_input_chars and messages:
            system = messages[:1]
            history = list(messages[1:])
            sys_chars = sum(len(m.content) for m in system)
            while history and sys_chars + sum(len(m.content) for m in history) > self.max_input_chars:
                history.pop(0)
            kept = system + history
            if sum(len(m.content) for m in kept) > self.max_input_chars:
                slack = max(1000, self.max_input_chars - sys_chars - 200)
                for i, m in enumerate(kept[1:], 1):
                    if len(m.content) > slack:
                        c = m.content
                        head = slack // 2
                        kept[i] = ChatMessage(
                            role=m.role,
                            content=c[:head]
                            + f"\n... [truncated {len(c) - slack} chars to fit input limit] ...\n"
                            + c[-(slack - head):],
                        )
                    if sum(len(mm.content) for mm in kept) <= self.max_input_chars:
                        break
            messages = kept
            total_chars = sum(len(m.content) for m in messages)
            estimated_tokens = total_chars // 4

        if estimated_tokens <= self.max_history_tokens:
            return messages

        # Keep system + last N messages that fit the token budget
        system = messages[:1]
        history = messages[1:]

        for keep in range(len(history), 0, -1):
            kept = system + history[-keep:]
            est = sum(len(m.content) for m in kept) // 4
            if est <= self.max_history_tokens:
                return kept

        return system  # Worst case: only system prompt

    @property
    def history(self) -> List[ChatMessage]:
        """Return typed message history."""
        return list(self._typed_messages)

    # ── Legacy dict-based context management ──

    def _manage_context(self) -> List[dict]:
        """
        Ensure messages fit within max_history_tokens.

        Strategy:
          - Always keep system message
          - Always keep recent N turns (user+assistant pairs)
          - Summarize older messages into a single "history summary" message
        """
        total_tokens = self._estimate_tokens(self.messages)

        if total_tokens <= self.max_history_tokens:
            # Under the token budget, but CJK-heavy histories can still exceed
            # the provider's hard character limit — enforce that too.
            return self._enforce_char_budget(list(self.messages))

        # Keep system + recent turns
        system = self.messages[0]
        # Each turn = 1 user + 1 assistant = 2 messages
        keep_count = self.keep_recent_turns * 2
        recent = self.messages[-keep_count:] if len(self.messages) > keep_count else self.messages[1:]
        old = self.messages[1:-keep_count] if len(self.messages) > keep_count + 1 else []

        if not old:
            # All messages are within the recent window — but they may still
            # blow the provider's char limit (e.g. several full PDF reads).
            return self._enforce_char_budget(list(self.messages))

        # Compress old messages into summary
        summary_lines = []
        for m in old:
            role = m["role"]
            content = m["content"][:150]
            summary_lines.append(f"[{role}] {content}...")

        summary = "[对话历史摘要]\n" + "\n".join(summary_lines)

        return self._enforce_char_budget(
            [system, {"role": "user", "content": summary}] + list(recent)
        )

    def _enforce_char_budget(self, messages: List[dict]) -> List[dict]:
        """Hard cap the rendered request size by CHARACTER count.

        Token estimates undercount CJK text, so a token-only budget can still
        exceed a provider's character limit (qwen: 258048). This guarantees the
        final list fits: it (1) drops the OLDEST non-system turns first, then
        (2) middle-truncates any single oversized message that remains.
        """
        budget = self.max_input_chars
        if not messages:
            return messages

        def total(msgs):
            return sum(len(m.get("content", "")) for m in msgs)

        if total(messages) <= budget:
            return messages

        system = messages[:1]
        history = messages[1:]
        sys_chars = total(system)

        # 1) Drop oldest history messages until we're under budget.
        while history and sys_chars + total(history) > budget:
            history.pop(0)

        kept = system + history

        # 2) If a single remaining message is still too big (e.g. one giant
        #    tool observation), middle-truncate it so the head+tail survive.
        if total(kept) > budget:
            slack = max(1000, budget - sys_chars - 200)
            for m in kept[1:]:
                if len(m.get("content", "")) > slack:
                    c = m["content"]
                    head = slack // 2
                    tail = slack - head
                    m["content"] = (
                        c[:head]
                        + f"\n... [truncated {len(c) - slack} chars to fit input limit] ...\n"
                        + c[-tail:]
                    )
                if total(kept) <= budget:
                    break

        return kept

    def _estimate_tokens(self, messages: List[dict]) -> int:
        """Rough token estimate: 4 chars ≈ 1 token."""
        return sum(len(m.get("content", "")) for m in messages) // 4

    def get_usage(self) -> dict:
        """Get cumulative token usage from the underlying provider's session.

        Delegates to provider.get_usage() (e.g. ClaudeCodeProvider reads the
        JSONL transcript). Returns zero dict if the provider does not support it.
        """
        empty = {
            "input_tokens": 0,
            "output_tokens": 0,
            "cache_creation_input_tokens": 0,
            "cache_read_input_tokens": 0,
        }
        fn = getattr(self.provider, "get_usage", None)
        if callable(fn):
            return fn()
        return empty

    def reset(self):
        """Clear conversation history, start fresh."""
        self.messages = [{"role": "system", "content": self.system_prompt}]
        self._typed_messages = []
        if hasattr(self.provider, 'reset_session'):
            self.provider.reset_session()

    def __rshift__(self, other):
        from lambdagent.primitives import Compose
        return Compose(self, other)
