from __future__ import annotations

from typing import Any, Dict, Iterator, List

from amplitude_ai import (
    AIConfig,
    EVENT_TOOL_CALL,
    PROP_ENV,
    MockAmplitudeAI,
    ToolCallTracker,
    tool,
)
from amplitude_ai.core.tracking import PROP_AGENT_DESCRIPTION
from amplitude_ai.core.privacy import redact_pii_patterns, sanitize_any_content
from amplitude_ai.providers.anthropic import WrappedMessages
from amplitude_ai.utils.model_tiers import (
    TIER_FAST,
    TIER_REASONING,
    TIER_STANDARD,
    infer_model_tier,
)


def test_bound_agent_merge_does_not_mutate_event_properties() -> None:
    mock = MockAmplitudeAI()
    parent = mock.agent(
        "parent",
        user_id="u1",
    )
    agent = parent.child("bot", device_id="dev1", browser_session_id="bs1")
    event_properties: Dict[str, Any] = {"custom": True}

    agent.track_user_message(
        content="hello",
        session_id="s1",
        event_properties=event_properties,
    )

    assert event_properties == {"custom": True}


def test_session_inject_does_not_mutate_event_properties() -> None:
    mock = MockAmplitudeAI()
    agent = mock.agent("bot", user_id="u1")
    event_properties: Dict[str, Any] = {"custom": True}

    with agent.session(
        session_id="s1",
        device_id="dev1",
        browser_session_id="bs1",
    ) as session:
        session.track_user_message(content="hello", event_properties=event_properties)

    assert event_properties == {"custom": True}


def test_tool_context_resolution_preserves_explicit_empty_string_env() -> None:
    mock = MockAmplitudeAI()
    ToolCallTracker.clear()
    ToolCallTracker.set_amplitude(
        amplitude=mock.amplitude,
        user_id="u1",
        session_id="s1",
        env="tracker-env",
    )

    def identity(value: str) -> str:
        return value

    wrapped = tool(
        identity,
        name="empty-env",
        amplitude=mock.amplitude,
        user_id="u1",
        session_id="s1",
        env="",
    )
    wrapped("ok")

    tool_event = mock.get_events(EVENT_TOOL_CALL)[0]
    props = tool_event.event_properties or {}
    assert props.get(PROP_ENV) != "tracker-env"


def test_model_tier_inference_handles_opus_and_short_o_series() -> None:
    assert infer_model_tier("claude-3-opus-20240229") == TIER_STANDARD
    assert infer_model_tier("o1-mini") == TIER_FAST
    assert infer_model_tier("o3-mini") == TIER_FAST
    assert infer_model_tier("o4-mini") == TIER_FAST
    assert infer_model_tier("o1-2024-12-17") == TIER_REASONING
    assert infer_model_tier("o100") == TIER_STANDARD
    assert infer_model_tier("co3-something") == TIER_STANDARD
    assert infer_model_tier("o13") == TIER_STANDARD


class _FakeObj:
    def __init__(self, **kwargs: Any) -> None:
        self.__dict__.update(kwargs)


class _FakeStreamingTracker:
    def __init__(self) -> None:
        self.token_usage: Dict[str, int] = {}
        self.finish_reason: str | None = None
        self.finalized_with_model: str | None = None

    def add_reasoning_chunk(self, _chunk: str) -> None:
        return None

    def add_content_chunk(self, _chunk: str) -> None:
        return None

    def add_tool_call(self, _tool_call: Dict[str, Any]) -> None:
        return None

    def update_token_usage(self, usage: Dict[str, int]) -> None:
        self.token_usage.update(usage)

    def set_finish_reason(self, reason: str) -> None:
        self.finish_reason = reason

    def set_error(self, _error_message: str) -> None:
        return None

    def finalize(self, model_name: str) -> None:
        self.finalized_with_model = model_name


class _FakeOriginalMessages:
    def __init__(self, events: List[Any]) -> None:
        self._events = events

    def create(self, **_kwargs: Any) -> Iterator[Any]:
        return iter(self._events)


class _FakeAnthropicClient:
    def __init__(self, events: List[Any], tracker: _FakeStreamingTracker) -> None:
        self._original_messages = _FakeOriginalMessages(events)
        self._propagate_context = False
        self._default_amplitude_user_id = None
        self.amplitude = MockAmplitudeAI().amplitude
        self._tracker = tracker

    def create_streaming_tracker(self, **_kwargs: Any) -> _FakeStreamingTracker:
        return self._tracker


