import asyncio
from copy import deepcopy
import json
import hashlib
from types import SimpleNamespace
from uuid import uuid4

import httpx
from langchain_core.messages import AIMessageChunk, HumanMessage
import pytest

from artemis.llm.diagnostics import (
    ACTIVE_ATTEMPT,
    RequestAttempt,
    error_text,
    input_images,
    observe_request_headers,
    observe_response_headers,
)
from artemis.llm.router import ModelEndpoint
from artemis.services import llm
from artemis.runtime.cancellation import CancellationReason


class Engine:
    current_session_id = "diagnostic-test"
    current_step_id = None

    def __init__(self):
        self.records = []

    def record_trace(self, **record):
        self.records.append(deepcopy(record))
        return record.get("trace_id", uuid4())

    def _publish(self, *_):
        pass


@pytest.mark.parametrize(
    ("headers", "expected"),
    [
        ({"X-Oneapi-Request-Id": "gateway-request"}, "gateway-request"),
        (
            {"X-Oneapi-Request-Id": "gateway-request", "X-Request-Id": "provider-request"},
            "gateway-request",
        ),
        ({"X-Request-Id": "provider-request"}, "provider-request"),
        ({"Request-Id": "legacy-request"}, "legacy-request"),
        (
            {"X-Oneapi-Request-Id": "invalid header value", "X-Request-Id": "standard-request"},
            "standard-request",
        ),
    ],
)
def test_gateway_request_id_preserves_the_standard_header_without_assuming_its_source(
    headers, expected
):
    attempt = RequestAttempt(
        Engine(), request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    attempt.headers(httpx.Response(200, headers=headers))
    assert attempt.data["server_request_id"] == expected
    assert attempt.data["response_x_request_id"] == headers.get("X-Request-Id")


def test_sdk_error_cannot_overwrite_the_observed_gateway_request_id():
    class RequestError(Exception):
        request_id = "sdk-request"

    attempt = RequestAttempt(
        Engine(), request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    attempt.headers(httpx.Response(400, headers={"X-Oneapi-Request-Id": "gateway-request"}))
    attempt.finish("failed", error=RequestError("failed"))
    assert attempt.data["server_request_id"] == "gateway-request"
    assert attempt.data["error_request_id"] == "sdk-request"


@pytest.mark.asyncio
async def test_request_hook_correlates_api_calls_without_logging_header_values():
    engine = Engine()
    attempt = RequestAttempt(
        engine, request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    token = ACTIVE_ATTEMPT.set(attempt)
    try:
        request = httpx.Request(
            "POST",
            "https://example.test/v1/responses",
            headers={
                "Thread-Id": "private-thread",
                "Session-Id": "private-session",
                "Authorization": "Bearer private-key",
            },
        )
        await observe_request_headers(request)
        assert request.headers["x-client-request-id"] == "local-request"
        assert request.headers["x-client-attempt-id"] == str(attempt.trace_id)
        assert attempt.data["client_attempt_id_matches"] is True
        assert attempt.data["client_request_id_matches"] is True
        assert len(attempt.data["routing_header_fingerprints"]) == 2
        assert "private" not in json.dumps(engine.records, default=str)
        request.headers["X-Client-Request-Id"] = "caller-id"
        request.headers["X-Client-Attempt-Id"] = "caller-attempt"
        await observe_request_headers(request)
        assert request.headers["x-client-request-id"] == "caller-id"
        assert attempt.data["client_request_id_matches"] is False
        assert request.headers["X-Client-Attempt-Id"] == "caller-attempt"
        assert attempt.data["client_attempt_id_matches"] is False
    finally:
        ACTIVE_ATTEMPT.reset(token)


@pytest.mark.asyncio
async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
    )
    llm._ENDPOINT_BREAKER._states.clear()
    calls = 0
    wire_ids = []

    class StalledBody(httpx.AsyncByteStream):
        async def __aiter__(self):
            await asyncio.sleep(10)
            yield b""

    async def transport(request):
        nonlocal calls
        calls += 1
        wire_ids.append(
            (request.headers["X-Client-Request-Id"], request.headers["X-Client-Attempt-Id"])
        )
        headers = {"x-request-id": f"server-{calls}", "set-cookie": "secret-cookie"}
        return httpx.Response(
            200, headers=headers, stream=StalledBody() if calls == 1 else httpx.ByteStream(b"ok")
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(transport),
        event_hooks={"request": [observe_request_headers], "response": [observe_response_headers]},
    ) as client:

        class Model:
            async def astream(self, *_args, **_kwargs):
                async with client.stream(
                    "POST",
                    "https://example.test/responses",
                    headers={"Authorization": "Bearer secret-key"},
                ) as response:
                    await response.aread()
                yield AIMessageChunk(
                    content="ready",
                    usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12},
                )

        wrapper = llm.RobustChatModelWrapper(
            Model(), endpoint=ModelEndpoint(timeout_seconds=0.05, max_retries=1)
        )
        result = await wrapper.complete([])
    assert result.content == "ready"
    snapshots = [r["payload"] for r in engine.records if r["type"] == "llm_diagnostic"]
    failed = next(p for p in snapshots if p["phase"] == "failed")
    final = snapshots[-1]
    assert failed["http_status"] == 200
    assert failed["server_request_id"] == "server-1"
    assert failed["failure_phase"] == "waiting_output"
    assert failed["error_type"] == "TimeoutError" and "0.05 秒" in failed["error"]
    assert failed["usage"] is None and "first_output_at" not in failed
    assert final["request_id"] == failed["request_id"]
    assert final["attempt_id"] != failed["attempt_id"] and final["attempt"] == 2
    assert final["previous_attempt_id"] == failed["attempt_id"]
    assert wire_ids == [
        (failed["request_id"], failed["attempt_id"]),
        (final["request_id"], final["attempt_id"]),
    ]
    retry = next(r["payload"] for r in engine.records if r["name"] == "llm_retry")
    assert retry["request_id"] == failed["request_id"]
    assert retry["attempt_id"] == failed["attempt_id"]
    assert final["phase"] == "completed" and final["usage"]["total_tokens"] == 12
    assert (
        final["started_at"] <= final["headers_at"] <= final["first_output_at"] <= final["ended_at"]
    )
    assert "secret" not in json.dumps(snapshots)
    assert ACTIVE_ATTEMPT.get() is None
    assert any(r["name"] == "llm_retry" and r["payload"]["error"] for r in engine.records)


def test_tool_deltas_are_counted_once_without_retaining_arguments():
    engine = Engine()
    attempt = RequestAttempt(
        engine, request_id="tools", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    deltas = [
        AIMessageChunk(
            content=[], tool_call_chunks=[{"name": "click", "args": "", "id": "c1", "index": 0}]
        ),
        AIMessageChunk(
            content=[{"type": "function_call", "arguments": '{"secret":'}],
            tool_call_chunks=[{"args": '{"secret":', "index": 0}],
        ),
        AIMessageChunk(content=[], tool_call_chunks=[{"args": '"private"}', "index": 0}]),
        AIMessageChunk(
            content=[],
            tool_call_chunks=[{"name": "wait", "args": '{"seconds":1}', "id": "c2", "index": 1}],
        ),
    ]
    combined = deltas[0]
    for i, delta in enumerate(deltas):
        attempt.chunk(delta)
        if i:
            combined += delta
    attempt.chunk(AIMessageChunk(content=[], response_metadata={"status": "completed"}))
    attempt.finish("completed", response=combined)
    assert attempt.data["output_chunks"] == 4
    assert attempt.data["tool_argument_chars"] == len('{"secret":"private"}{"seconds":1}')
    for tool, name, arguments, call_id, index, count in zip(
        attempt.data["tools"],
        ["click", "wait"],
        ['{"secret":"private"}', '{"seconds":1}'],
        ["c1", "c2"],
        [0, 1],
        [2, 1],
        strict=True,
    ):
        expected = {
            "name": name,
            "call_id": call_id,
            "chunk_index": index,
            "argument_chars": len(arguments),
            "argument_bytes": len(arguments.encode()),
            "argument_sha256": hashlib.sha256(arguments.encode()).hexdigest(),
            "argument_deltas": count,
            "argument_prefixes": [],
            "argument_encoding_valid": True,
        }
        assert {key: tool[key] for key in expected} == expected
    assert attempt.data["response_status"] == "completed"
    assert "private" not in json.dumps(engine.records, default=str)
    assert "secret" not in json.dumps(engine.records, default=str)


@pytest.mark.asyncio
async def test_timeout_during_tool_generation_keeps_last_output_and_tool_name(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    llm._ENDPOINT_BREAKER._states.clear()
    calls = 0

    class Model:
        async def astream(self, *_args, **_kwargs):
            nonlocal calls
            calls += 1
            yield AIMessageChunk(
                content=[],
                tool_call_chunks=[{"name": "click", "args": '{"x":', "id": "c1", "index": 0}],
            )
            if calls == 1:
                await asyncio.sleep(10)
            yield AIMessageChunk(content=[], tool_call_chunks=[{"args": "1}", "index": 0}])

    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
    )
    wrapper = llm.RobustChatModelWrapper(
        Model(), endpoint=ModelEndpoint(timeout_seconds=0.03, max_retries=1, max_tokens=12000)
    )
    await wrapper.complete([])
    failed = next(
        r["payload"]
        for r in engine.records
        if r["type"] == "llm_diagnostic" and r["payload"]["phase"] == "failed"
    )
    assert failed["failure_phase"] == "generating_tool_arguments"
    assert failed["last_output_at"] >= failed["first_output_at"]
    assert len(failed["tools"]) == 1
    tool = failed["tools"][0]
    assert tool["name"] == "click" and tool["argument_chars"] == tool["argument_bytes"] == 5
    assert tool["argument_sha256"] == hashlib.sha256(b'{"x":').hexdigest()
    assert failed["usage"] is None
    assert failed["max_output_tokens"] == 12000


@pytest.mark.asyncio
async def test_concurrent_requests_do_not_mix_server_ids():
    engine = Engine()

    async def run(name):
        attempt = RequestAttempt(
            engine, request_id=name, attempt=1, timeout=5, max_retries=0, metadata={"model": name}
        )
        token = ACTIVE_ATTEMPT.set(attempt)
        try:
            await asyncio.sleep(0)
            await observe_response_headers(httpx.Response(200, headers={"x-request-id": name}))
        finally:
            ACTIVE_ATTEMPT.reset(token)
        return attempt.data

    first, second = await asyncio.gather(run("one"), run("two"))
    assert first["server_request_id"] == "one"
    assert second["server_request_id"] == "two"


@pytest.mark.asyncio
async def test_diagnostic_ids_survive_streaming_and_structured_completion_without_header_bodies():
    engine = Engine()
    attempt = RequestAttempt(
        engine, request_id="local", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    token = ACTIVE_ATTEMPT.set(attempt)
    traceparent = "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01"
    try:
        request = httpx.Request(
            "POST", "https://example.test/responses", headers={"x-client-request-id": "caller-id"}
        )
        await observe_request_headers(request)
        await observe_response_headers(
            httpx.Response(
                200,
                headers={
                    "x-request-id": "server-id",
                    "x-trace-id": "gateway-id",
                    "traceparent": traceparent,
                    "set-cookie": "private cookie",
                    "authorization": "private credential",
                },
            )
        )
        attempt.chunk(AIMessageChunk(content="hello", response_metadata={"id": "resp-first"}))
        attempt.finish(
            "completed",
            response={
                "raw": AIMessageChunk(
                    content="hello", response_metadata={"id": "resp-final", "status": "completed"}
                ),
                "parsed": {"private": "business details"},
            },
        )
    finally:
        ACTIVE_ATTEMPT.reset(token)
    assert request.headers["x-client-request-id"] == "caller-id"
    assert attempt.data["client_request_id"] == "caller-id"
    assert attempt.data["server_request_id"] == "server-id"
    assert attempt.data["server_trace_id"] == "gateway-id"
    assert attempt.data["traceparent"] == traceparent
    assert attempt.data["response_id"] == "resp-final"
    assert attempt.data["response_status"] == "completed"
    assert "private" not in json.dumps(engine.records, default=str)


def test_invalid_or_overlong_correlation_headers_are_omitted():
    attempt = RequestAttempt(
        Engine(), request_id="local", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    attempt.headers(
        httpx.Response(
            200,
            headers={
                "x-request-id": "x" * 257,
                "x-trace-id": "not an id",
                "traceparent": "not a trace",
            },
        )
    )
    assert attempt.data["server_request_id"] is None
    assert attempt.data["server_trace_id"] is None
    assert "traceparent" not in attempt.data


def test_error_redaction_and_image_headers():
    import base64
    import io
    from PIL import Image

    out = io.BytesIO()
    Image.new("RGB", (32, 48)).save(out, "PNG")
    message = HumanMessage(
        content=[
            {
                "type": "image_url",
                "image_url": {
                    "url": "data:image/png;base64," + base64.b64encode(out.getvalue()).decode()
                },
            }
        ]
    )
    assert input_images(([message],), {}) == [{"width": 32, "height": 48}]
    text = error_text(
        RuntimeError("Authorization: Bearer sk-my-secret https://host/path?token=private")
    )
    assert "sk-my-secret" not in text and "private" not in text
    assert error_text(RuntimeError()) == "RuntimeError"


@pytest.mark.parametrize(
    "kind, key", [("llm_diagnostic", "requests"), ("final_review_diagnostic", "final_reviews")]
)
def test_snapshots_survive_out_of_order_writes_and_export(tmp_path, kind, key):
    from artemis.data_engine.models import TraceRecord
    from artemis.data_engine.storage import StorageManager
    from apps.admin_console.database.connection import db_session
    from apps.admin_console.services.request_diagnostics import request_diagnostics

    path = tmp_path / "data.db"
    storage = StorageManager(db_path=path, base_trace_dir=tmp_path)
    session_id, trace_id = uuid4(), uuid4()
    with db_session(path) as conn:
        conn.execute(
            "INSERT INTO sessions(session_id,status) VALUES (?, 'completed')", (str(session_id),)
        )
        conn.commit()
    for sequence, phase in [(3, "completed"), (2, "receiving_output"), (1, "waiting_output")]:
        storage.create_trace(
            TraceRecord(
                trace_id=trace_id,
                session_id=session_id,
                type=kind,
                name="llm_request",
                timestamp=10 + sequence,
                status="success" if sequence == 3 else "running",
                payload={
                    "sequence": sequence,
                    "phase": phase,
                    "status": "completed" if sequence == 3 else "running",
                    "request_id": "call-one",
                    "started_at": 10,
                    "ended_at": 13 if sequence == 3 else None,
                },
            )
        )
    result = request_diagnostics(str(session_id), db_path=path)
    assert len(result[key]) == 1
    assert result[key][0]["phase"] == "completed"
    assert result[key][0]["active"] is False
    assert result[key][0]["timestamps_utc"]["started_at"].endswith("+00:00")
    assert result["total"] == (1 if kind == "llm_diagnostic" else 0)


def test_review_diagnostics_show_evidence_stage_without_a_live_model_request(tmp_path, monkeypatch):
    from artemis.data_engine.storage import StorageManager
    from apps.admin_console.database.connection import db_session
    from apps.admin_console.services import request_diagnostics as service

    path = tmp_path / "data.db"
    StorageManager(db_path=path, base_trace_dir=tmp_path)
    monkeypatch.setattr(service.time, "time", lambda: 130)
    with db_session(path) as conn:
        conn.execute("INSERT INTO sessions(session_id,status) VALUES ('task', 'running')")
        value = {
            "review_id": "final-review",
            "status": "running",
            "stage": "reading_evidence",
            "started_at": 100,
            "deadline_at": 200,
            "total_seconds": 100,
            "updated_at": 110,
        }
        conn.execute(
            "INSERT INTO traces(trace_id,session_id,type,timestamp,payload) "
            "VALUES ('review', 'task', 'final_review_diagnostic', 100, ?)",
            (json.dumps(value),),
        )
        conn.commit()
    result = service.request_diagnostics("task", db_path=path, limit=100)
    assert result["requests"] == []
    review = result["final_reviews"][0]
    assert review["active"] is True
    assert review["stage"] == "reading_evidence"
    assert review["elapsed_seconds"] == 30
    assert review["remaining_seconds"] == 70
    assert review["updated_at"] == 110  # preserve the actual last persisted event time
    assert review["timestamps_utc"]["deadline_at"].endswith("+00:00")
    assert (
        service.request_diagnostics("task", db_path=path, offset=1)["final_reviews"][0]["active"]
        is False
    )


@pytest.mark.asyncio
async def test_cancellation_closes_attempt_and_does_not_retry(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    llm._ENDPOINT_BREAKER._states.clear()

    async def cancelled():
        raise asyncio.CancelledError()

    with pytest.raises(asyncio.CancelledError):
        await llm._run_with_recovery(cancelled, max_retries=5)
    assert engine.records[-1]["payload"]["phase"] == "cancelled"
    assert ACTIVE_ATTEMPT.get() is None


@pytest.mark.asyncio
async def test_cancelling_retry_wait_does_not_leave_a_live_request(monkeypatch):
    waiting = asyncio.Event()

    class RetryEngine(Engine):
        def record_trace(self, **record):
            result = super().record_trace(**record)
            if record["payload"].get("phase") == "retry_wait":
                waiting.set()
            return result

    engine = RetryEngine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 60, max_attempts=2)
    )

    async def timeout():
        raise TimeoutError()

    task = asyncio.create_task(llm._run_with_recovery(timeout, max_retries=1))
    await asyncio.wait_for(waiting.wait(), 2)
    task.cancel(CancellationReason.GUIDANCE_REPLACED)
    with pytest.raises(asyncio.CancelledError):
        await task
    assert engine.records[-1]["payload"]["phase"] == "cancelled"
    assert "模型请求" in engine.records[-1]["payload"]["error"]
    assert engine.records[-1]["payload"]["cancel_reason"] == "guidance_replaced"
    assert engine.records[-1]["payload"]["cancel_phase"] == "retry_wait"
    assert "timeout_scope" not in engine.records[-1]["payload"]


def test_diagnostic_api_limits_preview_but_downloads_complete_history(tmp_path, monkeypatch):
    from fastapi import FastAPI
    from fastapi.testclient import TestClient
    from apps.admin_console.database import connection
    from apps.admin_console.routers.sessions import router
    from artemis.data_engine.storage import StorageManager

    path = tmp_path / "data.db"
    StorageManager(db_path=path, base_trace_dir=tmp_path)
    monkeypatch.setattr(connection, "DB_PATH", path)
    with connection.db_session() as conn:
        conn.execute("INSERT INTO sessions(session_id,status) VALUES ('task', 'running')")
        for i in range(105):
            value = {
                "request_id": "one-request",
                "attempt_id": str(i),
                "started_at": i,
                "phase": "completed" if i == 104 else "retry_wait",
            }
            conn.execute(
                "INSERT INTO traces(trace_id,session_id,type,timestamp,payload) VALUES (?, 'task', 'llm_diagnostic', ?, ?)",
                (str(i), i, json.dumps(value)),
            )
        conn.commit()
    app = FastAPI()
    app.include_router(router)
    with TestClient(app) as client:
        preview = client.get("/api/sessions/task/request-diagnostics")
        assert preview.status_code == 200
        assert len(preview.json()["requests"]) == 100
        assert not any(r["active"] for r in preview.json()["requests"])
        download = client.get("/api/sessions/task/request-diagnostics?download=true")
        assert download.headers["content-disposition"].startswith("attachment;")
        assert len(download.json()["requests"]) == 105
        assert download.json()["session_id"] == "task"