import re
import tempfile
import warnings

import pytest
from chatlas import (
    AssistantTurn,
    ChatOpenAI,
    ChatOpenAICompletions,
    ContentToolRequest,
    ContentToolResult,
    ToolRejectError,
    Turn,
    UserTurn,
)
from chatlas._chat import ToolFailureWarning, finalize_assistant_turn
from pydantic import BaseModel


@pytest.mark.vcr
def test_simple_batch_chat():
    chat = ChatOpenAI(model="gpt-5.4")
    response = chat.chat("What's 1 + 1. Just give me the answer, no punctuation")
    assert str(response) == "2"


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_simple_async_batch_chat():
    chat = ChatOpenAI(model="gpt-5.4")
    response = await chat.chat_async(
        "What's 1 + 1. Just give me the answer, no punctuation",
    )
    assert "2" == await response.get_content()


@pytest.mark.vcr
def test_simple_streaming_chat():
    chat = ChatOpenAI(model="gpt-5.4")
    res = chat.stream(
        """
        What are the canonical colors of the ROYGBIV rainbow?
        Put each colour on its own line. Don't use punctuation.
    """
    )
    chunks = [chunk for chunk in res]
    assert len(chunks) > 2
    result = "".join(chunks)
    res = re.sub(r"\s+", "", result).lower()
    assert res == "redorangeyellowgreenblueindigoviolet"
    turn = chat.get_last_turn()
    assert turn is not None
    res = re.sub(r"\s+", "", turn.text).lower()
    assert res == "redorangeyellowgreenblueindigoviolet"


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_simple_streaming_chat_async():
    chat = ChatOpenAI(model="gpt-5.4")
    res = await chat.stream_async(
        """
        What are the canonical colors of the ROYGBIV rainbow?
        Put each colour on its own line. Don't use punctuation.
    """
    )
    chunks = [chunk async for chunk in res]
    assert len(chunks) > 2
    result = "".join(chunks)
    rainbow_re = "^red *\norange *\nyellow *\ngreen *\nblue *\nindigo *\nviolet *\n?$"
    assert re.match(rainbow_re, result.lower())
    turn = chat.get_last_turn()
    assert turn is not None
    assert re.match(rainbow_re, turn.text.lower())


def test_model_property_get_and_set():
    chat = ChatOpenAI(model="gpt-4o")
    assert chat.model == "gpt-4o"
    assert chat.provider.model == "gpt-4o"

    # The model name is intentionally not validated.
    chat.model = "some-unvalidated-model"
    assert chat.model == "some-unvalidated-model"
    assert chat.provider.model == "some-unvalidated-model"


def test_basic_repr(snapshot):
    chat = ChatOpenAI(
        model="gpt-5.4",
        system_prompt="You're a helpful assistant that returns very minimal output",
    )
    chat.set_turns(
        [
            UserTurn("What's 1 + 1? What's 1 + 2?"),
            AssistantTurn("2  3", tokens=(15, 5, 5), cost=0.001),
        ]
    )
    assert snapshot == repr(chat)


def test_basic_str(snapshot):
    chat = ChatOpenAI(
        model="gpt-5.4",
        system_prompt="You're a helpful assistant that returns very minimal output",
    )
    chat.set_turns(
        [
            UserTurn("What's 1 + 1? What's 1 + 2?"),
            AssistantTurn("2  3", tokens=(15, 5, 0), cost=0.001),
        ]
    )
    assert snapshot == str(chat)


def test_basic_export(snapshot):
    chat = ChatOpenAI(
        system_prompt="You're a helpful assistant that returns very minimal output",
    )
    chat.set_turns(
        [
            UserTurn("What's 1 + 1? What's 1 + 2?"),
            AssistantTurn("2  3", tokens=(15, 5, 0)),
        ]
    )
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpfile = tmpdir + "/chat.html"
        chat.export(tmpfile, title="My Chat")
        with open(tmpfile, "r") as f:
            assert snapshot == f.read()


