from __future__ import annotations

from pathlib import Path

import pytest
from msgspec import Struct
from pydantic import BaseModel
from sqlalchemy import Engine, String, select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column

from advanced_alchemy import base, mixins
from advanced_alchemy.repository import (
    SQLAlchemyAsyncRepository,
    SQLAlchemySyncRepository,
)
from advanced_alchemy.service import SQLAlchemyAsyncQueryService, SQLAlchemySyncQueryService
from advanced_alchemy.service._async import SQLAlchemyAsyncRepositoryService
from advanced_alchemy.service._sync import SQLAlchemySyncRepositoryService
from advanced_alchemy.service.typing import (
    is_msgspec_struct,
    is_msgspec_struct_with_field,
    is_msgspec_struct_without_field,
    is_pydantic_model,
    is_pydantic_model_with_field,
    is_pydantic_model_without_field,
)
from advanced_alchemy.utils.fixtures import open_fixture, open_fixture_async

pytestmark = [  # type: ignore
    pytest.mark.integration,
    pytest.mark.xdist_group("sqlquery_service"),
]
here = Path(__file__).parent
fixture_path = here.parent.parent / "examples"
state_registry = base.create_registry()


@pytest.fixture()
def sqlquery_test_tables(engine: Engine) -> None:
    """Create sqlquery test tables for sync engines."""
    if getattr(engine.dialect, "name", "") != "mock":
        state_registry.metadata.create_all(engine)


@pytest.fixture()
async def sqlquery_test_tables_async(async_engine: AsyncEngine) -> None:
    """Create sqlquery test tables for async engines."""
    if getattr(async_engine.dialect, "name", "") != "mock":
        async with async_engine.begin() as conn:
            await conn.run_sync(state_registry.metadata.create_all)


class UUIDBase(mixins.UUIDPrimaryKey, base.CommonTableAttributes, DeclarativeBase):
    """Base for all SQLAlchemy declarative models with UUID primary keys."""

    registry = state_registry


class USState(UUIDBase):
    __tablename__ = "us_state_lookup"  # type: ignore[assignment]
    abbreviation: Mapped[str] = mapped_column(String(5))
    name: Mapped[str] = mapped_column(String(50))


class USStateStruct(Struct):
    abbreviation: str
    name: str


class USStateBaseModel(BaseModel):
    abbreviation: str
    name: str


class USStateSyncRepository(SQLAlchemySyncRepository[USState]):
    """US State repository."""

    model_type = USState


class USStateSyncService(SQLAlchemySyncRepositoryService[USState, USStateSyncRepository]):
    """US State repository."""

    repository_type = USStateSyncRepository


class USStateAsyncRepository(SQLAlchemyAsyncRepository[USState]):
    """US State repository."""

    model_type = USState


class USStateAsyncService(SQLAlchemyAsyncRepositoryService[USState, USStateAsyncRepository]):
    """US State repository."""

    repository_type = USStateAsyncRepository


class StateQuery(base.SQLQuery):
    """Nonsensical query to test custom SQL queries."""

    __table__ = select(  # type: ignore
        USState.abbreviation.label("state_abbreviation"),
        USState.name.label("state_name"),
    ).alias("state_lookup")
    __mapper_args__ = {
        "primary_key": [USState.abbreviation],
    }
    state_abbreviation: str
    state_name: str


class StateQueryStruct(Struct):
    state_abbreviation: str
    state_name: str


class StateQueryBaseModel(BaseModel):
    state_abbreviation: str
    state_name: str


