import os

import pytest
import ydb
from ydb import issues


def pytest_addoption(parser):
    """Add custom command line options for pytest-docker compatibility."""
    parser.addoption(
        "--docker-compose-file",
        action="store",
        default="docker-compose.yml",
        help="Path to docker-compose file (relative to project root)",
    )


@pytest.fixture(scope="session")
def docker_compose_file(pytestconfig):
    """Path to docker-compose.yml for pytest-docker."""
    compose_file = pytestconfig.getoption("--docker-compose-file")
    return os.path.join(str(pytestconfig.rootdir), compose_file)


@pytest.fixture(scope="session")
def docker_cleanup():
    """Remove volumes after tests (equivalent to --docker-compose-remove-volumes)."""
    return ["down -v --remove-orphans"]


class DockerProject:
    """Compatibility wrapper for pytest-docker-compose docker_project fixture."""

    def __init__(self, docker_compose, docker_services, endpoint):
        self._docker_compose = docker_compose
        self._docker_services = docker_services
        self._endpoint = endpoint
        self._stopped = False

    def stop(self):
        """Stop all containers (marks as stopped, actual restart happens in start())."""
        self._stopped = True
        # Use 'kill' for instant stop (simulates network failure better than graceful stop)
        self._docker_compose.execute("kill")

    def start(self):
        """Restart containers and wait for YDB to be ready."""
        if self._stopped:
            # After kill, we need to recreate the container to restore YDB properly
            self._docker_compose.execute("up -d --force-recreate")
            self._stopped = False
        else:
            self._docker_compose.execute("start")

        self._docker_services.wait_until_responsive(
            timeout=120.0,
            pause=2.0,
            check=lambda: is_ydb_responsive(self._endpoint),
        )


@pytest.fixture(scope="module")
def docker_project(docker_services, endpoint):
    """Compatibility fixture providing stop/start methods like pytest-docker-compose."""
    return DockerProject(docker_services._docker_compose, docker_services, endpoint)


def is_ydb_responsive(endpoint):
    """Check if YDB is ready to accept connections."""
    try:
        with ydb.Driver(endpoint=endpoint, database="/local") as driver:
            driver.wait(timeout=5)
            with ydb.SessionPool(driver) as pool:
                with pool.checkout() as session:
                    session.execute_scheme("create table `.sys_health/test_table` (A int32, primary key(A));")
            return True
    except Exception:
        return False


def is_ydb_secure_responsive(endpoint, root_certificates):
    """Check if YDB TLS endpoint is ready to accept connections."""
    try:
        with ydb.Driver(
            endpoint=f"grpcs://{endpoint}",
            database="/local",
            root_certificates=root_certificates,
        ) as driver:
            driver.wait(timeout=5)
            with ydb.SessionPool(driver) as pool:
                with pool.checkout() as session:
                    session.execute_scheme("create table `.sys_health/test_table_tls` (A int32, primary key(A));")
            return True
    except Exception:
        return False


@pytest.fixture(scope="module")
def endpoint(docker_ip, docker_services):
    """Wait for YDB to be responsive and return endpoint."""
    port = docker_services.port_for("ydb", 2136)
    endpoint_url = f"{docker_ip}:{port}"
    docker_services.wait_until_responsive(
        timeout=60.0,
        pause=1.0,
        check=lambda: is_ydb_responsive(endpoint_url),
    )
    yield endpoint_url


@pytest.fixture(scope="session")
def secure_endpoint(pytestconfig, docker_ip, docker_services):
    """Wait for YDB TLS endpoint to be responsive."""
    ca_path = os.path.join(str(pytestconfig.rootdir), "ydb_certs/ca.pem")

    # Wait for certificate file to be generated by container
    def wait_for_certificate():
        return os.path.exists(ca_path)

    docker_services.wait_until_responsive(
        timeout=30.0,
        pause=1.0,
        check=wait_for_certificate,
    )

    assert os.path.exists(ca_path), f"Certificate not found at {ca_path}"
    os.environ["YDB_SSL_ROOT_CERTIFICATES_FILE"] = ca_path
    root_certificates = ydb.load_ydb_root_certificate()

    port = docker_services.port_for("ydb", 2135)
    # Use 'localhost' instead of docker_ip because SSL certificate is issued for 'localhost'
    endpoint_url = f"localhost:{port}"

    docker_services.wait_until_responsive(
        timeout=60.0,
        pause=1.0,
        check=lambda: is_ydb_secure_responsive(endpoint_url, root_certificates),
    )

    yield endpoint_url


@pytest.fixture(scope="module")
def database():
    return "/local"


@pytest.fixture()
async def aio_connection(endpoint, database):
    """A fixture to wait ydb start"""
    from ydb.aio.connection import Connection
    from ydb.driver import DriverConfig

    config = DriverConfig.default_from_endpoint_and_database(endpoint, database)
    connection = Connection(endpoint, config)
    await connection.connection_ready(ready_timeout=7)
    return connection


@pytest.fixture()
async def driver(endpoint, database, event_loop):
    driver_config = ydb.DriverConfig(
        endpoint,
        database,
    )

    driver = ydb.aio.Driver(driver_config=driver_config)
    await driver.wait(timeout=15)

    yield driver

    await driver.stop(timeout=10)


