import asyncio
import os

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine

from dot_search import (
    BM25Config,
    EmbeddingConfig,
    TableConfig,
    SearchEngine,
    serialize_row,
)

DB_URL = os.environ.get("DOT__DB_URL", "postgresql+asyncpg://postgres:postgres@localhost:5433/postgres")

PRODUCTS = [
    (1,  "Wireless Noise-Cancelling Headphones", "Premium over-ear headphones with active noise cancellation, 30-hour battery, and foldable design for travel.", "electronics", 249.99, True,  "audio bluetooth travel"),
    (2,  "Mechanical Keyboard",                  "Tenkeyless mechanical keyboard with Cherry MX Blue switches, RGB backlighting, and aluminum frame.",               "electronics", 129.00, True,  "typing gaming peripherals"),
    (3,  "Cast Iron Skillet",                    "Pre-seasoned 12-inch cast iron skillet perfect for searing steaks, baking cornbread, and stovetop cooking.",       "kitchen",      39.95, True,  "cooking baking stovetop"),
    (4,  "Yoga Mat",                             "Eco-friendly 6mm thick yoga mat made from natural rubber, non-slip surface, ideal for hot yoga and pilates.",      "fitness",      55.00, True,  "exercise wellness stretching"),
    (5,  "Espresso Machine",                     "Compact semi-automatic espresso machine with 15-bar pump, steam wand for milk frothing, and stainless steel body.", "kitchen",    199.00, True,  "coffee espresso brewing"),
    (6,  "Running Shoes",                        "Lightweight neutral running shoes with responsive foam midsole, breathable mesh upper, and reflective details.",   "footwear",     89.99, False, "running athletics outdoor"),
    (7,  "Portable Solar Charger",               "20000 mAh solar power bank with dual USB-A and USB-C outputs, waterproof casing, and foldable solar panels.",     "electronics",  44.99, True,  "solar outdoor camping"),
    (8,  "Stainless Steel Water Bottle",         "Vacuum-insulated 32oz water bottle keeps drinks cold 24 hours or hot 12 hours, leak-proof lid, BPA-free.",        "outdoors",     27.50, True,  "hydration hiking camping"),
    (9,  "Bamboo Cutting Board",                 "Large end-grain bamboo cutting board with juice groove, non-slip feet, and antimicrobial surface.",               "kitchen",      34.00, True,  "cooking food prep kitchen"),
    (10, "4K Action Camera",                     "Waterproof action camera recording 4K at 60fps with image stabilization, wide-angle lens, and Wi-Fi streaming.",  "electronics",  79.99, True,  "video photography adventure"),
    (11, "Foam Roller",                          "High-density EVA foam roller for deep tissue massage, muscle recovery, and myofascial release post-workout.",     "fitness",      22.00, True,  "recovery massage sports"),
    (12, "Electric Kettle",                      "Variable temperature 1.7L electric kettle with gooseneck spout, stainless steel interior, ideal for pour-over coffee and tea.", "kitchen", 49.00, False, "tea coffee kitchen"),
    (13, "Smart Watch",                          "GPS smartwatch with heart rate monitor, sleep tracking, 5-day battery, and smartphone notifications.",            "electronics", 199.95, True,  "fitness health wearable"),
    (14, "Backpacking Tent",                     "Ultralight 2-person backpacking tent weighing under 2kg, freestanding, with waterproof rainfly and vestibule.",   "outdoors",    159.00, True,  "camping hiking shelter"),
    (15, "Resistance Bands Set",                 "Set of 5 latex resistance bands with varying tension levels for strength training and physical therapy.",         "fitness",      18.99, True,  "exercise rehab strength"),
    (16, "Ceramic Pour-Over Coffee Dripper",     "Handcrafted ceramic dripper compatible with Chemex and standard carafes, produces clean bright coffee flavor.",   "kitchen",      28.00, True,  "coffee brewing pour-over"),
    (17, "USB-C Hub",                            "7-in-1 USB-C hub with HDMI 4K, 3x USB-A, SD card reader, USB-C PD charging, and Ethernet port.",               "electronics",  35.00, True,  "laptop accessories connectivity"),
    (18, "Hiking Boots",                         "Waterproof mid-cut hiking boots with Vibram outsole, ankle support, and Gore-Tex lining for wet terrain.",       "footwear",    149.00, False, "hiking outdoor waterproof"),
    (19, "Dutch Oven",                           "5.5-quart enameled cast iron Dutch oven for slow braising, soups, and artisan bread baking.",                    "kitchen",      79.00, True,  "cooking baking kitchen"),
    (20, "Noise Machine",                        "White noise machine with 30 sounds, timer, and night light, for better sleep and focus in open offices.",        "electronics",  39.99, True,  "sleep focus relaxation"),
    (21, "Climbing Harness",                     "Lightweight sport climbing harness with padded waist and leg loops, adjustable buckles, and gear loops.",        "outdoors",     65.00, True,  "climbing safety rock"),
    (22, "Cold Brew Coffee Maker",               "1-gallon glass cold brew system with fine mesh filter, produces smooth low-acid coffee concentrate overnight.",  "kitchen",      24.99, True,  "coffee cold-brew beverage"),
    (23, "Laptop Stand",                         "Adjustable aluminum laptop stand with 6 height settings, foldable design, and compatibility with 10-17 inch laptops.", "electronics", 29.99, True, "ergonomic desk accessories"),
    (24, "Trekking Poles",                       "Collapsible carbon fiber trekking poles with ergonomic cork grips, flick locks, and interchangeable tips.",      "outdoors",     75.00, False, "hiking trekking poles"),
    (25, "Air Purifier",                         "True HEPA air purifier covering up to 500 sq ft, removes 99.97% of allergens, dust, pollen, and pet dander.",   "home",        119.00, True,  "health air quality home"),
]

