# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import os
import uuid
from typing import Any, Coroutine, Sequence

import pytest
import pytest_asyncio
from langchain_core.documents import Document
from langchain_core.embeddings import DeterministicFakeEmbedding
from sqlalchemy import text
from sqlalchemy.engine.row import RowMapping

from langchain_google_cloud_sql_pg import Column, PostgresEngine
from langchain_google_cloud_sql_pg.async_vectorstore import AsyncPostgresVectorStore

DEFAULT_TABLE = "test_table" + str(uuid.uuid4())
DEFAULT_TABLE_SYNC = "test_table_sync" + str(uuid.uuid4())
CUSTOM_TABLE = "table-custom" + str(uuid.uuid4())
VECTOR_SIZE = 768

embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE)

texts = ["foo", "bar", "baz"]
metadatas = [{"page": str(i), "source": "google.com"} for i in range(len(texts))]
docs = [
    Document(page_content=texts[i], metadata=metadatas[i]) for i in range(len(texts))
]
id_column_as_metadata = [{"id": str(i)} for i in range(len(texts))]

embeddings = [embeddings_service.embed_query(texts[i]) for i in range(len(texts))]


def get_env_var(key: str, desc: str) -> str:
    v = os.environ.get(key)
    if v is None:
        raise ValueError(f"Must set env var {key} to: {desc}")
    return v


# Helper to bridge the Main Test Loop and the Engine Background Loop
async def run_on_background(engine: PostgresEngine, coro: Coroutine) -> Any:
    """Runs a coroutine on the engine's background loop."""
    if engine._loop:
        return await asyncio.wrap_future(
            asyncio.run_coroutine_threadsafe(coro, engine._loop)
        )
    return await coro


async def aexecute(engine: PostgresEngine, query: str) -> None:
    async def _impl():
        async with engine._pool.connect() as conn:
            await conn.execute(text(query))
            await conn.commit()

    # Run on background loop
    await run_on_background(engine, _impl())


async def afetch(engine: PostgresEngine, query: str) -> Sequence[RowMapping]:
    async def _impl():
        async with engine._pool.connect() as conn:
            result = await conn.execute(text(query))
            result_map = result.mappings()
            return result_map.fetchall()

    # Run on background loop
    return await run_on_background(engine, _impl())


