import pytest
from pydantic import BaseModel
from openai.resources.chat.completions import ChatCompletion

from instructor import openai_schema, OpenAISchema, Mode
from instructor.exceptions import IncompleteOutputException

@pytest.fixture
def test_model():
    class TestModel(OpenAISchema):
        name: str = "TestModel"
        data: str

    return TestModel
  
@pytest.fixture
def mock_completion(request):
    finish_reason = 'stop'
    data_content = "{\n\"data\": \"complete data\"\n}"

    if hasattr(request, 'param'):
        finish_reason = request.param.get('finish_reason', finish_reason)
        data_content = request.param.get('data_content', data_content)


    mock_choices = [{
        "index": 0,
        "message": {
            "role": "assistant",
            "function_call": {
                "name": "TestModel",
                "arguments": data_content
            },
            "content": data_content,
        },
        "finish_reason": finish_reason
    }]

    completion = ChatCompletion(
        id="test_id",
        choices=mock_choices,
        created=1234567890,
        model="gpt-3.5-turbo",
        object="chat.completion",
    )
    
    return completion

def test_openai_schema():
    @openai_schema
    class Dataframe(BaseModel):
        """
        Class representing a dataframe. This class is used to convert
        data into a frame that can be used by pandas.
        """

        data: str
        columns: str

        def to_pandas(self):
            pass

    assert hasattr(Dataframe, "openai_schema")
    assert hasattr(Dataframe, "from_response")
    assert hasattr(Dataframe, "to_pandas")
    assert Dataframe.openai_schema["name"] == "Dataframe"  # type: ignore


def test_openai_schema_raises_error():
    with pytest.raises(TypeError, match="must be a subclass of pydantic.BaseModel"):

        @openai_schema
        class Dummy:
            pass


def test_no_docstring():
    class Dummy(OpenAISchema):
        attr: str

    assert (
        Dummy.openai_schema["description"]
        == "Correctly extracted `Dummy` with all the required parameters with correct types"
    )

@pytest.mark.parametrize('mock_completion', [{'finish_reason': 'length', 'data_content': '{\n\"data\": \"incomplete dat\"\n}'}], indirect=True)
def test_incomplete_output_exception(test_model, mock_completion):
    with pytest.raises(IncompleteOutputException):
        test_model.from_response(mock_completion)
        
def test_complete_output_no_exception(test_model, mock_completion):
    test_model_instance = test_model.from_response(mock_completion)
    assert test_model_instance.data == "complete data"

@pytest.mark.asyncio
@pytest.mark.parametrize('mock_completion', [{'finish_reason': 'length', 'data_content': '{\n\"data\": \"incomplete dat\"\n}'}], indirect=True)
async def test_incomplete_output_exception(test_model, mock_completion):
    with pytest.raises(IncompleteOutputException):
        await test_model.from_response(mock_completion)
        
@pytest.mark.asyncio
async def test_complete_output_no_exception(test_model, mock_completion):
    test_model_instance = await test_model.from_response_async(mock_completion)
    assert test_model_instance.data == "complete data"