from __future__ import annotations

import json
from unittest.mock import patch

from a2a_llm_tracker.agent_proxy import (
    _tracking_headers,
    call_agent,
    call_orchestra_agent,
)
from a2a_llm_tracker.context import set_context
from a2a_llm_tracker.middleware import get_request_id, get_session_id, set_request_id, set_session_id


class _MockHeaders:
    def __init__(self, values: dict[str, str]):
        self._values = {k.lower(): v for k, v in values.items()}

    def get(self, key: str, default=None):
        return self._values.get(key.lower(), default)

    def items(self):
        return list(self._values.items())


class _MockResponse:
    def __init__(self, body: bytes, headers: dict[str, str] | None = None, status: int = 200):
        self._body = body
        self.headers = _MockHeaders(headers or {"content-type": "application/json"})
        self.status = status

    def read(self) -> bytes:
        return self._body

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


def test_tracking_headers_prefers_request_and_session_contextvars():
    set_request_id("req-1")
    set_session_id("sess-1")
    set_context(trace_id="trace-1")

    headers = _tracking_headers()

    assert headers["X-Request-ID"] == "req-1"
    assert headers["X-Session-ID"] == "sess-1"
    assert headers["X-Trace-ID"] == "trace-1"


def test_tracking_headers_falls_back_to_trace_for_request_id():
    set_request_id("")
    set_session_id("")
    set_context(trace_id="trace-fallback", session_id="sess-fallback")

    headers = _tracking_headers()

    assert headers["X-Request-ID"] == "trace-fallback"
    assert headers["X-Trace-ID"] == "trace-fallback"
    assert headers["X-Session-ID"] == "sess-fallback"


def test_call_agent_sends_json_payload_and_tracking_headers():
    set_request_id("req-22")
    set_session_id("sess-22")
    set_context(trace_id="trace-22")

    captured = {}

    def _fake_urlopen(req, timeout=30.0):
        captured["headers"] = {k.lower(): v for k, v in req.header_items()}
        captured["method"] = req.get_method()
        captured["timeout"] = timeout
        captured["body"] = req.data
        return _MockResponse(json.dumps({"ok": True}).encode("utf-8"))

    with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
        out = call_agent(
            "https://example.com/agent",
            payload={"hello": "world"},
            timeout=5.0,
        )

    assert out == {"ok": True}
    assert captured["method"] == "POST"
    assert captured["timeout"] == 5.0
    assert captured["headers"]["x-request-id"] == "req-22"
    assert captured["headers"]["x-session-id"] == "sess-22"
    assert captured["headers"]["x-trace-id"] == "trace-22"
    assert captured["headers"]["content-type"] == "application/json"
    assert json.loads(captured["body"].decode("utf-8")) == {"hello": "world"}


def test_call_orchestra_agent_builds_payload_and_extracts_text():
    captured = {}
    set_request_id("req-parent")
    set_session_id("sess-parent")
    set_context(trace_id="trace-parent")

    response_json = {
        "result": {
            "contextId": "ctx-1",
            "requestId": "agent-req-1",
            "artifacts": [{"parts": [{"text": "hello from agent"}]}],
        }
    }

    def _fake_urlopen(req, timeout=30.0):
        captured["url"] = req.full_url
        captured["headers"] = {k.lower(): v for k, v in req.header_items()}
        captured["body"] = json.loads(req.data.decode("utf-8"))
        return _MockResponse(
            json.dumps(response_json).encode("utf-8"),
            headers={
                "content-type": "application/json",
                "x-request-id": "req-from-server",
                "x-session-id": "sess-from-server",
                "x-trace-id": "trace-from-server",
            },
        )

    with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
        out = call_orchestra_agent(
            agent_id="77",
            input_message="hi there",
            agent_url="https://node.example.com/a2a",
            stream=False,
            user_token="token-1",
        )

    assert captured["url"] == "https://node.example.com/a2a/77/send"
    assert captured["headers"]["accept"] == "application/json"
    assert captured["headers"]["authorization"] == "Bearer token-1"
    assert "x-call-id" in captured["headers"]
    assert captured["headers"]["x-parent-call-id"] == "req-parent"
    assert captured["headers"]["x-sequence-number"] == "1"
    assert captured["body"]["method"] == "message/send"
    assert captured["body"]["params"]["message"]["parts"][0]["text"] == "hi there"
    assert out["success"] is True
    assert out["text"] == "hello from agent"
    assert out["context_id"] == "ctx-1"
    assert out["agent_request_id"] == "agent-req-1"
    assert get_request_id() == "req-from-server"
    assert get_session_id() == "sess-from-server"