@pytest.mark.xdist_group("sqlquery")
def test_sync_fixture_and_query(engine: Engine, sqlquery_test_tables: None) -> None:
    # Skip mock engines as they don't support proper query operations
    if getattr(engine.dialect, "name", "") == "mock":
        pytest.skip("Mock engines don't support proper query operations")

    with Session(engine) as session:
        state_service = USStateSyncService(session=session)
        query_service = SQLAlchemySyncQueryService(session=session)
        fixture = open_fixture(fixture_path, USStateSyncRepository.model_type.__tablename__)  # type: ignore[has-type]
        _add_objs = state_service.create_many(
            data=[USStateStruct(**raw_obj) for raw_obj in fixture],
        )
        _ordered_objs = state_service.list(order_by=(USState.name, True))
        assert _ordered_objs[0].name == "Wyoming"
        _ordered_objs_2 = state_service.list_and_count(order_by=[(USState.name, True)])
        assert _ordered_objs_2[0][0].name == "Wyoming"
        query_count = query_service.repository.count(statement=select(StateQuery))
        assert query_count > 0
        list_query_objs, list_query_count = query_service.repository.list_and_count(
            statement=select(StateQuery),
        )
        assert list_query_count >= 50
        _paginated_objs = query_service.to_schema(
            data=list_query_objs,
            total=list_query_count,
        )

        _pydantic_paginated_objs = query_service.to_schema(
            data=list_query_objs,
            total=list_query_count,
            schema_type=StateQueryBaseModel,
        )
        assert isinstance(_pydantic_paginated_objs.items[0], StateQueryBaseModel)
        _msgspec_paginated_objs = query_service.to_schema(
            data=list_query_objs,
            total=list_query_count,
            schema_type=StateQueryStruct,
        )
        assert isinstance(_msgspec_paginated_objs.items[0], StateQueryStruct)
        _list_service_objs = query_service.repository.list(statement=select(StateQuery))
        assert len(_list_service_objs) >= 50
        _get_ones = query_service.repository.list(statement=select(StateQuery), state_name="Alabama")
        assert len(_get_ones) == 1
        _get_one = query_service.repository.get_one(statement=select(StateQuery), state_name="Alabama")
        assert _get_one.state_name == "Alabama"
        _get_one_or_none_1 = query_service.repository.get_one_or_none(
            statement=select(StateQuery).where(StateQuery.state_name == "Texas"),  # type: ignore
        )
        assert _get_one_or_none_1 is not None
        assert _get_one_or_none_1.state_name == "Texas"
        _obj = query_service.to_schema(
            data=_get_one_or_none_1,
        )
        _pydantic_obj = query_service.to_schema(
            data=_get_one_or_none_1,
            schema_type=StateQueryBaseModel,
        )
        assert isinstance(_pydantic_obj, StateQueryBaseModel)
        assert is_pydantic_model(_pydantic_obj)
        assert is_pydantic_model_with_field(_pydantic_obj, "state_abbreviation")
        assert not is_pydantic_model_without_field(_pydantic_obj, "state_abbreviation")

        _msgspec_obj = query_service.to_schema(
            data=_get_one_or_none_1,
            schema_type=StateQueryStruct,
        )
        assert isinstance(_msgspec_obj, StateQueryStruct)
        assert is_msgspec_struct(_msgspec_obj)
        assert is_msgspec_struct_with_field(_msgspec_obj, "state_abbreviation")
        assert not is_msgspec_struct_without_field(_msgspec_obj, "state_abbreviation")

        _get_one_or_none = query_service.repository.get_one_or_none(
            statement=select(StateQuery).filter_by(state_name="Nope"),
        )
        assert _get_one_or_none is None


@pytest.mark.xdist_group("sqlquery")
async def test_async_fixture_and_query(async_engine: AsyncEngine, sqlquery_test_tables_async: None) -> None:
    # Skip mock engines as they don't support proper query operations
    if getattr(async_engine.dialect, "name", "") == "mock":
        pytest.skip("Mock engines don't support proper query operations")

    async with AsyncSession(async_engine) as session:
        state_service = USStateAsyncService(session=session)

        query_service = SQLAlchemyAsyncQueryService(session=session)
        fixture = await open_fixture_async(fixture_path, USStateSyncRepository.model_type.__tablename__)
        _add_objs = await state_service.create_many(
            data=[USStateBaseModel(**raw_obj) for raw_obj in fixture],
        )
        _ordered_objs = await state_service.list(order_by=(USState.name, True))
        assert _ordered_objs[0].name == "Wyoming"
        _ordered_objs_2 = await state_service.list_and_count(order_by=(USState.name, True))
        assert _ordered_objs_2[0][0].name == "Wyoming"
        query_count = await query_service.repository.count(statement=select(StateQuery))
        assert query_count > 0
        list_query_objs, list_query_count = await query_service.repository.list_and_count(
            statement=select(StateQuery),
        )
        assert list_query_count >= 50
        _paginated_objs = query_service.to_schema(
            list_query_objs,
            total=list_query_count,
        )

        _pydantic_paginated_objs = query_service.to_schema(
            data=list_query_objs,
            total=list_query_count,
            schema_type=StateQueryBaseModel,
        )
        assert isinstance(_pydantic_paginated_objs.items[0], StateQueryBaseModel)
        _msgspec_paginated_objs = query_service.to_schema(
            data=list_query_objs,
            total=list_query_count,
            schema_type=StateQueryStruct,
        )
        assert isinstance(_msgspec_paginated_objs.items[0], StateQueryStruct)
        _list_service_objs = await query_service.repository.list(statement=select(StateQuery))
        assert len(_list_service_objs) >= 50
        _get_ones = await query_service.repository.list(statement=select(StateQuery), state_name="Alabama")
        assert len(_get_ones) == 1
        _get_one = await query_service.repository.get_one(statement=select(StateQuery), state_name="Alabama")
        assert _get_one.state_name == "Alabama"
        _get_one_or_none_1 = await query_service.repository.get_one_or_none(
            statement=select(StateQuery).where(StateQuery.state_name == "Texas"),  # type: ignore
        )
        assert _get_one_or_none_1 is not None
        assert _get_one_or_none_1.state_name == "Texas"
        _obj = query_service.to_schema(
            data=_get_one_or_none_1,
        )
        _pydantic_obj = query_service.to_schema(
            data=_get_one_or_none_1,
            schema_type=StateQueryBaseModel,
        )
        assert isinstance(_pydantic_obj, StateQueryBaseModel)
        assert is_pydantic_model(_pydantic_obj)
        assert is_pydantic_model_with_field(_pydantic_obj, "state_abbreviation")
        assert not is_pydantic_model_without_field(_pydantic_obj, "state_abbreviation")

        _msgspec_obj = query_service.to_schema(
            data=_get_one_or_none_1,
            schema_type=StateQueryStruct,
        )
        assert isinstance(_msgspec_obj, StateQueryStruct)
        assert is_msgspec_struct(_msgspec_obj)
        assert is_msgspec_struct_with_field(_msgspec_obj, "state_abbreviation")
        _get_one_or_none = await query_service.repository.get_one_or_none(
            select(StateQuery).filter_by(state_name="Nope")
        )
        assert not is_msgspec_struct_without_field(_msgspec_obj, "state_abbreviation")
        assert _get_one_or_none is None