ARTICLES = [
    (1,  "The Rise of Quantum Computing",           "Quantum computers leverage qubits in superposition to solve optimization problems exponentially faster than classical machines. IBM and Google have demonstrated quantum advantage in specific tasks.",   "Alice Chen",        2023, "technology"),
    (2,  "Mediterranean Diet and Longevity",         "Studies spanning decades show that populations adhering to the Mediterranean diet, rich in olive oil, legumes, and fish, have significantly lower rates of cardiovascular disease and longer lifespans.", "Dr. Marco Rossi",   2022, "health"),
    (3,  "Reforestation Efforts in the Amazon",     "Brazil's federal reforestation program aims to restore 12 million hectares of degraded rainforest by 2030, sequestering carbon and protecting endemic species in the world's largest tropical ecosystem.", "Ana Lima",          2023, "environment"),
    (4,  "Machine Learning in Drug Discovery",      "AI models trained on molecular structures can predict protein-ligand binding affinities with remarkable accuracy, cutting years off early-stage pharmaceutical research and reducing development costs.",  "Dr. Sarah Kim",     2024, "science"),
    (5,  "The Art of Sourdough Fermentation",       "Wild-fermented sourdough bread relies on a complex microbiome of lactobacilli and wild yeasts. Long cold fermentation develops flavor complexity and improves digestibility compared to commercial yeast.", "Jean-Paul Blanc",   2021, "food"),
    (6,  "Deep Sea Exploration Breakthroughs",      "New pressure-resistant ROVs are mapping previously unexplored hadal zones at depths exceeding 10000 meters, discovering chemosynthetic ecosystems and novel extremophile organisms.",                      "Dr. Mia Torres",    2023, "science"),
    (7,  "Urban Vertical Farming",                  "Vertical farms using LED lighting and hydroponic systems can grow leafy greens year-round with 95% less water than traditional agriculture, addressing food security in dense urban environments.",         "James Wu",          2022, "food"),
    (8,  "The History of the Silk Road",            "The ancient Silk Road connected Han dynasty China to Rome through 4000 miles of overland and maritime trade routes, facilitating exchange of silk, spice, glass, and ideas across Eurasia for centuries.", "Prof. Omar Said",   2020, "history"),
    (9,  "Training for an Ultramarathon",           "Ultramarathon preparation requires base-building phases of 60-80 miles per week, back-to-back long runs, and targeted nutrition strategies including fat adaptation and carbohydrate periodization.",       "Lisa Nakamura",     2023, "fitness"),
    (10, "Fermentation and Gut Health",             "Fermented foods like kimchi, kefir, and sauerkraut introduce beneficial bacteria to the gut microbiome, reducing systemic inflammation and improving mental health outcomes through the gut-brain axis.",   "Dr. Priya Shah",    2022, "health"),
    (11, "Exploring Patagonia's Torres del Paine",  "Patagonia's W Trek winds through granite towers, teal glacial lakes, and windswept pampas. The five-day route crosses hanging glaciers and offers encounters with guanacos and Andean condors.",            "Carlos Mendez",     2023, "travel"),
    (12, "The Neuroscience of Sleep",               "During NREM slow-wave sleep the hippocampus replays memories for consolidation to the neocortex. Chronic sleep deprivation impairs synaptic plasticity and accelerates amyloid-beta accumulation linked to Alzheimer's.", "Dr. Emily Foster", 2024, "science"),
    (13, "Ramen: From Street Food to Global Icon",  "Ramen originated as Chinese-style wheat noodles adapted in postwar Japan. Regional varieties like Sapporo miso, Hakata tonkotsu, and Tokyo shoyu reflect local ingredients and evolving culinary traditions.", "Kenji Watanabe",   2021, "food"),
    (14, "Renewable Energy Storage Challenges",     "Intermittency of solar and wind power demands scalable storage. Grid-scale lithium-ion batteries, pumped hydro, and emerging iron-air batteries each present different cost-capacity-lifespan tradeoffs.",  "Dr. Fatima Al-Hassan", 2023, "technology"),
    (15, "Kyoto's Hidden Temples",                  "Beyond Kinkakuji and Fushimi Inari lie lesser-known gems: Daitoku-ji's moss gardens, Fushimi Momoyama's samurai heritage, and Kurama's mountain shrine accessible only by narrow forest trail.",           "Yuki Tanaka",       2022, "travel"),
]


