from dotenv import load_dotenv
from pinecone import Pinecone
from transformers import AutoTokenizer, AutoModel
import torch
import pandas as pd
import pandas.testing as pdt
import boto3
import os
from io import StringIO
from gradio_client import Client

load_dotenv()

MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)
pc = Pinecone(api_key=os.environ["PINECONE_KEY"])
index = pc.Index("product-recommender")


def compare_s3_to_setupfiles(file_name: str) -> bool:
    expected_table = pd.read_csv("setupfiles/test_" + file_name)
    actual_table = pd.read_csv(
        StringIO(
            boto3.client("s3")
            .get_object(Bucket="charityshop-media-assistant", Key=file_name)["Body"]
            .read()
            .decode("utf-8")
        )
    )
    actual_test_table = actual_table[
        actual_table.store_name == "TestStore"
    ].reset_index(drop=True)
    return pdt.assert_frame_equal(expected_table, actual_test_table)


def text_to_embedding(text):
    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1).numpy()
    return embeddings.flatten().tolist()


def find_similar_media(query_text, top_k=1):
    query_embedding = text_to_embedding(query_text)
    matches = index.query(
        vector=query_embedding,
        top_k=top_k,
        include_metadata=True,
        filter={"vector_type": "test"},
    )["matches"]
    df = pd.DataFrame(
        {"barcode": [m["id"] for m in matches], "score": [m["score"] for m in matches]}
    )
    for col in ["artist", "title", "media_type", "description", "vector_type"]:
        df[col] = [m["metadata"][col] for m in matches]
    return df


def test_InventoryDatabaseMatchesSampleRecords():
    compare_s3_to_setupfiles("inventory.csv")


def test_ChatHistoryDatabaseMatchesSampleRecords():
    compare_s3_to_setupfiles("chat_history.csv")


def test_VectorDatabaseMatchesSampleRecords():
    expected_table = pd.read_csv(
        "setupfiles/test_vector_database_query.csv", dtype={"barcode": str}
    )
    actual_test_table = find_similar_media("denver", top_k=5).drop(columns=["score"])
    pdt.assert_frame_equal(expected_table, actual_test_table)


def test_CanQueryLLM():
    current_site = "media_assistant"
    client = Client("ml-usecases/" + current_site, hf_token=os.environ["HF_TOKEN"])
    try:
        result = client.predict(
            user_query="Hello!!", api_name="/combined_recommendation"
        )
        return len(str(result)) > 0
    except:
        return False