@pytest.mark.xdist_group("sqlquery")
async def test_async_query_repository_instantiation(async_engine: AsyncEngine) -> None:
    """Test that SQLAlchemyAsyncQueryRepository can be instantiated without super().__init__() error."""
    from advanced_alchemy.repository import SQLAlchemyAsyncQueryRepository

    async with AsyncSession(async_engine) as session:
        # Test direct instantiation - this should not raise TypeError
        repository = SQLAlchemyAsyncQueryRepository(session=session)
        assert repository is not None
        assert repository.session == session
        assert repository.error_messages is None
        assert repository.wrap_exceptions is True

        # Test with optional parameters
        repository_with_params = SQLAlchemyAsyncQueryRepository(
            session=session, error_messages={"not_found": "Custom not found"}, wrap_exceptions=False
        )
        assert repository_with_params.session == session
        assert repository_with_params.error_messages == {"not_found": "Custom not found"}
        assert repository_with_params.wrap_exceptions is False


@pytest.mark.xdist_group("sqlquery")
def test_sync_query_repository_instantiation(engine: Engine) -> None:
    """Test that SQLAlchemySyncQueryRepository can be instantiated without super().__init__() error."""
    from advanced_alchemy.repository import SQLAlchemySyncQueryRepository

    with Session(engine) as session:
        # Test direct instantiation - this should not raise TypeError
        repository = SQLAlchemySyncQueryRepository(session=session)
        assert repository is not None
        assert repository.session == session
        assert repository.error_messages is None
        assert repository.wrap_exceptions is True

        # Test with optional parameters
        repository_with_params = SQLAlchemySyncQueryRepository(
            session=session, error_messages={"not_found": "Custom not found"}, wrap_exceptions=False
        )
        assert repository_with_params.session == session
        assert repository_with_params.error_messages == {"not_found": "Custom not found"}
        assert repository_with_params.wrap_exceptions is False


@pytest.mark.xdist_group("sqlquery")
async def test_async_query_service_with_repository_instantiation(async_engine: AsyncEngine) -> None:
    """Test that SQLAlchemyAsyncQueryService using the repository works correctly."""
    from advanced_alchemy.service import SQLAlchemyAsyncQueryService

    async with AsyncSession(async_engine) as session:
        # This should not raise TypeError when creating the repository internally
        query_service = SQLAlchemyAsyncQueryService(session=session)
        assert query_service is not None
        assert query_service.repository is not None
        assert query_service.repository.session == session


@pytest.mark.xdist_group("sqlquery")
def test_sync_query_service_with_repository_instantiation(engine: Engine) -> None:
    """Test that SQLAlchemySyncQueryService using the repository works correctly."""
    from advanced_alchemy.service import SQLAlchemySyncQueryService

    with Session(engine) as session:
        # This should not raise TypeError when creating the repository internally
        query_service = SQLAlchemySyncQueryService(session=session)
        assert query_service is not None
        assert query_service.repository is not None
        assert query_service.repository.session == session