async def setup_tables(engine: sa.ext.asyncio.AsyncEngine) -> None:
    async with engine.begin() as conn:
        await conn.execute(sa.text("DROP TABLE IF EXISTS products CASCADE"))
        await conn.execute(sa.text("DROP TABLE IF EXISTS articles CASCADE"))
        await conn.execute(sa.text("DROP TABLE IF EXISTS ds_configs CASCADE"))
        await conn.execute(sa.text("""
            CREATE TABLE products (
                id SERIAL PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT NOT NULL,
                category TEXT NOT NULL,
                price FLOAT NOT NULL,
                in_stock BOOLEAN NOT NULL,
                tags TEXT NOT NULL
            )
        """))
        await conn.execute(sa.text("""
            CREATE TABLE articles (
                id SERIAL PRIMARY KEY,
                title TEXT NOT NULL,
                body TEXT NOT NULL,
                author TEXT NOT NULL,
                published_year INT NOT NULL,
                topic TEXT NOT NULL
            )
        """))

    async with engine.begin() as conn:
        await conn.execute(
            sa.text("INSERT INTO products (id, name, description, category, price, in_stock, tags) VALUES (:id, :name, :description, :category, :price, :in_stock, :tags)"),
            [{"id": r[0], "name": r[1], "description": r[2], "category": r[3], "price": r[4], "in_stock": r[5], "tags": r[6]} for r in PRODUCTS],
        )
        await conn.execute(
            sa.text("INSERT INTO articles (id, title, body, author, published_year, topic) VALUES (:id, :title, :body, :author, :published_year, :topic)"),
            [{"id": r[0], "title": r[1], "body": r[2], "author": r[3], "published_year": r[4], "topic": r[5]} for r in ARTICLES],
        )

    print(f"  Inserted {len(PRODUCTS)} products and {len(ARTICLES)} articles.")


def print_results(results, text_key: str) -> None:
    for i, r in enumerate(results, 1):
        text_preview = r.row.get(text_key, "")[:80].replace("\n", " ")
        print(f"  {i}. [id={r.id}] score={r.score:.4f}  {text_preview}")