@pytest.fixture()
async def driver_sync(endpoint, database, event_loop):
    driver_config = ydb.DriverConfig(
        endpoint,
        database,
    )

    driver = ydb.Driver(driver_config=driver_config)
    driver.wait(timeout=15)

    yield driver

    driver.stop(timeout=10)


@pytest.fixture()
def table_name(driver_sync, database):
    table_name = "table"

    with ydb.SessionPool(driver_sync) as pool:

        def create_table(s):
            try:
                s.drop_table(database + "/" + table_name)
            except ydb.SchemeError:
                pass

            s.execute_scheme(
                """
CREATE TABLE %s (
id Int64 NOT NULL,
i64Val Int64,
PRIMARY KEY(id)
)
"""
                % table_name
            )

        pool.retry_operation_sync(create_table)
    return table_name


@pytest.fixture()
def table_path(database, table_name) -> str:
    return database + "/" + table_name


@pytest.fixture
def column_table_name(driver_sync, database):
    table_name = "column_table"

    with ydb.SessionPool(driver_sync) as pool:

        def create_table(s):
            try:
                s.drop_table(database + "/" + table_name)
            except ydb.SchemeError:
                pass

            s.execute_scheme(
                """
CREATE TABLE %s (
id Int64 NOT NULL,
i64Val Int64,
PRIMARY KEY(id)
)
PARTITION BY HASH(id)
WITH (
    STORE = COLUMN
)
"""
                % table_name
            )

        pool.retry_operation_sync(create_table)
    return table_name


@pytest.fixture()
def column_table_path(database, column_table_name) -> str:
    return database + "/" + column_table_name


@pytest.fixture()
def topic_consumer():
    return "fixture-consumer"


@pytest.fixture()
@pytest.mark.asyncio()
async def topic_path(driver, topic_consumer, database) -> str:
    topic_path = database + "/test-topic"

    try:
        await driver.topic_client.drop_topic(topic_path)
    except issues.SchemeError:
        pass

    await driver.topic_client.create_topic(
        path=topic_path,
        consumers=[topic_consumer],
    )

    return topic_path


@pytest.fixture()
@pytest.mark.asyncio()
async def topic2_path(driver, topic_consumer, database) -> str:
    topic_path = database + "/test-topic2"

    try:
        await driver.topic_client.drop_topic(topic_path)
    except issues.SchemeError:
        pass

    await driver.topic_client.create_topic(
        path=topic_path,
        consumers=[topic_consumer],
    )

    return topic_path


@pytest.fixture()
@pytest.mark.asyncio()
async def topic_with_two_partitions_path(driver, topic_consumer, database) -> str:
    topic_path = database + "/test-topic-two-partitions"

    try:
        await driver.topic_client.drop_topic(topic_path)
    except issues.SchemeError:
        pass

    await driver.topic_client.create_topic(
        path=topic_path,
        consumers=[topic_consumer],
        min_active_partitions=2,
        partition_count_limit=2,
    )

    return topic_path


@pytest.fixture()
@pytest.mark.asyncio()
async def topic_with_messages(driver, topic_consumer, database):
    topic_path = database + "/test-topic-with-messages"
    try:
        await driver.topic_client.drop_topic(topic_path)
    except issues.SchemeError:
        pass

    await driver.topic_client.create_topic(
        path=topic_path,
        consumers=[topic_consumer],
    )

    writer = driver.topic_client.writer(topic_path, producer_id="fixture-producer-id", codec=ydb.TopicCodec.RAW)
    await writer.write_with_ack(
        [
            ydb.TopicWriterMessage(data="123".encode()),
            ydb.TopicWriterMessage(data="456".encode()),
        ]
    )
    await writer.write_with_ack(
        [
            ydb.TopicWriterMessage(data="789".encode()),
            ydb.TopicWriterMessage(data="0".encode()),
        ]
    )
    await writer.close()
    return topic_path


@pytest.fixture()
@pytest.mark.asyncio()
async def topic_with_messages_with_metadata(driver, topic_consumer, database):
    topic_path = database + "/test-topic-with-messages-with-metadata"
    try:
        await driver.topic_client.drop_topic(topic_path)
    except issues.SchemeError:
        pass

    await driver.topic_client.create_topic(
        path=topic_path,
        consumers=[topic_consumer],
    )

    writer = driver.topic_client.writer(topic_path, producer_id="fixture-producer-id", codec=ydb.TopicCodec.RAW)
    await writer.write_with_ack(
        [
            ydb.TopicWriterMessage(data="123".encode(), metadata_items={"key": "value"}),
            ydb.TopicWriterMessage(data="456".encode(), metadata_items={"key": b"value"}),
        ]
    )
    await writer.close()
    return topic_path


@pytest.fixture()
@pytest.mark.asyncio()
async def topic_reader(driver, topic_consumer, topic_path) -> ydb.TopicReaderAsyncIO:
    reader = driver.topic_client.reader(topic=topic_path, consumer=topic_consumer)
    yield reader
    await reader.close()


@pytest.fixture()
def topic_reader_sync(driver_sync, topic_consumer, topic_path) -> ydb.TopicReader:
    reader = driver_sync.topic_client.reader(topic=topic_path, consumer=topic_consumer)
    yield reader
    reader.close()
