# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors

import os
import shutil
import uuid
from pathlib import Path

import lancedb
import pytest
import pytest_asyncio

# What the quickstart's download step leaves in an application's directory.
IMAGE = (
    Path(__file__).resolve().parents[2]
    / "web/static/assets/images/quickstart/sir-lancelot.jpg"
)
DEPLOYMENT = ("LANCEDB_URI", "LANCEDB_HOST_OVERRIDE", "LANCEDB_API_KEY")
# The connection settings alone never send these tests to a deployment.
OPT_IN = "LANCEDB_DOCS_TEST_DEPLOYMENT"


def deployment_settings():
    if os.environ.get(OPT_IN) != "1":
        pytest.skip(f"set {OPT_IN}=1 and {', '.join(DEPLOYMENT)} to use a deployment")
    missing = [name for name in DEPLOYMENT if not os.environ.get(name)]
    if missing:
        pytest.fail(f"{OPT_IN}=1 needs {', '.join(missing)}")
    return dict(
        uri=os.environ["LANCEDB_URI"],
        api_key=os.environ["LANCEDB_API_KEY"],
        region=os.environ.get("LANCEDB_REGION", "us-east-1"),
        host_override=os.environ["LANCEDB_HOST_OVERRIDE"],
    )


class OwnedTables:
    """A deployment connection on which a test creates only tables of its own.

    Each name gets a suffix unique to the test and is claimed only while no
    table has it; the test drops what it claimed, whether or not it passed.
    """

    def __init__(self, db):
        self.db = db
        self.suffix = uuid.uuid4().hex[:12]
        self.claimed = []

    def claim(self, name, existing, mode):
        owned = f"{name}_{self.suffix}"
        if owned in self.claimed:
            return owned, mode
        if owned in existing:
            raise RuntimeError(f"{owned} already exists and is not this test's")
        self.claimed.append(owned)
        return owned, "create"

    def table_names(self):
        names, token = [], None
        while True:
            page = self.db.list_tables(page_token=token)
            names += page.tables
            if not (token := page.page_token):
                return names

    def create_table(self, name, *args, mode="create", **kwargs):
        owned, mode = self.claim(name, self.table_names(), mode)
        return self.db.create_table(owned, *args, mode=mode, **kwargs)

    def drop_claimed(self):
        existing = self.table_names()
        for name in self.claimed:
            if name in existing:
                self.db.drop_table(name)


class AsyncOwnedTables(OwnedTables):
    async def table_names(self):
        names, token = [], None
        while True:
            page = await self.db.list_tables(page_token=token)
            names += page.tables
            if not (token := page.page_token):
                return names

    async def create_table(self, name, *args, mode="create", **kwargs):
        owned, mode = self.claim(name, await self.table_names(), mode)
        return await self.db.create_table(owned, *args, mode=mode, **kwargs)

    async def drop_claimed(self):
        existing = await self.table_names()
        for name in self.claimed:
            if name in existing:
                await self.db.drop_table(name)


@pytest.fixture
def deployment():
    owned = OwnedTables(lancedb.connect(**deployment_settings()))
    yield owned
    owned.drop_claimed()


@pytest.fixture(params=["local", "enterprise"])
def db(request, db_path_factory):
    if request.param == "local":
        yield lancedb.connect(db_path_factory("quickstart_db"))
    else:
        yield request.getfixturevalue("deployment")


@pytest_asyncio.fixture(params=["local", "enterprise"])
async def async_db(request, db_path_factory):
    if request.param == "local":
        yield await lancedb.connect_async(db_path_factory("quickstart_async_db"))
        return
    owned = AsyncOwnedTables(await lancedb.connect_async(**deployment_settings()))
    yield owned
    await owned.drop_claimed()