@pytest.mark.vcr
def test_chat_structured():
    chat = ChatOpenAI(model="gpt-5.4")

    class Person(BaseModel):
        name: str
        age: int

    data = chat.chat_structured("John, age 15, won first prize", data_model=Person)
    assert data == Person(name="John", age=15)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_chat_structured_async():
    chat = ChatOpenAI(model="gpt-5.4")

    class Person(BaseModel):
        name: str
        age: int

    data = await chat.chat_structured_async(
        "John, age 15, won first prize", data_model=Person
    )
    assert data == Person(name="John", age=15)


@pytest.mark.vcr
def test_stream_with_data_model():
    from chatlas._content import ContentJson

    chat = ChatOpenAI(model="gpt-5.4")

    class Person(BaseModel):
        name: str
        age: int

    chunks = list(chat.stream("John, age 15, won first prize", data_model=Person))
    result = "".join(chunks)
    person = Person.model_validate_json(result)
    assert person == Person(name="John", age=15)

    # Verify the last turn contains ContentJson with the structured data
    turn = chat.get_last_turn()
    assert turn is not None
    assert len(turn.contents) == 1
    assert isinstance(turn.contents[0], ContentJson)
    assert turn.contents[0].value == {"name": "John", "age": 15}


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_stream_async_with_data_model():
    from chatlas._content import ContentJson

    chat = ChatOpenAI(model="gpt-5.4")

    class Person(BaseModel):
        name: str
        age: int

    chunks = [
        chunk
        async for chunk in await chat.stream_async(
            "John, age 15, won first prize", data_model=Person
        )
    ]
    result = "".join(chunks)
    person = Person.model_validate_json(result)
    assert person == Person(name="John", age=15)

    # Verify the last turn contains ContentJson with the structured data
    turn = chat.get_last_turn()
    assert turn is not None
    assert len(turn.contents) == 1
    assert isinstance(turn.contents[0], ContentJson)
    assert turn.contents[0].value == {"name": "John", "age": 15}


@pytest.mark.vcr
def test_last_turn_retrieval():
    chat = ChatOpenAI(model="gpt-5.4")
    assert chat.get_last_turn(role="user") is None
    assert chat.get_last_turn(role="assistant") is None

    chat.chat("Hi")
    user_turn = chat.get_last_turn(role="user")
    assert user_turn is not None and user_turn.role == "user"
    turn = chat.get_last_turn(role="assistant")
    assert turn is not None and turn.role == "assistant"


def test_system_prompt_retrieval():
    chat1 = ChatOpenAI()
    assert chat1.system_prompt is None
    assert chat1.get_last_turn(role="system") is None

    chat2 = ChatOpenAI(system_prompt="You are from New Zealand")
    assert chat2.system_prompt == "You are from New Zealand"
    turn = chat2.get_last_turn(role="system")
    assert turn is not None and turn.text == "You are from New Zealand"


def test_modify_system_prompt():
    chat = ChatOpenAI()
    chat.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello"),
        ]
    )

    # NULL -> NULL
    chat.system_prompt = None
    assert chat.system_prompt is None

    # NULL -> string
    chat.system_prompt = "x"
    assert chat.system_prompt == "x"

    # string -> string
    chat.system_prompt = "y"
    assert chat.system_prompt == "y"

    # string -> NULL
    chat.system_prompt = None
    assert chat.system_prompt is None


@pytest.mark.vcr
def test_json_serialize():
    chat = ChatOpenAI(model="gpt-5.4")
    chat.chat("Tell me a short joke", echo="none")
    turns = chat.get_turns()
    turns_json = [x.model_dump_json() for x in turns]
    turns_restored = [Turn.model_validate_json(x) for x in turns_json]

    assert len(turns) == 2
    # Verify correct types were restored
    assert type(turns_restored[0]) == type(turns[0])
    assert type(turns_restored[1]) == type(turns[1])

    # Completion objects, at least of right now, aren't included in the JSON
    # Need to create new turn without completion for comparison
    turns_for_comparison = [turns[0]]
    if isinstance(turns[1], AssistantTurn):
        turns_for_comparison.append(
            AssistantTurn(
                turns[1].contents,
                tokens=turns[1].tokens,
                finish_reason=turns[1].finish_reason,
                completion=None,
                cost=turns[1].cost,
            )
        )
    else:
        turns_for_comparison.append(turns[1])
    assert turns_for_comparison == turns_restored


