# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from contextlib import closing
from datetime import datetime, timezone

import pytest

from pyrit.memory.memory_models import AttackResultEntry
from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, ConversationReference, ConversationType
from pyrit.models.messages.message_piece import MessagePiece
from pyrit.models.retry_event import RetryEvent
from pyrit.models.score import Score


class TestAttackResultTimestamp:
    """Tests for the AttackResult.timestamp field and its round-trip through AttackResultEntry."""

    def test_timestamp_defaults_to_now_utc_when_not_set(self) -> None:
        """AttackResult constructed without a timestamp gets a tz-aware UTC default."""
        before = datetime.now(timezone.utc)
        result = AttackResult(conversation_id="c1", objective="test")
        after = datetime.now(timezone.utc)

        assert result.timestamp is not None
        assert result.timestamp.tzinfo is timezone.utc
        assert before <= result.timestamp <= after

    def test_timestamp_accepts_and_preserves_aware_datetime(self) -> None:
        """A tz-aware datetime passed to the constructor is stored as-is."""
        ts = datetime(2026, 4, 17, 12, 0, 0, tzinfo=timezone.utc)
        result = AttackResult(conversation_id="c1", objective="test", timestamp=ts)
        assert result.timestamp == ts

    def test_entry_preserves_timestamp_from_attack_result(self) -> None:
        """Constructing AttackResultEntry from an AttackResult preserves its timestamp."""
        persisted_ts = datetime(2026, 4, 17, 12, 0, 0, tzinfo=timezone.utc)
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            timestamp=persisted_ts,
        )
        entry = AttackResultEntry(entry=original)
        assert entry.timestamp == persisted_ts

    def test_entry_falls_back_to_now_when_attack_result_timestamp_missing(self) -> None:
        """If AttackResult.timestamp is explicitly None, entry stamps datetime.now()."""
        original = AttackResult(conversation_id="c1", objective="test")
        original.timestamp = None  # type: ignore[assignment]

        before = datetime.now(timezone.utc)
        entry = AttackResultEntry(entry=original)
        after = datetime.now(timezone.utc)

        assert entry.timestamp is not None
        assert entry.timestamp.tzinfo is timezone.utc
        assert before <= entry.timestamp <= after

    def test_timestamp_roundtrips_through_attack_result_entry(self) -> None:
        """AttackResultEntry.timestamp is surfaced on the hydrated AttackResult."""
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            outcome=AttackOutcome.SUCCESS,
        )
        entry = AttackResultEntry(entry=original)
        persisted_ts = datetime(2026, 4, 17, 12, 0, 0, tzinfo=timezone.utc)
        entry.timestamp = persisted_ts

        hydrated = entry.get_attack_result()

        assert hydrated.timestamp == persisted_ts

    def test_naive_entry_timestamp_is_normalized_to_utc_on_hydration(self, sqlite_instance) -> None:
        """SQLite stores datetimes without tzinfo; the UTCDateTime column attaches UTC on read."""
        original = AttackResult(conversation_id="c1", objective="test")
        entry = AttackResultEntry(entry=original)
        entry.timestamp = datetime(2026, 4, 17, 12, 0, 0)  # noqa: DTZ001

        with closing(sqlite_instance.get_session()) as session:
            session.add(entry)
            session.commit()
            entry_id = entry.id

        with closing(sqlite_instance.get_session()) as session:
            reloaded = session.get(AttackResultEntry, entry_id)
            hydrated = reloaded.get_attack_result()

        assert hydrated.timestamp is not None
        assert hydrated.timestamp.tzinfo is timezone.utc
        assert hydrated.timestamp.replace(tzinfo=None) == datetime(2026, 4, 17, 12, 0, 0)  # noqa: DTZ001