def test_anthropic_streaming_uses_message_start_input_tokens_and_message_delta_finish_reason() -> None:
    tracker = _FakeStreamingTracker()
    events = [
        _FakeObj(
            type="message_start",
            message=_FakeObj(
                model="claude-3-5-sonnet",
                usage=_FakeObj(
                    input_tokens=42,
                    cache_read_input_tokens=3,
                    cache_creation_input_tokens=2,
                ),
            ),
        ),
        _FakeObj(
            type="message_delta",
            delta=_FakeObj(stop_reason="end_turn"),
            usage=_FakeObj(output_tokens=7),
        ),
        _FakeObj(type="message_stop"),
    ]

    client = _FakeAnthropicClient(events=events, tracker=tracker)
    wrapped = WrappedMessages(client)  # type: ignore[arg-type]

    stream = wrapped._create_streaming(  # pylint: disable=protected-access
        amplitude_user_id="u1",
        model="claude-3-5-sonnet",
    )
    for _event in stream:
        pass

    assert tracker.token_usage.get("input_tokens") == 47
    assert tracker.token_usage.get("output_tokens") == 7
    assert tracker.finish_reason == "end_turn"
    assert tracker.finalized_with_model == "claude-3-5-sonnet"


def test_debug_and_dry_run_both_log_when_enabled(capsys: Any) -> None:
    mock = MockAmplitudeAI(config=AIConfig(debug=True, dry_run=True))

    mock.track_user_message(
        user_id="u1",
        content="hello",
        session_id="s1",
    )

    captured = capsys.readouterr()
    assert "[amplitude-ai]" in captured.err
    assert '"event"' in captured.err
    # dry_run should still suppress transport forwarding.
    assert len(mock.events) == 0


def test_sanitize_any_content_extracts_text_from_structured_text_blocks() -> None:
    result = sanitize_any_content({"type": "text", "text": "Hello world"})
    llm_message = result.get("$llm_message", {})
    assert isinstance(llm_message, dict)
    assert llm_message.get("text") == "Hello world"


def test_sanitize_any_content_omits_llm_message_for_empty_content() -> None:
    assert sanitize_any_content("") == {}
    assert sanitize_any_content(None) == {}
    assert sanitize_any_content([]) == {}
    assert sanitize_any_content({"type": "text", "text": ""}) == {}


def test_redact_pii_patterns_does_not_treat_pipe_as_valid_tld_char() -> None:
    assert redact_pii_patterns("Invalid: user@example.c|m") == "Invalid: user@example.c|m"


def test_description_propagates_through_agent_to_events() -> None:
    mock = MockAmplitudeAI()
    agent = mock.agent("chatbot", user_id="u1", description="Customer support chatbot")
    agent.track_user_message(content="Hello", session_id="s1")
    agent.track_ai_message(
        content="Hi",
        session_id="s1",
        model="gpt-4o",
        provider="openai",
        latency_ms=100,
    )

    for event in mock.events:
        props = event.event_properties or {}
        assert props.get(PROP_AGENT_DESCRIPTION) == "Customer support chatbot"


def test_description_not_inherited_by_child_agents() -> None:
    mock = MockAmplitudeAI()
    parent = mock.agent("orchestrator", user_id="u1", description="Main orchestrator")
    child = parent.child("worker")

    child.track_user_message(content="From child", session_id="s1")
    props = mock.events[0].event_properties or {}
    assert props.get(PROP_AGENT_DESCRIPTION) is None


def test_description_overridable_on_child() -> None:
    mock = MockAmplitudeAI()
    parent = mock.agent("orchestrator", user_id="u1", description="Main orchestrator")
    child = parent.child("worker", description="Specialized worker")

    child.track_user_message(content="From child", session_id="s1")
    props = mock.events[0].event_properties or {}
    assert props.get(PROP_AGENT_DESCRIPTION) == "Specialized worker"


def test_description_omitted_when_not_set() -> None:
    mock = MockAmplitudeAI()
    agent = mock.agent("chatbot", user_id="u1")
    agent.track_user_message(content="Hello", session_id="s1")

    props = mock.events[0].event_properties or {}
    assert PROP_AGENT_DESCRIPTION not in props