# Chat can be deepcopied/forked
@pytest.mark.vcr
def test_deepcopy_chat():
    import copy

    chat = ChatOpenAI(model="gpt-5.4")
    chat.chat("Hi", echo="none")
    chat_fork = copy.deepcopy(chat)

    assert len(chat.get_turns()) == 2
    assert len(chat_fork.get_turns()) == 2

    chat_fork.chat("Bye", echo="none")

    assert len(chat.get_turns()) == 2
    assert len(chat_fork.get_turns()) == 4


@pytest.mark.vcr
def test_chat_callbacks():
    chat = ChatOpenAI(model="gpt-5.4")

    def test_tool(user: str) -> str:
        "Find out a user's favorite color"
        return "red"

    chat.register_tool(test_tool)

    last_request = None
    cb_count_request = 0
    cb_count_result = 0

    def on_tool_request(request: ContentToolRequest):
        nonlocal cb_count_request, last_request
        cb_count_request += 1
        assert isinstance(request, ContentToolRequest)
        assert request.name == "test_tool"
        last_request = request

    def on_tool_result(result: ContentToolResult):
        nonlocal cb_count_result, last_request
        cb_count_result += 1
        assert isinstance(result, ContentToolResult)
        assert result.request == last_request

    chat.on_tool_request(on_tool_request)
    chat.on_tool_result(on_tool_result)
    chat.chat("What are Joe and Hadley's favorite colors?")

    assert cb_count_request == 2
    assert cb_count_result == 2


@pytest.mark.vcr
@pytest.mark.filterwarnings("ignore", category=ToolFailureWarning)
def test_chat_tool_request_reject():
    chat = ChatOpenAI(model="gpt-4.1")

    def test_tool(user: str) -> str:
        "Find out a user's favorite color"
        return "red"

    chat.register_tool(test_tool)

    def on_tool_request(request: ContentToolRequest):
        if request.arguments["user"] == "Joe":
            raise ToolRejectError("Joe denied the request.")

    chat.on_tool_request(on_tool_request)

    response = chat.chat(
        "What are Joe and Hadley's favorite colors? ",
        "Write 'Joe ____ Hadley ____'. Use 'unknown' if you don't know. ",
        "Don't ever include punctuation in your answers.",
    )

    assert str(response).lower() == "joe unknown hadley red"


@pytest.mark.vcr
@pytest.mark.filterwarnings("ignore", category=ToolFailureWarning)
def test_chat_tool_request_reject2(capsys):
    chat = ChatOpenAI(model="gpt-4.1")

    def test_tool(user: str) -> str:
        "Find out a user's favorite color"
        if "joe" in user.lower():
            raise ToolRejectError("Joe denied the request.")
        return "red"

    chat.register_tool(test_tool)

    response = chat.chat(
        "What are Joe and Hadley's favorite colors? ",
        "Write 'Joe ____ Hadley ____'. Use 'unknown' if you don't know. ",
        "Don't ever include punctuation in your answers.",
    )

    assert str(response).lower() == "joe unknown hadley red"
    assert "Joe denied the request." in capsys.readouterr().out