def create_frames(db):
    # --8<-- [start:quickstart_data]
    data = [
        {
            "frame_id": "1",
            "video_id": "clip-1",
            "caption": "a robot arm reaching toward a cup",
            "stats": {"motion": 3, "sharpness": 4, "exposure": 4, "framing": 3},
            "embedding": [0.7, 0.1, 0.9, 0.7],
        },
        {
            "frame_id": "2",
            "video_id": "clip-1",
            "caption": "a two-finger gripper closing on a cup",
            "stats": {"motion": 5, "sharpness": 4, "exposure": 5, "framing": 2},
            "embedding": [0.2, 0.9, 0.4, 0.9],
        },
        {
            "frame_id": "3",
            "video_id": "clip-2",
            "caption": "a robot arm holding a cup steady",
            "stats": {"motion": 2, "sharpness": 3, "exposure": 3, "framing": 4},
            "embedding": [0.9, 0.1, 0.5, 0.4],
        },
    ]
    # --8<-- [end:quickstart_data]

    # --8<-- [start:quickstart_create_table]
    table = db.create_table("frames", data=data, mode="overwrite")
    # --8<-- [end:quickstart_create_table]
    assert table.count_rows() == 3

    # --8<-- [start:quickstart_add_data]
    table.add(
        [
            {
                "frame_id": "4",
                "video_id": "clip-2",
                "caption": "a gripper lowering a cup onto the table",
                "stats": {"motion": 5, "sharpness": 4, "exposure": 4, "framing": 2},
                "embedding": [0.3, 0.9, 0.6, 0.8],
            }
        ]
    )
    # --8<-- [end:quickstart_add_data]
    assert table.count_rows() == 4
    return table


def test_quickstart(db):
    table = create_frames(db)

    # --8<-- [start:quickstart_vector_search_1]
    # Search for frames similar to "a gripper closing on an object"
    query_vector = [0.2, 0.8, 0.4, 0.9]

    # Ensure you run `pip install polars` beforehand
    result = (
        table.search(query_vector)
        .select(["frame_id", "video_id", "caption", "_distance"])
        .limit(2)
        .to_polars()
    )
    print(result)
    # --8<-- [end:quickstart_vector_search_1]
    assert result["frame_id"].to_list() == ["2", "4"]

    # --8<-- [start:quickstart_output_pandas]
    # Ensure you run `pip install pandas` beforehand
    result = table.search(query_vector).limit(2).to_pandas()
    print(result)
    # --8<-- [end:quickstart_output_pandas]
    assert result["frame_id"].to_list() == ["2", "4"]

    # --8<-- [start:quickstart_curate_with_metadata]
    curated = (
        table.search(query_vector)
        .where("stats.motion >= 4")
        .select(["frame_id", "video_id", "caption", "_distance"])
        .limit(2)
        .to_polars()
    )
    print(curated)
    # --8<-- [end:quickstart_curate_with_metadata]
    assert curated["frame_id"].to_list() == ["2", "4"]

    curated = table.search(query_vector).where("stats.motion < 4").limit(2).to_list()
    assert [row["frame_id"] for row in curated] == ["1", "3"]


@pytest.mark.asyncio
async def test_quickstart_async_api(async_db):
    data = [
        {
            "frame_id": "2",
            "video_id": "clip-1",
            "caption": "a two-finger gripper closing on a cup",
            "stats": {"motion": 5, "sharpness": 4, "exposure": 5, "framing": 2},
            "embedding": [0.2, 0.9, 0.4, 0.9],
        }
    ]

    # --8<-- [start:quickstart_create_table_async]
    async_table = await async_db.create_table(
        "frames",
        data=data,
        mode="overwrite",
    )
    # --8<-- [end:quickstart_create_table_async]

    # --8<-- [start:quickstart_vector_search_1_async]
    # Search for frames similar to "a gripper closing on an object"
    query_vector = [0.2, 0.8, 0.4, 0.9]

    # Ensure you run `pip install polars` beforehand
    async_result = await (
        await async_table.search(query_vector)
    ).select(["frame_id", "video_id", "caption", "_distance"]).limit(2).to_polars()
    print(async_result)
    # --8<-- [end:quickstart_vector_search_1_async]
    assert async_result["frame_id"].to_list() == ["2"]