async def main() -> None:
    print("=" * 70)
    print("dot-search integration test")
    print("=" * 70)

    # ------------------------------------------------------------------ setup
    print("\n--- Setting up tables ---")
    raw_engine = create_async_engine(DB_URL)
    await setup_tables(raw_engine)
    await raw_engine.dispose()

    engine = SearchEngine(DB_URL)

    # --------------------------------------------------------------- indexing
    print("\n--- Indexing products (ds_desc + ds_name, BM25 on description) ---")
    products_config = TableConfig(
        table="products",
        embeddings=[
            EmbeddingConfig(target_column="ds_products_description_default", source_column="description"),
            EmbeddingConfig(target_column="ds_products_name_default", source_column="name"),
        ],
        pk_column="id",
        bm25=[BM25Config(source_column="description")],
        batch_size=50,
    )
    await engine.index(config=products_config)
    print("  products indexed.")

    print("\n--- Indexing articles (ds_body, BM25 on body) ---")
    articles_config = TableConfig(
        table="articles",
        embeddings=[
            EmbeddingConfig(target_column="ds_articles_body_default", source_column="body"),
        ],
        pk_column="id",
        bm25=[BM25Config(source_column="body")],
        batch_size=50,
    )
    await engine.index(config=articles_config)
    print("  articles indexed.")

    # --------------------------------------------------------------- searches
    print("\n" + "=" * 70)
    print("SEARCH QUERIES")
    print("=" * 70)

    # Q1: products — audio / headphones
    print("\n[Q1] Products: 'over-ear headphones with noise cancellation'")
    results = await engine.search(
        "over-ear headphones with noise cancellation",
        "products",
        limit=5,
    )
    print_results(results, "description")
    assert results, "Expected at least one result"
    assert results[0].id == 1, f"Top result should be headphones (id=1), got id={results[0].id}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q2: products — coffee brewing
    print("\n[Q2] Products: 'brewing espresso coffee at home'")
    results = await engine.search(
        "brewing espresso coffee at home",
        "products",
        limit=5,
    )
    print_results(results, "description")
    assert results
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)
    top_ids = [r.id for r in results[:3]]
    coffee_ids = {5, 12, 16, 22}
    assert coffee_ids & set(top_ids), f"Expected a coffee product in top 3, got {top_ids}"

    # Q3: products with WHERE in_stock = true
    print("\n[Q3] Products (in_stock=true): 'outdoor hiking gear'")
    results = await engine.search(
        "outdoor hiking gear",
        "products",
        limit=5,
        where="in_stock = true",
    )
    print_results(results, "description")
    assert results
    for r in results:
        assert r.row["in_stock"] is True, f"Result id={r.id} has in_stock={r.row['in_stock']}, expected True"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q4: products with WHERE price < 50
    print("\n[Q4] Products (price<50): 'kitchen cooking tools'")
    results = await engine.search(
        "kitchen cooking tools",
        "products",
        limit=5,
        where="price < 50",
    )
    print_results(results, "description")
    assert results
    for r in results:
        assert r.row["price"] < 50, f"Result id={r.id} has price={r.row['price']}, expected < 50"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q5: products with WHERE category = 'electronics'
    print("\n[Q5] Products (category=electronics): 'wireless connectivity device'")
    results = await engine.search(
        "wireless connectivity device",
        "products",
        limit=5,
        where="category = 'electronics'",
    )
    print_results(results, "description")
    assert results
    for r in results:
        assert r.row["category"] == "electronics", f"Result id={r.id} has category={r.row['category']!r}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q6: articles — quantum computing
    print("\n[Q6] Articles: 'quantum computers and computational advantage'")
    results = await engine.search(
        "quantum computers and computational advantage",
        "articles",
        limit=5,
    )
    print_results(results, "body")
    assert results
    assert results[0].id == 1, f"Top result should be quantum computing article (id=1), got id={results[0].id}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q7: articles — gut health and fermented food
    print("\n[Q7] Articles: 'fermented foods gut microbiome health'")
    results = await engine.search(
        "fermented foods gut microbiome health",
        "articles",
        limit=5,
    )
    print_results(results, "body")
    assert results
    top_ids = [r.id for r in results[:2]]
    assert 10 in top_ids or 5 in top_ids, f"Expected fermentation article (id=5 or 10) in top 2, got {top_ids}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q8: articles — travel in Japan
    print("\n[Q8] Articles: 'travel Japan temples ancient culture'")
    results = await engine.search(
        "travel Japan temples ancient culture",
        "articles",
        limit=5,
    )
    print_results(results, "body")
    assert results
    assert results[0].id == 15, f"Top result should be Kyoto temples article (id=15), got id={results[0].id}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q9: BM25 search on articles
    print("\n[Q9] Articles BM25: 'fermentation'")
    results = await engine.search(
        "fermentation",
        "articles",
        strategy="bm25",
        limit=5,
    )
    print_results(results, "body")
    assert results, "BM25 search returned no results"
    top_ids = [r.id for r in results[:2]]
    assert 5 in top_ids or 10 in top_ids, f"Expected fermentation article in top 2, got {top_ids}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q10: hybrid search — vector + BM25 fused via RRF
    print("\n[Q10] Articles hybrid (vector + BM25): 'sourdough bread microbiome'")
    results = await engine.search(
        "sourdough bread microbiome",
        "articles",
        limit=5,
    )
    print_results(results, "body")
    assert results
    assert results[0].id == 5, f"Top hybrid result should be sourdough article (id=5), got id={results[0].id}"
    assert [r.score for r in results] == sorted([r.score for r in results], reverse=True)

    # Q11: source_column=None — auto-serialize all columns
    print("\n[Q11] Products auto-serialize (source_column=None): 'lightweight outdoor shelter'")
    products_all_config = TableConfig(
        table="products",
        index_id="products_all",
        embeddings=[
            EmbeddingConfig(target_column="ds_products_all_auto_default", source_column=None),
        ],
        pk_column="id",
    )
    await engine.index(config=products_all_config)
    results = await engine.search(
        "lightweight outdoor shelter",
        "products_all",
        limit=3,
    )
    print_results(results, "description")
    assert results
    assert results[0].id == 14, f"Top result should be backpacking tent (id=14), got id={results[0].id}"

    # Q12: small batch_size forces multiple batches — re-index articles with batch_size=4
    print("\n[Q12] Re-indexing articles with batch_size=4 (forces multiple batches over 15 rows)")
    articles_rebatch_config = TableConfig(
        table="articles",
        embeddings=[
            EmbeddingConfig(target_column="ds_articles_body_default", source_column="body"),
        ],
        pk_column="id",
        bm25=[BM25Config(source_column="body")],
        batch_size=4,
    )
    await engine.index(config=articles_rebatch_config)
    results = await engine.search(
        "quantum computers and computational advantage",
        "articles",
        limit=3,
    )
    print_results(results, "body")
    assert results[0].id == 1, f"After re-index, top result should still be id=1, got {results[0].id}"

    # Q13: BM25 + WHERE filter
    print("\n[Q13] Articles BM25 + WHERE (topic=food): 'fermentation'")
    results = await engine.search(
        "fermentation",
        "articles",
        strategy="bm25",
        limit=5,
        where="topic = 'food'",
    )
    print_results(results, "body")
    assert results
    for r in results:
        assert r.row["topic"] == "food", f"Result id={r.id} has topic={r.row['topic']!r}, expected 'food'"
    assert results[0].id == 5, f"Top BM25+WHERE result should be sourdough article (id=5), got {results[0].id}"

    # Q14: divergent RRF — name-biased vs description-biased targets
    print("\n[Q14] RRF fusion — name-biased vs description-biased targets")
    await engine.index(config=TableConfig(
        table="products",
        index_id="products_name_only",
        embeddings=[EmbeddingConfig(target_column="ds_products_name_only_name_default", source_column="name")],
        pk_column="id",
    ))
    name_only = await engine.search("Dutch Oven", "products_name_only", limit=5)

    await engine.index(config=TableConfig(
        table="products",
        index_id="products_desc_only",
        embeddings=[EmbeddingConfig(target_column="ds_products_desc_only_description_default", source_column="description")],
        pk_column="id",
    ))
    desc_only = await engine.search("Dutch Oven", "products_desc_only", limit=5)

    await engine.index(config=TableConfig(
        table="products",
        index_id="products_fused",
        embeddings=[
            EmbeddingConfig(target_column="ds_products_fused_name_default", source_column="name"),
            EmbeddingConfig(target_column="ds_products_fused_description_default", source_column="description"),
        ],
        pk_column="id",
    ))
    fused = await engine.search("Dutch Oven", "products_fused", limit=5)

    print(f"  name-only top3: {[r.id for r in name_only[:3]]}")
    print(f"  desc-only top3: {[r.id for r in desc_only[:3]]}")
    print(f"  fused     top3: {[r.id for r in fused[:3]]}")
    assert name_only[0].id == 19, f"name-only top should be Dutch Oven (id=19), got {name_only[0].id}"
    assert fused[0].id == 19, f"fused top should be Dutch Oven (id=19), got {fused[0].id}"
    fused_top_score = fused[0].score
    name_top_score = name_only[0].score
    assert fused_top_score != name_top_score, "Fused score should differ from single-target score"

    # ----------------------------------------------------------------- cleanup
    await engine.close()

    print("\n" + "=" * 70)
    print("All assertions passed.")
    print("=" * 70)


if __name__ == "__main__":
    asyncio.run(main())