def test_get_cost():
    chat = ChatOpenAI(api_key="fake_key")
    chat.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello", tokens=(2, 10, 2)),
            UserTurn("Hi"),
            AssistantTurn("Hello", tokens=(14, 10, 2)),
        ]
    )

    with pytest.raises(
        ValueError,
        match=re.escape(
            "Expected `include` to be one of 'all' or 'last', not 'bad_option'"
        ),
    ):
        chat.get_cost(include="bad_option")  # type: ignore

    # Checking that these have the right form vs. the actual calculation because the price may change
    cost = chat.get_cost(include="all")
    assert isinstance(cost, float)
    assert cost > 0

    last = chat.get_cost(include="last")
    assert isinstance(last, float)
    assert last > 0

    assert cost > last

    # User-specified cost values
    byoc = (2.0, 3.0, 0.1)

    expected_cost = (
        (10 * byoc[1] / 1e6) + (2 * byoc[0] / 1e6) + (2 * byoc[2] / 1e6)
    ) + ((10 * byoc[1] / 1e6) + (14 * byoc[0] / 1e6) + (2 * byoc[2] / 1e6))
    cost2 = chat.get_cost(include="all", token_price=byoc)
    assert cost2 == expected_cost

    # get_cost(include="last") returns the full cost of the last assistant turn
    # Last turn has tokens=(14, 10, 2) -> input=14, output=10, cached=2
    last_expected_cost = (
        (14 * byoc[0] / 1e6) + (10 * byoc[1] / 1e6) + (2 * byoc[2] / 1e6)
    )
    last2 = chat.get_cost(include="last", token_price=byoc)
    assert last2 == last_expected_cost

    chat2 = ChatOpenAI(api_key="fake_key", model="BADBAD")
    chat2.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello", tokens=(2, 10, 0)),
            UserTurn("Hi"),
            AssistantTurn("Hello", tokens=(14, 10, 0)),
        ]
    )
    with pytest.raises(
        KeyError,
        match="We could not locate pricing information for model 'BADBAD' from provider 'OpenAI'. If you know the pricing for this model, specify it in `token_price`.",
    ):
        chat2.get_cost(include="all")


# -----------------------------------------------------------------------------
# Tests for repr formatting
# -----------------------------------------------------------------------------


def test_turn_repr_user():
    """Test that UserTurn repr shows content without header."""
    turn = UserTurn("Hello, world!")
    assert repr(turn) == "Hello, world!"


def test_turn_repr_system():
    """Test that SystemTurn repr shows content without header."""
    from chatlas import SystemTurn

    turn = SystemTurn("You are a helpful assistant.")
    assert repr(turn) == "You are a helpful assistant."


def test_turn_repr_assistant_with_tokens():
    """Test that AssistantTurn repr shows content without header/token info."""
    turn = AssistantTurn("The answer is 42.", tokens=(100, 50, 20), cost=0.0025)
    assert repr(turn) == "The answer is 42."


def test_turn_repr_assistant_without_tokens():
    """Test that AssistantTurn repr works without token info."""
    turn = AssistantTurn("Hello!")
    assert repr(turn) == "Hello!"


def test_turn_repr_assistant_no_cached_tokens():
    """Test that AssistantTurn repr shows content without token info."""
    turn = AssistantTurn("Test", tokens=(100, 50, 0), cost=0.001)
    assert repr(turn) == "Test"


def test_chat_repr_header_format():
    """Test Chat repr header shows correct format."""
    chat = ChatOpenAI(api_key="fake_key", system_prompt="Be helpful")
    chat.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello!", tokens=(10, 20, 5), cost=0.001),
        ]
    )
    result = repr(chat)

    first_line = result.split("\n")[0]
    assert first_line.startswith("<Chat")
    assert "OpenAI" in first_line
    assert "turns=3" in first_line
    assert "input=10+5" in first_line
    assert "output=20" in first_line
    assert "cost=" in first_line


def test_chat_repr_no_tokens():
    """Test Chat repr when no token info is available."""
    chat = ChatOpenAI(api_key="fake_key")
    chat.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello!"),
        ]
    )
    result = repr(chat)

    first_line = result.split("\n")[0]
    assert "turns=2" in first_line
    assert "input=" not in first_line


def test_turn_repr_with_tool_request():
    """Test repr formatting for turns with tool requests."""
    from chatlas import ContentToolRequest

    turn = AssistantTurn(
        [
            ContentToolRequest(id="123", name="get_weather", arguments={"city": "NYC"}),
        ],
        tokens=(50, 30, 0),
    )
    result = repr(turn)

    assert "## Assistant" not in result
    assert "🔧 tool request (123)" in result
    assert 'get_weather(city="NYC")' in result