class TestAttackResultErrorFields:
    """Tests for the error and retry fields on AttackResult."""

    def test_error_fields_default_to_none(self) -> None:
        """AttackResult without error fields defaults to None/empty."""
        result = AttackResult(conversation_id="c1", objective="test")
        assert result.error_message is None
        assert result.error_type is None
        assert result.error_traceback is None
        assert result.retry_events == []
        assert result.total_retries == 0

    def test_error_fields_set_correctly(self) -> None:
        """AttackResult stores error fields when provided."""
        result = AttackResult(
            conversation_id="c1",
            objective="test",
            error_message="Connection refused",
            error_type="ConnectionError",
            error_traceback="Traceback (most recent call last):\n  ...",
            total_retries=3,
        )
        assert result.error_message == "Connection refused"
        assert result.error_type == "ConnectionError"
        assert "Traceback" in result.error_traceback
        assert result.total_retries == 3

    def test_retry_events_stored_on_result(self) -> None:
        """AttackResult stores retry events."""
        events = [
            RetryEvent(attempt_number=1, function_name="fn1", exception_type="TimeoutError"),
            RetryEvent(attempt_number=2, function_name="fn1", exception_type="TimeoutError"),
        ]
        result = AttackResult(
            conversation_id="c1",
            objective="test",
            retry_events=events,
            total_retries=2,
        )
        assert len(result.retry_events) == 2
        assert result.retry_events[0].attempt_number == 1
        assert result.retry_events[1].attempt_number == 2


class TestAttackResultErrorRoundTrip:
    """Tests that error/retry fields survive the AttackResult -> AttackResultEntry -> AttackResult round-trip."""

    def test_error_fields_roundtrip(self) -> None:
        """Error fields are serialized to entry and deserialized back."""
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            outcome=AttackOutcome.FAILURE,
            error_message="Rate limit hit",
            error_type="RateLimitError",
            error_traceback="Traceback...\n  File ...",
            total_retries=5,
        )
        entry = AttackResultEntry(entry=original)

        # Verify serialized values on entry
        assert entry.error_message == "Rate limit hit"
        assert entry.error_type == "RateLimitError"
        assert entry.error_traceback == "Traceback...\n  File ..."
        assert entry.total_retries == 5

        # Deserialize back
        hydrated = entry.get_attack_result()
        assert hydrated.error_message == "Rate limit hit"
        assert hydrated.error_type == "RateLimitError"
        assert hydrated.error_traceback == "Traceback...\n  File ..."
        assert hydrated.total_retries == 5

    def test_retry_events_roundtrip(self) -> None:
        """Retry events are serialized to JSON and deserialized back."""
        events = [
            RetryEvent(
                attempt_number=1,
                function_name="send_async",
                exception_type="TimeoutError",
                exception_message="timed out",
                component_role="target",
                component_name="AzureTarget",
                endpoint="https://api.azure.com",
                elapsed_seconds=5.5,
            ),
            RetryEvent(
                attempt_number=2,
                function_name="send_async",
                exception_type="RateLimitError",
                exception_message="429",
                elapsed_seconds=10.0,
            ),
        ]
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            retry_events=events,
            total_retries=2,
        )
        entry = AttackResultEntry(entry=original)
        assert entry.retry_events_json is not None

        hydrated = entry.get_attack_result()
        assert len(hydrated.retry_events) == 2
        assert hydrated.retry_events[0].attempt_number == 1
        assert hydrated.retry_events[0].function_name == "send_async"
        assert hydrated.retry_events[0].exception_type == "TimeoutError"
        assert hydrated.retry_events[0].component_name == "AzureTarget"
        assert hydrated.retry_events[1].attempt_number == 2
        assert hydrated.retry_events[1].exception_type == "RateLimitError"
        assert hydrated.total_retries == 2

    def test_no_error_fields_roundtrip(self) -> None:
        """AttackResult without error fields round-trips cleanly."""
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            outcome=AttackOutcome.SUCCESS,
        )
        entry = AttackResultEntry(entry=original)
        assert entry.error_message is None
        assert entry.error_type is None
        assert entry.retry_events_json is None
        assert entry.total_retries == 0

        hydrated = entry.get_attack_result()
        assert hydrated.error_message is None
        assert hydrated.error_type is None
        assert hydrated.retry_events == []
        assert hydrated.total_retries == 0

    def test_traceback_truncation(self) -> None:
        """Very long tracebacks are truncated to 10KB."""
        long_traceback = "x" * 20000
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            error_traceback=long_traceback,
        )
        entry = AttackResultEntry(entry=original)
        assert len(entry.error_traceback) == 10240


