"""Postgres + Mongo materialization into the session DuckDB.

Self-contained: builds a tiny pg_dump-style .sql and a tiny mongodump .bson
in tmp_path. Postgres needs the embedded `pgserver`; Mongo needs `bson`
(pymongo) — both in the `connectors` extra. Each test skips cleanly if its
dependency isn't installed.
"""

from __future__ import annotations

import sqlite3

import duckdb
import pytest

from lexsi_ds.agent.datasource import AttachSpec, load_attached_handle

bson = pytest.importorskip("bson", reason="needs pymongo/bson (connectors extra)")


def _mongo_dump(folder, collection, docs):
    folder.mkdir(parents=True, exist_ok=True)
    with open(folder / f"{collection}.bson", "wb") as f:
        for d in docs:
            f.write(bson.encode(d))
    # mongodump also drops a metadata.json — ensure it's ignored
    (folder / f"{collection}.metadata.json").write_text("{}")


def test_mongo_dump_materializes_with_nested_as_json(tmp_path):
    dump = tmp_path / "reviews_dump"
    _mongo_dump(dump, "reviews", [
        {"_id": bson.ObjectId(), "book_id": 1, "stars": 5, "tags": ["scifi", "classic"]},
        {"_id": bson.ObjectId(), "book_id": 2, "stars": 3, "tags": ["horror"]},
    ])
    handle = load_attached_handle(
        [AttachSpec(alias="rev", db_type="mongo", dump_folder=dump, db_name="reviews_db")],
        dataset_id="t:mongo",
    )
    names = {t.name for t in handle.tables}
    assert "rev_reviews" in names
    con = duckdb.connect(str(handle.duckdb_path), read_only=True)
    rows = con.execute("SELECT book_id, stars, tags FROM rev_reviews ORDER BY book_id").fetchall()
    con.close()
    assert rows[0][:2] == (1, 5)
    assert rows[0][2] == '["scifi", "classic"]'   # nested array → JSON text
    # queryable as JSON
    con = duckdb.connect(str(handle.duckdb_path), read_only=True)
    n = con.execute("SELECT count(*) FROM rev_reviews "
                    "WHERE list_contains(json_transform(tags, '[\"VARCHAR\"]'), 'horror')").fetchone()[0]
    con.close()
    assert n == 1


def test_sqlite_mixed_type_blob_column_falls_back(tmp_path):
    """Reproduces the krama bug: a dynamically-typed sqlite column holding a
    float in one row and an invalid-unicode blob in another. DuckDB's typed
    scanner rejects it; the Python sqlite3 fallback must still load both rows."""
    import sqlite3
    p = tmp_path / "dyn.db"
    c = sqlite3.connect(p)
    c.execute("CREATE TABLE t (id INTEGER, shape)")  # no affinity → dynamic typing
    c.execute("INSERT INTO t VALUES (1, 3.14)")
    c.execute("INSERT INTO t VALUES (2, ?)", (b"\xff\xfeGP",))  # non-utf8 blob
    c.commit(); c.close()
    handle = load_attached_handle(
        [AttachSpec(alias="d", db_type="sqlite", path=p)], dataset_id="t:dyn")
    assert "d_t" in {t.name for t in handle.tables}
    con = duckdb.connect(str(handle.duckdb_path), read_only=True)
    n = con.execute("SELECT count(*) FROM d_t").fetchone()[0]
    con.close()
    assert n == 2


def test_postgres_dump_materializes(tmp_path):
    pytest.importorskip("pgserver", reason="needs pgserver (connectors extra)")
    sql = tmp_path / "books.sql"
    sql.write_text(
        "CREATE TABLE books (id integer, title text, author text);\n"
        "COPY books (id, title, author) FROM stdin;\n"
        "1\tDune\tHerbert\n"
        "2\tIt\tKing\n"
        "\\.\n"
    )
    handle = load_attached_handle(
        [AttachSpec(alias="books_database", db_type="postgres", sql_file=sql, db_name="bookreview_db")],
        dataset_id="t:pg",
    )
    assert "books_database_books" in {t.name for t in handle.tables}
    con = duckdb.connect(str(handle.duckdb_path), read_only=True)
    rows = con.execute("SELECT title, author FROM books_database_books ORDER BY id").fetchall()
    con.close()
    assert rows == [("Dune", "Herbert"), ("It", "King")]


def test_cross_db_join_pg_mongo_sqlite(tmp_path):
    pytest.importorskip("pgserver", reason="needs pgserver (connectors extra)")
    sql = tmp_path / "books.sql"
    sql.write_text(
        "CREATE TABLE books (id integer, title text);\n"
        "COPY books (id, title) FROM stdin;\n1\tDune\n2\tIt\n\\.\n"
    )
    dump = tmp_path / "rev_dump"
    _mongo_dump(dump, "reviews", [
        {"_id": bson.ObjectId(), "book_id": 1, "stars": 5},
        {"_id": bson.ObjectId(), "book_id": 2, "stars": 3},
    ])
    sq = tmp_path / "meta.db"
    c = sqlite3.connect(sq)
    c.execute("CREATE TABLE genre (book_id int, name text)")
    c.executemany("INSERT INTO genre VALUES (?,?)", [(1, "scifi"), (2, "horror")])
    c.commit(); c.close()

    handle = load_attached_handle([
        AttachSpec(alias="b", db_type="postgres", sql_file=sql),
        AttachSpec(alias="r", db_type="mongo", dump_folder=dump),
        AttachSpec(alias="m", db_type="sqlite", path=sq),
    ], dataset_id="t:cross")
    con = duckdb.connect(str(handle.duckdb_path), read_only=True)
    rows = con.execute(
        "SELECT b.title, r.stars, m.name FROM b_books b "
        "JOIN r_reviews r ON r.book_id = b.id "
        "JOIN m_genre m ON m.book_id = b.id ORDER BY b.id"
    ).fetchall()
    con.close()
    assert rows == [("Dune", 5, "scifi"), ("It", 3, "horror")]