def test_call_orchestra_agent_uses_ccs_token_when_user_token_missing():
    captured = {}

    class _Client:
        accessToken = "ccs-token-xyz"

    def _fake_urlopen(req, timeout=30.0):
        captured["headers"] = {k.lower(): v for k, v in req.header_items()}
        return _MockResponse(
            json.dumps({"result": {"artifacts": [{"parts": [{"text": "ok"}]}]}}).encode("utf-8"),
            headers={"content-type": "application/json"},
        )

    with patch("a2a_llm_tracker.config.get_ccs_client", return_value=_Client()):
        with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
            out = call_orchestra_agent(agent_id="22", input_message="hello", stream=False)

    assert out["success"] is True
    assert captured["headers"]["authorization"] == "Bearer ccs-token-xyz"


def test_call_orchestra_agent_stream_parses_sse():
    set_request_id("req-stream-parent")
    set_session_id("sess-stream-parent")
    set_context(trace_id="trace-stream-parent")

    sse = (
        'data: {"result":{"artifacts":[{"parts":[{"text":"Hello "}]}]}}\n\n'
        'data: {"result":{"artifacts":[{"parts":[{"text":"world"}]}], "contextId":"ctx-2"}}\n\n'
        "data: [DONE]\n\n"
    )

    def _fake_urlopen(req, timeout=30.0):
        return _MockResponse(
            sse.encode("utf-8"),
            headers={
                "content-type": "text/event-stream",
                "x-request-id": "req-stream",
                "x-session-id": "sess-stream",
            },
        )

    with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
        out = call_orchestra_agent(
            agent_id="11",
            input_message="stream this",
            stream=True,
        )

    assert out["success"] is True
    assert out["text"] == "Hello world"
    assert len(out["events"]) == 2
    assert out["context_id"] == "ctx-2"


def test_call_orchestra_agent_sequence_increments_for_same_trace():
    captured_headers = []
    set_request_id("req-seq")
    set_session_id("sess-seq")
    set_context(trace_id="trace-seq")

    def _fake_urlopen(req, timeout=30.0):
        captured_headers.append({k.lower(): v for k, v in req.header_items()})
        return _MockResponse(
            json.dumps({"result": {"artifacts": [{"parts": [{"text": "ok"}]}]}}).encode("utf-8"),
            headers={"content-type": "application/json"},
        )

    with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
        out1 = call_orchestra_agent(agent_id="1", input_message="one", stream=False)
        out2 = call_orchestra_agent(agent_id="1", input_message="two", stream=False)

    assert out1["sequence_number"] == 1
    assert out2["sequence_number"] == 2
    assert captured_headers[0]["x-sequence-number"] == "1"
    assert captured_headers[1]["x-sequence-number"] == "2"


def test_call_orchestra_agent_accepts_agent_card_url_without_agent_id():
    captured = {}

    def _fake_urlopen(req, timeout=30.0):
        captured["url"] = req.full_url
        return _MockResponse(
            json.dumps({"result": {"artifacts": [{"parts": [{"text": "ok"}]}]}}).encode("utf-8"),
            headers={"content-type": "application/json"},
        )

    with patch("a2a_llm_tracker.agent_proxy.transport.urlopen", side_effect=_fake_urlopen):
        out = call_orchestra_agent(
            input_message="hello",
            agent_url="https://node.a2aorchestra.com/api/v1/agents/sendmessage/103219967/.well-known/agent-card.json",
        )

    assert captured["url"] == "https://node.a2aorchestra.com/api/v1/agents/sendmessage/103219967/send"
    assert out["success"] is True
    assert out["text"] == "ok"


def test_apply_response_ids_preserves_existing_request_trace_session_context():
    from a2a_llm_tracker.agent_proxy.headers import _apply_response_ids
    from a2a_llm_tracker.context import get_context

    set_request_id("req-root")
    set_session_id("sess-root")
    set_context(trace_id="trace-root", session_id="sess-root")

    out = _apply_response_ids(
        {
            "X-Request-ID": "req-downstream",
            "X-Trace-ID": "trace-downstream",
            "X-Session-ID": "sess-downstream",
        }
    )

    assert out["request_id"] == "req-root"
    assert out["trace_id"] == "trace-root"
    assert out["session_id"] == "sess-root"
    assert get_request_id() == "req-root"
    assert get_session_id() == "sess-root"
    assert get_context().trace_id == "trace-root"