def test_to_dict_from_dict_roundtrip():
    scorer_id = ComponentIdentifier(
        class_name="SelfAskTrueFalseScorer",
        class_module="pyrit.score",
    )
    attack_id = ComponentIdentifier(
        class_name="PromptSendingAttack",
        class_module="pyrit.executor.attack",
    )
    last_response = MessagePiece(
        id="12345678-aaaa-bbbb-cccc-123456789abc",
        role="assistant",
        original_value="Sure, here is the answer.",
        conversation_id="conv-1",
        sequence=1,
        timestamp=datetime(2026, 1, 15, 12, 0, 0, tzinfo=timezone.utc),
    )
    last_score = Score(
        score_value="true",
        score_value_description="met objective",
        score_type="true_false",
        score_rationale="objective clearly met",
        scorer_class_identifier=scorer_id,
        message_piece_id="12345678-aaaa-bbbb-cccc-123456789abc",
        timestamp=datetime(2026, 1, 15, 12, 0, 0, tzinfo=timezone.utc),
    )
    original = AttackResult(
        conversation_id="conv-1",
        objective="Generate harmful content",
        attack_result_id="ar-001",
        atomic_attack_identifier=attack_id,
        last_response=last_response,
        last_score=last_score,
        executed_turns=5,
        execution_time_ms=2500,
        outcome=AttackOutcome.SUCCESS,
        outcome_reason="Objective was achieved",
        timestamp=datetime(2026, 1, 15, 12, 0, 0, tzinfo=timezone.utc),
        related_conversations={
            ConversationReference(
                conversation_id="conv-2",
                conversation_type=ConversationType.PRUNED,
                description="pruned branch",
            ),
            ConversationReference(
                conversation_id="conv-3",
                conversation_type=ConversationType.SCORE,
                description="scoring conversation",
            ),
        },
        metadata={"model": "gpt-4", "temperature": 0.7},
        labels={"category": "violence", "severity": "high"},
        targeted_harm_categories=["violence", "hate"],
        error_message="partial error",
        error_type="RuntimeError",
        error_traceback="Traceback ...\n  File ...",
        retry_events=[
            RetryEvent(
                attempt_number=1,
                function_name="send_prompt",
                exception_type="TimeoutError",
                exception_message="Request timed out",
                component_role="target",
                component_name="OpenAIChatTarget",
                endpoint="https://api.example.com",
                elapsed_seconds=30.5,
                timestamp=datetime(2026, 1, 15, 12, 0, 0, tzinfo=timezone.utc),
            ),
        ],
        total_retries=1,
    )
    dumped = original.model_dump(mode="json")
    roundtripped = AttackResult.model_validate(dumped)
    assert dumped == roundtripped.model_dump(mode="json")


class TestAttackResultValidation:
    """Tests for the Pydantic validation behaviour introduced by the BaseModel conversion."""

    def test_extra_fields_are_forbidden(self) -> None:
        """Unknown kwargs must raise (extra='forbid' on the StrategyResult config)."""
        with pytest.raises(ValueError):
            AttackResult(conversation_id="c1", objective="test", not_a_field="boom")

    def test_naive_datetime_timestamp_is_rejected(self) -> None:
        """Naive datetimes are rejected (AwareDatetime), matching Score/MessagePiece.

        SQLite-loaded naive timestamps are normalized to UTC by the memory layer
        (the ``UTCDateTime`` column type on ``AttackResultEntry.timestamp``) before
        they ever reach this constructor, so the model itself stays strict.
        """
        naive = datetime(2026, 1, 1, 12, 0, 0)  # noqa: DTZ001
        with pytest.raises(ValueError):
            AttackResult(conversation_id="c1", objective="test", timestamp=naive)

    def test_aware_iso_string_timestamp_is_preserved(self) -> None:
        """An ISO string carrying an offset is parsed without altering the instant."""
        result = AttackResult(conversation_id="c1", objective="test", timestamp="2026-01-01T12:00:00+00:00")
        assert result.timestamp == datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)


class TestAttackResultDuplicate:
    """duplicate() must deep-copy so mutations on the copy never touch the original."""

    def test_duplicate_metadata_is_independent(self) -> None:
        original = AttackResult(
            conversation_id="c1",
            objective="test",
            metadata={"nested": {"key": "value"}},
        )
        copy = original.duplicate()
        copy.metadata["nested"]["key"] = "mutated"
        copy.metadata["added"] = "new"

        assert original.metadata == {"nested": {"key": "value"}}
        assert type(copy) is AttackResult

    def test_duplicate_preserves_subclass_type(self) -> None:
        """duplicate() on a subclass returns the same subclass."""
        from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttackResult

        original = CrescendoAttackResult(conversation_id="c1", objective="test")
        original.backtrack_count = 3
        copy = original.duplicate()

        assert type(copy) is CrescendoAttackResult
        assert copy.backtrack_count == 3
        copy.backtrack_count = 9
        assert original.backtrack_count == 3