def test_turn_repr_with_tool_result():
    """Test repr formatting for turns with tool results."""
    request = ContentToolRequest(
        id="123", name="get_weather", arguments={"city": "NYC"}
    )
    turn = UserTurn(
        [
            ContentToolResult(value="72°F and sunny", request=request),
        ]
    )
    result = repr(turn)

    assert "## User" not in result
    assert "✅ tool result (123)" in result
    assert "72°F and sunny" in result


def test_str_unchanged():
    """Verify that __str__ still uses the original emoji-based format."""
    chat = ChatOpenAI(api_key="fake_key")
    chat.set_turns(
        [
            UserTurn("Hi"),
            AssistantTurn("Hello!"),
        ]
    )
    result = str(chat)

    assert "👤" in result or "User" in result
    assert "🤖" in result or "Assistant" in result


def test_json_serialize_with_tool_failure():
    """Test that get_turns() is serializable when a tool call results in failure."""
    # Create a tool request
    request = ContentToolRequest(
        id="test-123",
        name="my_tool",
        arguments={"x": 1},
    )

    # Create a tool result with an error (simulating tool failure)
    error = Exception("Something went wrong")
    result = ContentToolResult(
        value=None,
        error=error,
        request=request,
    )

    # Build turns that include the failed tool result
    chat = ChatOpenAI(api_key="fake_key")
    chat.set_turns(
        [
            UserTurn("Call the tool"),
            AssistantTurn([request]),
            UserTurn([result]),
        ]
    )

    # Verify serialization works
    turns = chat.get_turns()
    turns_json = [x.model_dump_json() for x in turns]

    # Verify we can deserialize them back
    turns_restored = [Turn.model_validate_json(x) for x in turns_json]

    assert len(turns_restored) == 3

    # Check the tool result turn was properly restored
    tool_result_turn = turns_restored[2]
    assert len(tool_result_turn.contents) == 1
    restored_result = tool_result_turn.contents[0]
    assert isinstance(restored_result, ContentToolResult)
    assert restored_result.error is not None
    # After serialization, the Exception becomes a string representation
    assert "Something went wrong" in str(restored_result.error)


@pytest.mark.vcr
def test_partial_turn_preserved_on_close():
    """Closing a streaming generator mid-response preserves the partial turn."""
    chat = ChatOpenAI(model="gpt-5.4")
    gen = chat.stream(
        """
        What are the canonical colors of the ROYGBIV rainbow?
        Put each colour on its own line. Don't use punctuation.
    """
    )
    # Consume a few chunks then close
    chunks = []
    for chunk in gen:
        chunks.append(chunk)
        if len(chunks) >= 3:
            break
    gen.close()

    turns = chat.get_turns()
    assert len(turns) == 2  # user + partial assistant
    assistant_turn = turns[1]
    assert isinstance(assistant_turn, AssistantTurn)
    assert assistant_turn.is_partial
    assert assistant_turn.partial_reason == "interrupted"
    assert len(assistant_turn.text) > 0


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_partial_turn_preserved_on_close_async():
    """Closing an async streaming generator mid-response preserves the partial turn."""
    chat = ChatOpenAI(model="gpt-5.4")
    gen = await chat.stream_async(
        """
        What are the canonical colors of the ROYGBIV rainbow?
        Put each colour on its own line. Don't use punctuation.
    """
    )
    # Consume a few chunks then close
    chunks = []
    async for chunk in gen:
        chunks.append(chunk)
        if len(chunks) >= 3:
            break
    await gen.aclose()

    turns = chat.get_turns()
    assert len(turns) == 2  # user + partial assistant
    assistant_turn = turns[1]
    assert isinstance(assistant_turn, AssistantTurn)
    assert assistant_turn.is_partial
    assert assistant_turn.partial_reason == "interrupted"
    assert len(assistant_turn.text) > 0