@pytest.mark.asyncio(scope="class")
class TestVectorStore:
    @pytest.fixture(scope="module")
    def db_project(self) -> str:
        return get_env_var("PROJECT_ID", "project id for google cloud")

    @pytest.fixture(scope="module")
    def db_region(self) -> str:
        return get_env_var("REGION", "region for cloud sql instance")

    @pytest.fixture(scope="module")
    def db_instance(self) -> str:
        return get_env_var("INSTANCE_ID", "instance for cloud sql")

    @pytest.fixture(scope="module")
    def db_name(self) -> str:
        return get_env_var("DATABASE_ID", "database name on cloud sql instance")

    @pytest_asyncio.fixture(scope="class")
    async def engine(self, db_project, db_region, db_instance, db_name):
        engine = await PostgresEngine.afrom_instance(
            project_id=db_project,
            instance=db_instance,
            region=db_region,
            database=db_name,
        )

        yield engine
        await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"')
        await aexecute(engine, f'DROP TABLE IF EXISTS "{CUSTOM_TABLE}"')
        await engine.close()

    @pytest_asyncio.fixture(scope="class")
    async def vs(self, engine):
        # Wrap private init method
        await run_on_background(
            engine, engine._ainit_vectorstore_table(DEFAULT_TABLE, VECTOR_SIZE)
        )
        # Wrap creation of the async vectorstore
        vs = await run_on_background(
            engine,
            AsyncPostgresVectorStore.create(
                engine,
                embedding_service=embeddings_service,
                table_name=DEFAULT_TABLE,
            ),
        )
        yield vs

    @pytest_asyncio.fixture(scope="class")
    async def vs_custom(self, engine):
        # Wrap private init method
        await run_on_background(
            engine,
            engine._ainit_vectorstore_table(
                CUSTOM_TABLE,
                VECTOR_SIZE,
                id_column="myid",
                content_column="mycontent",
                embedding_column="myembedding",
                metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],
                metadata_json_column="mymeta",
            ),
        )

        # Wrap creation of the async vectorstore
        vs = await run_on_background(
            engine,
            AsyncPostgresVectorStore.create(
                engine,
                embedding_service=embeddings_service,
                table_name=CUSTOM_TABLE,
                id_column="myid",
                content_column="mycontent",
                embedding_column="myembedding",
                metadata_columns=["page", "source"],
                metadata_json_column="mymeta",
            ),
        )
        yield vs

    async def test_init_with_constructor(self, engine):
        with pytest.raises(Exception):
            AsyncPostgresVectorStore(
                engine,
                embedding_service=embeddings_service,
                table_name=CUSTOM_TABLE,
                id_column="myid",
                content_column="noname",
                embedding_column="myembedding",
                metadata_columns=["page", "source"],
                metadata_json_column="mymeta",
            )

    async def test_post_init(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="noname",
                    embedding_column="myembedding",
                    metadata_columns=["page", "source"],
                    metadata_json_column="mymeta",
                ),
            )

    async def test_id_metadata_column(self, engine):
        table_name = "id_metadata" + str(uuid.uuid4())
        await run_on_background(
            engine,
            engine._ainit_vectorstore_table(
                table_name,
                VECTOR_SIZE,
                metadata_columns=[Column("id", "TEXT")],
            ),
        )
        custom_vs = await run_on_background(
            engine,
            AsyncPostgresVectorStore.create(
                engine,
                embedding_service=embeddings_service,
                table_name=table_name,
                metadata_columns=["id"],
            ),
        )
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        # Wrap aadd_texts
        await run_on_background(
            engine, custom_vs.aadd_texts(texts, id_column_as_metadata, ids)
        )

        results = await afetch(engine, f'SELECT * FROM "{table_name}"')
        assert len(results) == 3
        assert results[0]["id"] == "0"
        assert results[1]["id"] == "1"
        assert results[2]["id"] == "2"
        await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}"')

    async def test_aadd_texts(self, engine, vs):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        # Wrap aadd_texts
        await run_on_background(engine, vs.aadd_texts(texts, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 3

        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        # Wrap aadd_texts
        await run_on_background(engine, vs.aadd_texts(texts, metadatas, ids))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 6
        await aexecute(engine, f'TRUNCATE TABLE "{DEFAULT_TABLE}"')

    async def test_aadd_texts_edge_cases(self, engine, vs):
        texts = ["Taylor's", '"Swift"', "best-friend"]
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        # Wrap aadd_texts
        await run_on_background(engine, vs.aadd_texts(texts, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 3
        await aexecute(engine, f'TRUNCATE TABLE "{DEFAULT_TABLE}"')

    async def test_aadd_docs(self, engine, vs):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        # Wrap aadd_documents
        await run_on_background(engine, vs.aadd_documents(docs, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 3
        await aexecute(engine, f'TRUNCATE TABLE "{DEFAULT_TABLE}"')

    async def test_aadd_docs_no_ids(self, engine, vs):
        # Wrap aadd_documents
        await run_on_background(engine, vs.aadd_documents(docs))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 3
        await aexecute(engine, f'TRUNCATE TABLE "{DEFAULT_TABLE}"')

    async def test_adelete(self, engine, vs):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        await run_on_background(engine, vs.aadd_texts(texts, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 3
        await run_on_background(engine, vs.adelete([ids[0]]))
        results = await afetch(engine, f'SELECT * FROM "{DEFAULT_TABLE}"')
        assert len(results) == 2
        result = await run_on_background(engine, vs.adelete())
        assert result == False

    ##### Custom Vector Store  #####

    async def test_aadd_texts_custom(self, engine, vs_custom):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        await run_on_background(engine, vs_custom.aadd_texts(texts, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{CUSTOM_TABLE}"')
        assert len(results) == 3
        assert results[0]["mycontent"] == "foo"
        assert results[0]["myembedding"]
        assert results[0]["page"] is None
        assert results[0]["source"] is None

        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        await run_on_background(engine, vs_custom.aadd_texts(texts, metadatas, ids))
        results = await afetch(engine, f'SELECT * FROM "{CUSTOM_TABLE}"')
        assert len(results) == 6
        await aexecute(engine, f'TRUNCATE TABLE "{CUSTOM_TABLE}"')

    async def test_aadd_docs_custom(self, engine, vs_custom):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        docs = [
            Document(
                page_content=texts[i],
                metadata={"page": str(i), "source": "google.com"},
            )
            for i in range(len(texts))
        ]
        await run_on_background(engine, vs_custom.aadd_documents(docs, ids=ids))

        results = await afetch(engine, f'SELECT * FROM "{CUSTOM_TABLE}"')
        assert len(results) == 3
        assert results[0]["mycontent"] == "foo"
        assert results[0]["myembedding"]
        assert results[0]["page"] == "0"
        assert results[0]["source"] == "google.com"
        await aexecute(engine, f'TRUNCATE TABLE "{CUSTOM_TABLE}"')

    async def test_adelete_custom(self, engine, vs_custom):
        ids = [str(uuid.uuid4()) for i in range(len(texts))]
        await run_on_background(engine, vs_custom.aadd_texts(texts, ids=ids))
        results = await afetch(engine, f'SELECT * FROM "{CUSTOM_TABLE}"')
        content = [result["mycontent"] for result in results]
        assert len(results) == 3
        assert "foo" in content
        await run_on_background(engine, vs_custom.adelete([ids[0]]))
        results = await afetch(engine, f'SELECT * FROM "{CUSTOM_TABLE}"')
        content = [result["mycontent"] for result in results]
        assert len(results) == 2
        assert "foo" not in content

    async def test_ignore_metadata_columns(self, engine):
        column_to_ignore = "source"
        vs = await run_on_background(
            engine,
            AsyncPostgresVectorStore.create(
                engine,
                embedding_service=embeddings_service,
                table_name=CUSTOM_TABLE,
                ignore_metadata_columns=[column_to_ignore],
                id_column="myid",
                content_column="mycontent",
                embedding_column="myembedding",
                metadata_json_column="mymeta",
            ),
        )
        assert column_to_ignore not in vs.metadata_columns

    async def test_create_vectorstore_with_invalid_parameters_1(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="mycontent",
                    embedding_column="myembedding",
                    metadata_columns=["random_column"],  # invalid metadata column
                ),
            )

    async def test_create_vectorstore_with_invalid_parameters_2(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="langchain_id",  # invalid content column type
                    embedding_column="myembedding",
                    metadata_columns=["random_column"],
                ),
            )

    async def test_create_vectorstore_with_invalid_parameters_3(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="mycontent",
                    embedding_column="random_column",  # invalid embedding column
                    metadata_columns=["random_column"],
                ),
            )

    async def test_create_vectorstore_with_invalid_parameters_4(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="mycontent",
                    embedding_column="langchain_id",  # invalid embedding column data type
                    metadata_columns=["random_column"],
                ),
            )

    async def test_create_vectorstore_with_invalid_parameters_5(self, engine):
        with pytest.raises(ValueError):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore.create(
                    engine,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="mycontent",
                    embedding_column="langchain_id",
                    metadata_columns=["random_column"],
                    ignore_metadata_columns=[
                        "one",
                        "two",
                    ],  # invalid use of metadata_columns and ignore columns
                ),
            )

    async def test_create_vectorstore_with_init(self, engine):
        with pytest.raises(Exception):
            await run_on_background(
                engine,
                AsyncPostgresVectorStore(
                    engine._pool,
                    embedding_service=embeddings_service,
                    table_name=CUSTOM_TABLE,
                    id_column="myid",
                    content_column="mycontent",
                    embedding_column="myembedding",
                    metadata_columns=["random_column"],  # invalid metadata column
                ),
            )