def test_quickstart_derived_feature(db):
    table = create_frames(db)

    # --8<-- [start:quickstart_add_feature]
    table.add_columns(
        {
            "quality_score": "cast(((stats.motion + stats.sharpness + stats.exposure + stats.framing) / 4.0) as float)"
        }
    )
    # --8<-- [end:quickstart_add_feature]

    # --8<-- [start:quickstart_query_feature]
    features = table.search().select(["frame_id", "video_id", "quality_score"]).to_polars()
    print(features)
    # --8<-- [end:quickstart_query_feature]
    assert sorted(zip(features["frame_id"], features["quality_score"])) == [
        ("1", 3.5),
        ("2", 4.0),
        ("3", 3.0),
        ("4", 3.75),
    ]


def caption_lengths(table):
    rows = table.search().select(["caption", "caption_length"]).to_list()
    return {row["caption"]: row["caption_length"] for row in rows}


def test_quickstart_computed_column(db_path_factory):
    table = create_frames(lancedb.connect(db_path_factory("computed")))

    # --8<-- [start:quickstart_computed_feature]
    table.add_columns(computed={"caption_length": "length(caption)"})
    table.refresh_column("caption_length")
    # --8<-- [end:quickstart_computed_feature]
    lengths = caption_lengths(table)
    assert lengths == {text: len(text) for text in lengths}


def test_quickstart_computed_column_on_a_deployment(deployment):
    table = create_frames(deployment)

    # --8<-- [start:quickstart_computed_feature_enterprise]
    table.add_columns(computed={"caption_length": "length(caption)"})
    job = table.refresh_column_async("caption_length")
    job.wait()
    # --8<-- [end:quickstart_computed_feature_enterprise]
    lengths = caption_lengths(table)
    assert len(lengths) == 4
    assert lengths == {text: len(text) for text in lengths}
    with pytest.raises(Exception):
        table.refresh_column("caption_length")


def test_quickstart_multimodal(db, tmp_path, monkeypatch):
    app = tmp_path / "app"
    (app / "data").mkdir(parents=True)
    shutil.copy(IMAGE, app / "data/sample-frame.jpg")
    monkeypatch.chdir(app)

    # --8<-- [start:quickstart_multimodal_bytes]
    from pathlib import Path

    image_bytes = Path("data/sample-frame.jpg").read_bytes()

    multimodal_table = db.create_table(
        "frame_images",
        data=[
            {
                "frame_id": "thumb-1",
                "caption": "Sample frame thumbnail",
                "image": image_bytes,
                "embedding": [0.9, 0.1, 0.5, 0.4],
            }
        ],
        mode="overwrite",
    )
    # --8<-- [end:quickstart_multimodal_bytes]
    assert multimodal_table.count_rows() == 1
    stored = multimodal_table.search().select(["image"]).to_list()[0]["image"]
    assert stored == IMAGE.read_bytes()


def test_quickstart_training(db):
    table = create_frames(db)

    # --8<-- [start:quickstart_training]
    import torch
    from lancedb.permutation import Permutation

    table.add_columns(
        {
            "motion": "cast(stats.motion as float)",
            "sharpness": "cast(stats.sharpness as float)",
        }
    )
    dataset = Permutation.identity(table).select_columns(["motion", "sharpness"])
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1024, shuffle=True)
    # --8<-- [end:quickstart_training]
    batches = list(dataloader)
    assert len(batches) == 1
    assert set(batches[0]) == {"motion", "sharpness"}
    assert all(isinstance(value, torch.Tensor) for value in batches[0].values())
    assert sorted(batches[0]["motion"].tolist()) == [2.0, 3.0, 5.0, 5.0]