import asyncio
from datasette_test import Datasette
import os
import pytest
import sqlite_utils


async def _cookies(datasette, path="/-/enrich/data/items/embeddings"):
    cookies = {"ds_actor": datasette.sign({"a": {"id": "root"}}, "actor")}
    csrftoken = (await datasette.client.get(path, cookies=cookies)).cookies[
        "ds_csrftoken"
    ]
    cookies["ds_csrftoken"] = csrftoken
    return cookies


@pytest.mark.vcr(ignore_localhost=True)
@pytest.mark.asyncio
async def test_enrichment(tmpdir):
    data = str(tmpdir / "data.db")
    datasette = Datasette([data])
    db = sqlite_utils.Database(
        data,
    )
    rows = [
        {"id": 1, "name": "One", "description": "First item"},
        {"id": 2, "name": "Two", "description": "Second item"},
        {"id": 3, "name": "Three", "description": "Third item"},
    ]
    db["items"].insert_all(rows, pk="id")

    api_key = os.environ.get("OPENAI_API_KEY") or "sk-mock-api-key"

    cookies = await _cookies(datasette)
    post = {
        "model": "text-embedding-3-large-256",
        "template": "{{ name }} {{ description }}",
        "api_key": api_key,
    }
    post["csrftoken"] = cookies["ds_csrftoken"]
    response = await datasette.client.post(
        "/-/enrich/data/items/embeddings",
        data=post,
        cookies=cookies,
    )
    assert response.status_code == 302
    # Poll for completion
    db = datasette.get_database("data")
    attempts = 0
    while True:
        jobs = await db.execute("select * from _enrichment_jobs")
        job = dict(jobs.first())
        if job["status"] != "finished":
            await asyncio.sleep(0.3)
            attempts += 1
            assert attempts < 1000
        else:
            break

    errors = [
        dict(row) for row in (await db.execute("select * from _enrichment_errors")).rows
    ]
    assert not errors

    assert job["status"] == "finished"
    assert job["enrichment"] == "embeddings"
    assert job["done_count"] == 3
    assert job["error_count"] == 0

    results = await db.execute("select * from items order by id")
    rows = [dict(r) for r in results.rows]
    assert rows == [
        {
            "id": 1,
            "name": "One",
            "description": "First item",
            "emb_text_embedding_3_large_256": b'\x90\xbfw\xbc\xe7_\xe5<+\x82Z\xbc\xcb\xa2p\xb9\xed\xa9\xbd\xbd\x1e\xc4\x0c\xbeG?O\xbcD\x905>\x0f\xb1\x8a=\x03\xfa\xd4;\xa5\xa6\xff;\xd6J|<\xcb;\xc8<\x92\xf1&\xbd\xd3\x87\x87=\xe4&\x9e;\x81R\x10=6*f=t\xee\xa7\xbd[V\x9f\xbd>#X<R:(=#\x98\x12\xbd\xf7\x0b?=\x8dc\xab<\x9d2\x8a<\xab\xf0W\xbd\x8c\xa7N<\x1d~\x82\xbc\x1eN\xba\xbch\xdd\x8c<Lz}=\x02\xb4J\xbd\x06\x1f\xc1\xbd\xfa\xc9\x02\xbd?\xbc\xaf\xbbO\x8b\x8e=\xf0\xce)\xbd\xf8tN=\x00K;=\xaa\xcdR\xbbG\x1cJ\xbd\x9eU\x0f\xbe\xde\x0c\x8e=\xe9\xeby=\xf7\xa4\x96<\xfa\xdd\xdd\xbbb\x93\xb4\xbb\xcb\xb1\x9a\xbd\x1f-\x9c<\xfb\xec\x87=~-\xa4;\xd3\xee\xaf\xbc>\x00S\xbb\x80/\x8b\xbc\xf0\xce\xa9:XT\xb8<\xac"\x07>\x85\x8d>\xbc\xfb#h\xbdTP\xea\xbc\xae8\xc9\xbd\x99\xa4\x8e\xbd\xb1\x80:\xbdx6\x99\xbd\x98\x0b7=\xe1\xde,:\xbc{\x93=\xd2\xa8\xa5\xbd_nH=\xdds\xb6\xbdr\x99\xf3\xbc\xf3M{=^(\xbe=d\xfcC=\xba\x12\x84\xbdO\x8b\x0e>\xd6J|=\xe8n\x0f\xbc\x8b\xd7\x16=#\xff\xba\xbb\xa3\x1ak\xbb\x9cv-9\x03p\'\xbc\xcb\xa2\xf0;\x1d~\x82\xbc\x11`\xa4\xbd\x91\xe2|\xbez[\x85<\x10\x1a\x1a\xbcLf"<\xce=\xaf=\xad|\xec\xbc\xe5\xd3P=:\x81\x01=4\x9e\xd1<\xd3d\x82=1e\n=\xde\x96;\xbcS\x802=E)\r\xbd\x84\xd1a\xbcD\x90\xb5\xbd\x01\x07\x18\xbe\x82EM\xbc\xf8tN=S\n\xe0=\x1c\xe5\xaa=\xcf\x83\xb9\xbd\x1f\xb7\xc9\xbdI\x1e\xb1;\x9d\xbc\xb7<\xbdnP\xbd+\xa5_=;\xc7\x8b=\xbb5\x89\xbd\xf5\x18\x82\xbd\x9cS\xa8<.\xa7\xc6;F\xd6\xbf\xbb<\xea\x10=\x06\x95\x93\xbd\x06\xb8\x98\xbd\xc4\xfe2<\x82EM\xbcFo\x97\xbcl\xf55\xbdH\x85Y=i\xadD=\xd7|+\xbc\xdc\xb7\xd9\xbc\xcf\x1c\x11=\x1eN\xba\xbd{\x92e\xbcg\xce\xe2=\xad|l=\xc1ct=*_U\xbd\x19\x13\x8c\xba\x96\xe6J<\x9e\x02\xc2\xbdI\xcb\xe3\xbd\xbb\xe2;=\xbf\xb4\xda=Q\xa1P\xbd#\xac\xed=\xa3\x06\x10\xbd\x91\xce!=\xfaS0\xbd)\xf6E=9\xe8)\xbc\x19\'\xe7\xbd\xedB\x95\xbd\x1aY\x96<\x16\xee\x1f\xbay\x06\xd1<\x8e\xfc\x82<7Mk<G\xb5!=\xcak\x10\xbdz\xc2-;Q~\xcb=|\xfb\xf4\xbc]\xe2\xb3\xbdx\xc0\xc6\xbc:\x81\x81=\\I\\=\x11=\x9f\xbb\xb9V\'\xbe \xda\xce=Hb\xd4\xbc\xc4\xdb-\xbcj\x8c\xa6=z8\x00\xbc\xad|l=\x05\xd96<\xedB\x15>`*%\xbcx\xe3\xcb\xbb\xd2\xa8%\xbd\x00(6=R\xe7\xda=\xf0g\x01>\xfe\xbe&=-\xeb\xe9\xbd\x0e\x183\xbd\x86\xb0\xc3\xbc?\xbc\xaf=\xe0\x98"\xbd\xacYg\xbdm\x8e\x8d=\x06\x1f\xc1\xbd\xf2Z\xbe=\x85\x17\xec=\xae\xc2\xf6=8\x7f\x1a=\x05r\x8e\xbd\xea\x0e\x7f<\xfa0\xab=\x0f;8=\x05cd=\'\xe0\x03\xbcF\xd6\xbf=6*\xe6=\x8b\xb4\x91\xbdK4\xf3\xbat\xdf\xfd=\x1d~\x82=\x08!(=\x05\xd96=\xd1\x85\xa0<?2\x82\xbd\x86:q\xbc\xacY\xe7<\xb6\x98c\xba|N\xc2=xY\x9e\xbd\x9b\r\x9e=\xad|l;\x19\x9d\xb9;6\xf3\x85=\xe3\xe0\x93;\x8b\xd7\x96\xbcB\x04\xa1=\x18\x04\xe2\xbd\xacE\x8c\xbc[V\x9f<\x1d\x080=\x83\x8bW=3\xce\x19=\xf9\x97S\xbd\xb6\x98\xe3;Lf"\xbd\x9e\x02\xc2\xbc\x18\x04\xe2\xbc\xc6\x8aG=V\xa5\x9e\xbd\xdc\xb7Y\xbc*_\xd5\xbc\x10\x81B=R\xc4U=\xba\x12\x84<"/\x03\xbe\x10\x0b\xf0<\xf5\x18\x02<\xdbqO=h\x8a?\xbd\xf0{\\=\xc5!8=v\xcd\x89=\x93\xad\x03>(\xb0;\xbc\xb1]\xb5<\xfb\x00\xe3<\x1c\xe5\xaa=\xf3M\xfb<',
        },
        {
            "id": 2,
            "name": "Two",
            "description": "Second item",
            "emb_text_embedding_3_large_256": b'\x14i`\xbc\x82\xa3\x8c;at\xbd\xbc\x88\x12\xb9\xbc\x8ae\x8a\xbd\t\xed;\xbe3vx\xbbN\n">\x974\x80\xbd\x8d\x81e<\xefSq;\xf2\xd2\xbb\xbb\\g\xc5=SyN\xbd;8\xf6=\xe5\xd7L=\x8b)\xf3\xbci\x00\x00>\x88tm\xbcsP+\xbd\x93\xebp;]\x93\xbe;\xeb\xda\x82\xba?\xe39=\xd6\xb0\xe4<\x81\xad\xce<\x82\xa3\x0c\xbd\x969\xa1<\xadS!\xbd".\x94=\xc9z\x99;\x969\xa1<-\xd61<\xc4\xd4v\xbd\x899\x91\xbcQ!\xdc\xbc/.\xa4=\xa0\xba\xe6\xbb\x88\r\x18\xbc\xbaX\xd2\xbcg\xdeH<\xaam\x81\xb9\xe0\x94\x19\xbe\xf5\xbd|\xbcY\xe3Y=$\x86\x06=\x8ae\x8a\xbc\xc2\xe4\x14\xbc=Z-=>\x86\xa6<[@\xed<\xf8\xda\x12<\xe1\xf1\xac\xbd`H\xc4<\xcd[\x18\xbc@E\xee\xbcH\xa0\x96=\x85\xbf\xe7=\xc7\x89|\xbd\n\xe8\x9a<\x974\x00;\\\x98\xdf\xbd/.$\xbdK\xbcq\xbc\xcf\xe9\xc5\xbd\xd4\xc0\x82<\xb0\x08\xa7<D!L=\x7fU\xdc<\xf2\x08w=\x10!\x0c\xbe}\x9b5\xbd>\xed{=\xb6\xa8m=\xa0\x89L=\x08\x90\xa8<\x05B\xf8=x\x8e\xbd\xba\x9bF\x19\xbd\xce\xee\xe6<\xe0\xca\xd4\xbb\x11~\x1f\xbdC\xf5\xd2\xbb\xd5SQ\xbd$\xe8\xba<\xedc\x0f\xbda\xaa\xf8\xbd\x8a\xcc_\xbe\x83\x00 =\x0b\x14\x14\xbd\x81w\x13=\xa0\xba\xe6<\xcd\x8c\xb2\xbc&\x0f\x93;\xaeN\x80\xbd\xbf`)=\xb3\x8c\x92=\x8f\xd9\xd7<\x0c\xa7b\xbd\xca\xd7\xac=\x84gu\xbcm\x12\x99<\xf8\xda\x92\xbd\xe2\x84\xfb\xbd\xd7u\x88=3q\xd7=r\x86\xe6=\xe2\x1d&<\xadX\xc2\xbca>\x82\xbd\x06\xd6\x81=\xda[\xa8\xbc\xbe\x03\x96\xbd3\n\x02\xbd\xfd\x7fz\xbd"\x95i\xbdw\x98\x7f\xbd{\x12\xa9\xbc/_>\xbd5\x98/\xbc1\x86\x96=%\x144\xbd\x0fW\xc7\xbcI\xfd\xa9=y\xebP\xbd\xbbN\x10\xbd\x83\xd4\xa6\xbc\x1b\x98\x8f<Q&\xfd=-\x0c\xed\xbdE\x17\n\xbe\x9eb\xf4;>\xe8\xda\xbd\xa1\x7f\n\xbc\xf7\xae\x99=V.\xd4=\x0bv\xc8<\xa491\xbc\x98\x91\x93<\x8f\xd9\xd7<j\x93\xce\xbd\x8fr\x82\xbd\x10R\xa6=e$\xa2=\xa8\x1a\xb0\xbdY\xe8z=i6\xbb\xbc\xcd\x8c\xb2<\xdb\xb8\xbb\xbcG\xa57=\xdfh\xa0\xbcK\xbc\xf1\xbd\x1c\xf5"\xbe-\xd6\xb1=\xe2\x84{\xbc+~\xbf=:\xa5\xa7\xbc\xdc\x82\x00\xbd\xed\xc5\xc3<\xa5\xc7\xde\xbb\xb2e:=\x8c\xee\x96=\xac\xf6\x8d<\xc6,i\xbd\xe6\xcd\x8a\xbc\x12\x0c\xcd<\x97e\x9a=8M5<\x1dR6\xbdq)\xd3=\xac\'(\xbd:t\r\xbd\xc2\xe4\x14=\xaeN\x00\xbd\x9d\x9e\x0b>\xe5\xa1\x91\xbd\xc1\xbd<>\xf8\x10\xce\xbd\x00\xd3K\xbd\xda\xc2}\xbc^\x8e\x1d=^\xba\x96<\xad"\x07>\n\x195\xbcj\x93\xce\xbd\x04\xe5d\xbd\x143\xa5\xbc\xa2\xdc\x9d=\x95Cc\xbc[@m\xbd5\x98/<k\x8e\xad\xbdSC\x13>\xf4\xf9\x93=\xfb\xc5S=\x8b)s=y\xb5\x15\xbb\xe7*\x9e\xbc\xcc4@=\x97\x9bU=\x04\xe5\xe4;o\xa0F\xbc\x1b\xc9\xa9=\xebAX\xbc-\x07L\xbd_\xe6\x8f\xbc{\r\x88=xX\x02<\x03\xb4J=\x05B\xf8\xbc\xff\xd7l=\x10\xb9\xfb\xbb\x0f\\h\xbd\xcc\x96t<\xb1\x03\x06<H\xd1\xb0=\x8b\x91\x83\xbd\x07\x9a\xea<\x80\xb2\xef\xbbBg%=\x9a\x1fA<l\xba&\xbd<_N=\x88\r\x98=d\xf8\xa8\xbd\xf6\xb3:\xbdx]\xa3;\x11\xaf\xb9=+~?=\x06\xd6\x81<\xa8\x15\x8f\xbdV.\xd4\xbd\x13\xd6\x91\xbd\xd1m1\xbd,\xaa\xb8\xbd+\x1c\x0b\xbb\xaam\x81=/_>\xbd\x80K\x9a9\xc5\x99\x9a=\xb9\xfb\xbe<\x96\xa0\xf6<)\xf5\xb2\xbd\x95\xe1.\xbd\xc8SA\xbd\x93\x84\x1b=\xdb\xbd\xdc\xbc\x13\xd6\x91=\xe6\x03\xc6<\x92\x8e\xdd;\xe1\'h=\xc1\xb8\x1b<\xdb\x8cB;\x98\xf8h=H\x07l=\xf3\x9c\x00>',
        },
        {
            "id": 3,
            "name": "Three",
            "description": "Third item",
            "emb_text_embedding_3_large_256": b"\xa7\x07\xae\xbc%>F\xbd\x80s0\xbd\xc6`\x11=\xd6\xac\x85\xbd\xf3\xdb\x03\xbe\x87Z\xa5\xbd\xec\x9e\xfc=\xe0\xbb\x97=\x98e\xcc\xbc\xf7-a=\xc3\xcby\xbb#\xea =I\xd4\x8b\xbc\xdb\xfe\x87=\xeb\x89\xdc\xbb\x9f\xcc\x13\xbdF\x95\xe1=9\xb2W\xbdW \xdb\xbd\x05\x11\x10\xbd\x8c\x02\x95;W [=\x13\x89\x0c>NP\xce=\xf4/)\xbc\x99\xb9\xf1<\x90T\xf2:nh\xe4\xbc\x1d\xd7~=l\xff\x9e\xbcTx\x90<5u\x1a<\xad\x03\xc3\xbd\xb1\x81M\xbd<0b\xbd\xd8\x95\x1d>\xa3\x1eq\xbb\xee3\x14=\x9f\xcc\x93\xbc|u\xd3\xb9\x8e\x15\xed;\xd0Z\x83\xbd\x19\x05O\xbd\xa4\x9eC=\xa2_><\xf3\xf0#=\xd8\xd4\xfd;\xffS\xdb<\x11J\x87=q\xbc\xae<\xf8-\x06\xbc\x89m\xfd<~\xc9x=\x9b\xf8\xf6\xbc\x19p\x81\xbc\x9b\xe3\xd6=\xd1\x196>\x85\x06\x00\xbd\xb2\xab\xb2=\xe5\xa2\xe7\xbc\xa2J\x1e\xbe6\xc9?\xbc\xcb\xc7\xb3<\xb5\xa9\xea\xbdx\xcdc<h\xec!\xbb\xd7\xeb\x8a=\x11\xdfT\xbc36\x15<M\xfc(\xbd\xa8\x9c\xa0\xba\x94=\xaf\xbco\xd3\x16\xbd!\xab\x9b<\xd6\xac\x05=\x1cY\x19>T\xa2P=\xc2\xf7&\xbc\xb8'\xf5:y\xe2\xa8<\xb9\x92'\xbd\xf5\xee\xdb<_\x1c\x15<\x11_\xa7;\xeb_\x9c\xbc\xef\xc8\x06\xbe\xdf\xbbr\xbe\\\xb3\xaa;\xa6]\x1b=\xbafz\xbc\x8d\xd6\xe7<:G\xca;\xf7\x18\xc1\xbc\xb9'\x1a\xbdc\xc4_\xbdD,\x9c=`\xc6'=\xda\xbf\x02\xbd\x9d\r<=\xbc\x102\xbd\xfd\x14V\xbd\xf5\xc4\x9b\xbdR9\x8b\xbdq\xe6\xee=4\xf5G=\xd4\x97@=Z\xdf2=\xe9\x8b\xa4\xbd\x8c\x97\x07\xbe\xa5\xc8\xa8=\xfal\x8b=\x18\xdb\x8e\xbd\xe7\xe1\xec:\xfc\xc00\xbd\xd7\x80\xd8\xbc{!\xae\xbd~\xc9x\xbd\x8e\xd6\x0c\xbd\x00i\xa0=\x97\xa6\x19=\xf8B\xa6\xbd\xcfZ^\xbcw\xa3\xa3<l>\x7f\xbco\xe86\xbd\x99\x8f\xb1\xbd|u\xd3<I\x13l=\x8f\x95\xbf\xbd\x06\xd0\xc2\xbd\xc2\xe2\x86=\x832\x88\xbdV\xcc\xb5\xbc\x94(\x0f<\x81\xf3\x02>*\x90H=\xe8 r\xbd\xe2\x8f\x0f=BX$\xbdV\xe1\xd5\xbd\xeeH\xb4\xbdLR\x96=V\xb7\x95=\x07zU\xbd\x19\x85\xa1=\x10\xca4\xbd\xb1\x16@=\xa2_><\xf8l\xe6<\x88\x048<\xf5\x03|\xbd\x88\x048\xbdAC\x84=|`3\xbd\xdf\xbb\xf2<\xdai\x15=\xe8\xe1\x91\xbb_\x1c\x15\xbd~\x8a\x18\xbd\xc9\xde\x9b=\xffh\xfb=\xc3\x8c\x99\xbcM\xe7\x88<\x14\xc8\x91\xbdw\xa3\xa3\xbc@.\xbf=\xdc\xbd\xba<\x80s\xb0\xbd\"j\xce=\x07P\x15\xbb\xaf-\xa8\xbd\xa1\xca\xcb=\xb7S\xa2\xbaO\xe5\xc0=\xd7\x80\xd8\xbc\xf8l\xe6=\xdd=\x8d\xbd\xdb\x13(=\xb3j\xe5;\xa4\x89\xa3=Z_\x05\xbd\xe1e\xaa=\x14H?\xbdl\x14\xbf\xbd\x93\xa8\xbc\xbc\xc5\xf5\xde<\xa1\xa0\x8b=\x08\xfa'\xbdo}\xa9\xbd\xf3\xdb\x83;\xc0\x0e\x8f\xbdx8\x96=)'\x83<\xfe\xa9\xc8=?\xael=\xf5\xee\xdb\xbc\x02\xd2\xe5<\x16\x07\x97=\xcf\xc5\x90\xbc\xda\xfeb<\xc5\xcb\x9e\xbd\ra\xca=\xa9F3=\xfc\xeap\xbc\xb0\x01{=\x17\xdbi=w\x0e1;\xa9\x05\xe6<\xa1\xdfk\xbdvy>=v\xa3~=\xb1W\x8d\xbdW [\xbc\xaf-(\xba\xab\x1a\xab=\n\xce\x9f\xbc\x9eL\xc1<\x96;g=\xad\xee\"=s\x90&=\xf3\x85\x16<\xb2+`=\xf7\x18\xc1=\xa0\xa0\xe6\xbd\x18F\x1c=\xaf-\xa8=\xcb\xdc\xd3=\xb5j\x8a\xbc)f\xe3<\xe1P\x8a\xbc\xfe\x94\xa8<\x80\xde=\xbd\xd4m\x80\xbd\xc3\xcb\xf9\xbdS\xe3\x1d=\xe3$\x82\xbc\x13\xb3\xcc\xbd\xceE\xbe<\x143\x9f=\x97\xbb9\xbc\x832\x88<q\xd1\xce\xbd2\xcb\xe2<1\x0c\xb0\xbc\xba\xe6\xcc\xbc\x1a\x1a\x14\xbd\xcb\xf1\xf3=\x92\xfe)=\xd2\xc3H=G\xd4f=q\xbc.\xbc6I\xed\xbc\x1d\xad\xbe;N;.=\x9a9\xc4=",
        },
    ]
