"""Opt-in live probe using synthetic input only; never prints keys or provider bodies.

Run from backend: .venv/Scripts/python scripts/probe_context_recovery.py --live --model NAME --rows 40000
"""

import argparse
import json
import logging
import re
import time

from langchain.agents.middleware.types import ModelRequest, ModelResponse
from langchain_core.messages import HumanMessage, message_chunk_to_message

from deerflow.agents.middlewares.context_window_guard_middleware import ContextWindowGuardMiddleware, estimate_model_request_tokens
from deerflow.agents.middlewares.llm_error_handling_middleware import is_llm_error_message
from deerflow.config.app_config import get_app_config
from deerflow.models.context_recovery import CONTEXT_RECOVERY_MESSAGE, is_context_length_error
from deerflow.models.factory import create_chat_model


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--live", action="store_true", required=True)
    parser.add_argument("--model", required=True)
    parser.add_argument("--rows", type=int, default=100)
    parser.add_argument("--stream", action="store_true")
    args = parser.parse_args()
    if not 0 <= args.rows <= 100000:
        parser.error("rows must be between 0 and 100000")
    logging.disable(logging.CRITICAL)
    started = time.monotonic()
    attempts = []
    report = {"model": args.model, "rows": args.rows, "stream": args.stream, "attempts": attempts}
    try:
        config = get_app_config().model_copy(deep=True)
        selected = config.get_model_config(args.model)
        selected.max_retries = 0
        selected.request_timeout = 90
        selected.max_tokens = 1024
        model = create_chat_model(name=args.model, app_config=config, thinking_enabled=False, force_disable_thinking=True)
        text = "以下是合成测试材料。只回复：验证通过。\n" + "".join(f"记录{i:06d}：甲乙丙丁，气温正常。\n" for i in range(args.rows)) + "\n只回复：验证通过。"
        request = ModelRequest(model=model, messages=[HumanMessage(content=text)])

        def invoke(current):
            attempt = {"estimated_input": estimate_model_request_tokens(current)}
            attempts.append(attempt)
            try:
                messages = ([current.system_message] if current.system_message else []) + current.messages
                if args.stream:
                    chunks = list(model.stream(messages, **current.model_settings))
                    response = message_chunk_to_message(sum(chunks[1:], chunks[0]))
                else:
                    response = model.invoke(messages, **current.model_settings)
                attempt["outcome"] = "response"
                return ModelResponse(result=[response])
            except Exception as exc:
                attempt["outcome"] = "context_length" if is_context_length_error(exc) else type(exc).__name__
                body = getattr(exc, "body", None)
                if isinstance(body, dict):
                    error = body.get("error", body)
                    if isinstance(error, dict):
                        attempt["provider_code"] = error.get("code")
                        attempt["provider_message"] = str(error.get("message", ""))[:500]
                raise

        response = ContextWindowGuardMiddleware().wrap_model_call(request, invoke)
        message = response.result[0] if isinstance(response, ModelResponse) else response
        text = message.text.strip()
        text = re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()
        if "<think>" in text:
            text = text.split("<think>", 1)[0].strip()
        report["outcome"] = "answered" if "验证通过" in text and not is_llm_error_message(message) else "clarification" if text == CONTEXT_RECOVERY_MESSAGE else "unexpected_response"
        report["visible_chars"] = len(text)
    except Exception as exc:
        report["outcome"] = type(exc).__name__
    report["elapsed_seconds"] = round(time.monotonic() - started, 2)
    print(json.dumps(report, ensure_ascii=True))
    return 0 if report["outcome"] == "answered" else 1


if __name__ == "__main__":
    raise SystemExit(main())