def test_partial_turn_display(snapshot):
    chat = ChatOpenAI(model="gpt-5.4")
    chat.set_turns(
        [
            UserTurn("hello"),
            AssistantTurn("response", tokens=(10, 5, 0), cost=0.001),
            UserTurn("more"),
            AssistantTurn("partial response", partial_reason="interrupted"),
        ]
    )
    assert snapshot == repr(chat)


def test_partial_turns_excluded_from_cost():
    chat = ChatOpenAI()
    chat.set_turns(
        [
            UserTurn("hello"),
            AssistantTurn("response", tokens=(10, 5, 0), cost=0.001),
            UserTurn("more"),
            AssistantTurn("partial", partial_reason="interrupted"),
        ]
    )
    # Cost should only include the complete turn
    cost = chat.get_cost()
    assert cost == 0.001

    # get_tokens should skip partial
    tokens = chat.get_tokens()
    assert len(tokens) == 2  # user + assistant from the complete turn only


def test_partial_turns_excluded_from_tokens_mid_conversation():
    """Partial turns in the middle of history (not just trailing) are excluded."""
    chat = ChatOpenAI()
    chat.set_turns(
        [
            UserTurn("hello"),
            AssistantTurn("partial", partial_reason="interrupted"),
            UserTurn("retry"),
            AssistantTurn("response", tokens=(10, 5, 0), cost=0.001),
        ]
    )
    # Cost should only include the complete turn
    cost = chat.get_cost()
    assert cost == 0.001

    # get_tokens should skip the partial pair entirely
    tokens = chat.get_tokens()
    assert len(tokens) == 2  # user + assistant from the complete turn only


def test_content_add():
    from chatlas._content import ContentText, ContentThinking, ContentThinkingDelta

    # ContentText addition
    a = ContentText.model_construct(text="a")
    b = ContentText.model_construct(text="b")
    result = a + b
    assert isinstance(result, ContentText)
    assert result.text == "ab"

    # ContentThinking addition (keeps latest extra)
    t1 = ContentThinking(thinking="a", extra={"k": 1})
    t2 = ContentThinking(thinking="b", extra={"k": 2})
    result = t1 + t2
    assert isinstance(result, ContentThinking)
    assert result.thinking == "ab"
    assert result.extra == {"k": 2}

    # ContentThinkingDelta addition (keeps first phase)
    d1 = ContentThinkingDelta(thinking="x", phase="start")
    d2 = ContentThinkingDelta(thinking="y", phase="body")
    result = d1 + d2
    assert isinstance(result, ContentThinkingDelta)
    assert result.thinking == "xy"
    assert result.phase == "start"

    # Mismatched types return NotImplemented
    assert a.__add__(t1) is NotImplemented


def test_token_count_complete_includes_history_and_system():
    chat = ChatOpenAICompletions(
        model="gpt-4o", system_prompt="You are a terse assistant."
    )
    new_only = chat.token_count("hello there", include="new")

    chat.set_turns(
        [
            UserTurn("a much earlier and quite lengthy question"),
            AssistantTurn("an earlier answer", tokens=(5, 5, 0)),
        ]
    )
    complete = chat.token_count("hello there", include="complete")

    assert new_only > 0
    assert complete > new_only


def test_token_count_invalid_include_raises():
    chat = ChatOpenAICompletions(model="gpt-4o")

    with pytest.raises(
        ValueError,
        match=re.escape(
            "Expected `include` to be one of 'new' or 'complete', not 'bad_option'"
        ),
    ):
        chat.token_count("hello", include="bad_option")  # type: ignore


def test_truncated_plain_chat_warns():
    # Structured extraction errors on a truncated response; plain chat keeps the
    # partial text but shouldn't hand it back silently (gh-315).
    chat = ChatOpenAI()
    turn = AssistantTurn("partial ans", finish_reason="max_tokens")

    with pytest.warns(UserWarning, match="max_tokens"):
        finalize_assistant_turn(chat.provider, turn)


def test_complete_plain_chat_does_not_warn():
    chat = ChatOpenAI()
    turn = AssistantTurn("full answer", finish_reason="success")

    with warnings.catch_warnings():
        warnings.simplefilter("error")
        finalize_assistant_turn(chat.provider, turn)
