# Initialize the OpenAI client
import os
from typing import List

from getpass import getpass
import openai
from tenacity import retry, wait_exponential

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass("OpenAI API Key: ")

client = openai.OpenAI()

# This introduces exponential backoff in case ChatGPT is being rate limited/has errors
@retry(wait=wait_exponential(multiplier=1.2, min=4, max=30))
def chat_gpt_completion(prompt: str) -> str:
    return client.chat.completions.create(
        messages=[{
            "role": "user",
            "content": prompt,
        }],
        model="gpt-3.5-turbo",
    ).choices[0].message.content

@retry(wait=wait_exponential(multiplier=1.2, min=4, max=30))
def ada_embedding(text: str) -> List[float]:
    return client.embeddings.create(
        input=text,
        model="text-embedding-3-small",
        timeout=10,
    ).data[0].embedding

# Initialize the Astra DB client
import os

from astrapy import DataAPIClient

# Get the Astra DB API endpoint and Token from the environment or user input
if "ASTRA_DB_API_ENDPOINT" in os.environ:
    ASTRA_DB_API_ENDPOINT = os.environ["ASTRA_DB_API_ENDPOINT"]
else:
    ASTRA_DB_API_ENDPOINT = input("Astra DB Endpoint: ")
if "ASTRA_DB_APPLICATION_TOKEN" in os.environ:
    ASTRA_DB_APPLICATION_TOKEN = os.environ["ASTRA_DB_APPLICATION_TOKEN"]
else:
    ASTRA_DB_APPLICATION_TOKEN = getpass("Astra DB Token: ")
# set a different keyspace here programmatically if needed:
ASTRA_DB_KEYSPACE = os.environ.get("ASTRA_DB_KEYSPACE")

astra_db_client = DataAPIClient()
astra_db = astra_db_client.get_database(
    ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
    keyspace=ASTRA_DB_KEYSPACE,
)

import redshift_connector

RS_HOST_URL = os.getenv(
    "REDSHIFT_HOST_URL",
    getpass("RedShift Host URL: "),
)
RS_USER = os.getenv(
    "REDSHIFT_USER",
    input("RedShift User: "),
)
RS_PASSWORD = os.getenv(
    "REDSHIFT_PASSWORD",
    getpass("RedShift Password: "),
)
RS_DB_NAME = os.getenv(
    "REDSHIFT_DB_NAME",
    input("RedShift DB Name: "),
)

def rs_execute(*statements, return_as_df: bool = True, commit: bool = True) -> pd.DataFrame:
    """Execute RedShift statements and return the result of the final query"""
    with redshift_connector.connect(
        host=RS_HOST_URL,
        database=RS_DB_NAME,
        user=RS_USER,
        password=RS_PASSWORD,
    ) as conn:
        with conn.cursor() as cursor:
            for statement in statements:
                cursor.execute(statement)

            try:
                if return_as_df:
                    res = cursor.fetch_dataframe()
                else:
                    res = cursor.fetchall()
            except redshift_connector.ProgrammingError:
                # No result set (final statement was not a SELECT)
                res = None

            if commit:
                conn.commit()

            return res

from datasets import load_dataset

spider = load_dataset("spider", split="validation")
spider_df = spider.to_pandas()
spider_schema = load_dataset("richardr1126/spider-schema", split="train")
spider_schema_df = spider_schema.to_pandas()

spider_df.head(5)

# Remove some questions to test the pipeline with
test_ndxs = [21, 28, 60, 63, 75]  # Randomly pick some queries the model succeeds and fails on
ndxs_to_drop = test_ndxs.copy()
test_df = spider_df.loc[test_ndxs]

for ndx in test_ndxs:
    row = spider_df.loc[ndx]
    print(f"Question ({row['db_id']}):\n    {row['question']}")
    print(f"Gold Query SQL:\n    {row['query']}")
    print("=" * 80)

    # Remove any rows that have the exact same query as the answer
    # This prevents data leakage so that the dynamic few-shot approach can't see the correct query
    same_query_rows = spider_df[spider_df['query'] == row['query']]
    for same_query_ndx in same_query_rows.index:
        ndxs_to_drop.append(same_query_ndx)

spider_df = spider_df.drop(ndxs_to_drop)

# Define concert_singer database
CREATE_TABLES_SQL = """
-- Create the stadium table
CREATE TABLE stadium (
    Stadium_ID INT PRIMARY KEY,
    Location TEXT,
    Name TEXT,
    Capacity INT,
    Highest INT,
    Lowest INT,
    Average INT
);

-- Create the singer table
CREATE TABLE singer (
    Singer_ID INT PRIMARY KEY,
    Name TEXT,
    Country TEXT,
    Song_Name TEXT,
    Song_release_year TEXT,
    Age INT,
    Is_male BOOLEAN
);

-- Create the concert table
CREATE TABLE concert (
    concert_ID INT PRIMARY KEY,
    concert_Name TEXT,
    Theme TEXT,
    Stadium_ID INT,
    Year TEXT,
    FOREIGN KEY (Stadium_ID) REFERENCES stadium(Stadium_ID)
);

-- Create the singer_in_concert table
CREATE TABLE singer_in_concert (
    concert_ID INT,
    Singer_ID INT,
    PRIMARY KEY (concert_ID, Singer_ID),
    FOREIGN KEY (concert_ID) REFERENCES concert(concert_ID),
    FOREIGN KEY (Singer_ID) REFERENCES singer(Singer_ID)
);
"""

POPULATE_DATA_SQL = """
-- Populate the stadium table
INSERT INTO stadium (Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average) VALUES
(1, 'New York, USA', 'Liberty Stadium', 50000, 1000, 500, 750),
(2, 'London, UK', 'Royal Arena', 60000, 1500, 600, 900),
(3, 'Tokyo, Japan', 'Sunshine Dome', 55000, 1200, 550, 800),
(4, 'Sydney, Australia', 'Ocean Field', 40000, 900, 400, 650),
(5, 'Berlin, Germany', 'Eagle Grounds', 45000, 1100, 450, 700);

-- Populate the singer table
INSERT INTO singer (Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male) VALUES
(1, 'John Doe', 'USA', 'Freedom Song', '2018', 28, TRUE),
(2, 'Emma Stone', 'UK', 'Rolling Hills', '2019', 25, FALSE),
(3, 'Haruki Tanaka', 'Japan', 'Tokyo Lights', '2020', 30, TRUE),
(4, 'Alice Johnson', 'Australia', 'Ocean Waves', '2021', 27, FALSE),
(5, 'Max Müller', 'Germany', 'Berlin Nights', '2017', 32, TRUE);

-- Populate the concert table
INSERT INTO concert (concert_ID, concert_Name, Theme, Stadium_ID, Year) VALUES
(1, 'Freedom Fest', 'Pop', 1, '2021'),
(2, 'Rock Mania', 'Rock', 2, '2022'),
(3, 'Electronic Waves', 'Electronic', 3, '2020'),
(4, 'Jazz Evenings', 'Jazz', 3, '2019'),
(5, 'Classical Mornings', 'Classical', 5, '2023');

-- Populate the singer_in_concert table
INSERT INTO singer_in_concert (concert_ID, Singer_ID) VALUES
(1, 1),
(1, 2),
(2, 3),
(3, 4),
(4, 5),
(5, 1),
(2, 2),
(3, 3),
(4, 4),
(5, 5);
"""

statements = [
    statement.strip() for statement in (CREATE_TABLES_SQL + POPULATE_DATA_SQL).split(";")
    if len(statement.strip()) > 0
]
rs_execute(*statements)

# Define pets_1 database
CREATE_TABLES_SQL = """
-- Create the Student table
CREATE TABLE Student (
    StuID INT PRIMARY KEY,
    LName VARCHAR(255),
    Fname VARCHAR(255),
    Age INT,
    Sex VARCHAR(10),
    Major INT,
    Advisor INT,
    city_code VARCHAR(50)
);

-- Create the Pets table
CREATE TABLE Pets (
    PetID INT PRIMARY KEY,
    PetType VARCHAR(255),
    pet_age INT,
    weight INT
);

-- Create the Has_Pet table
CREATE TABLE Has_Pet (
    StuID INT,
    PetID INT,
    FOREIGN KEY (StuID) REFERENCES Student(StuID),
    FOREIGN KEY (PetID) REFERENCES Pets(PetID)
);
"""

POPULATE_DATA_SQL = """
-- Populate the Student table
INSERT INTO Student (StuID, LName, Fname, Age, Sex, Major, Advisor, city_code) VALUES
(101, 'Smith', 'John', 20, 'M', 501, 301, 'NYC'),
(102, 'Johnson', 'Emma', 22, 'F', 502, 302, 'LAX'),
(103, 'Williams', 'Michael', 21, 'M', 503, 303, 'CHI'),
(104, 'Brown', 'Sarah', 23, 'F', 504, 304, 'HOU'),
(105, 'Jones', 'David', 19, 'M', 505, 305, 'PHI');

-- Populate the Pets table
INSERT INTO Pets (PetID, PetType, pet_age, weight) VALUES
(201, 'dog', 3, 20.5),
(202, 'cat', 5, 10.2),
(203, 'dog', 2, 8.1),
(204, 'parrot', 4, 0.5),
(205, 'hamster', 1, 0.7);

-- Populate the Has_Pet table
INSERT INTO Has_Pet (StuID, PetID) VALUES
(101, 201),
(101, 202),
(105, 203),
(103, 204),
(104, 205),
(105, 201);
"""

statements = [
    statement.strip() for statement in (CREATE_TABLES_SQL + POPULATE_DATA_SQL).split(";")
    if len(statement.strip()) > 0
]
rs_execute(*statements)

def eval_generated_queries(queries: List[str]) -> pd.DataFrame:
    """Evaluate the given queries against the test set, and report the performance of each row"""
    report = []

    for ndx, pred_sql in enumerate(queries):
        row = test_df.iloc[ndx]
        gold_sql = row["query"]
        gold_results = rs_execute(gold_sql)

        pred_sql = queries[ndx]
        try:
            pred_results = rs_execute(pred_sql)
            err = None

            # Figure out correctness by deciding between set matching (no order) and list matching (ordered)
            correct = len(gold_results) == len(pred_results)
            if correct:
                # Sort columns so the arrays match up
                pred_results = pred_results.reindex(sorted(pred_results.columns), axis=1)
                gold_results = gold_results.reindex(sorted(gold_results.columns), axis=1)
                if "order by" in gold_sql.lower():
                    # All row orderings must match
                    correct = (gold_results.values == pred_results.values).all()
                else:
                    for col in gold_results.columns:
                        if set(gold_results[col].tolist()) != set(pred_results[col].tolist()):
                            correct = False

        except redshift_connector.ProgrammingError as e:
            pred_results = None
            err = e
            correct = False

        report.append({
            "DB ID": row["db_id"],
            "Question": row["question"],
            "Correct": correct,
            "Error": err,
        })

    return pd.DataFrame(report)

# Create a collection for vector search
collection = astra_db.create_collection(
    "text2sql_examples_test1",
    dimension=1536,
)

from tqdm.auto import tqdm

tqdm.pandas()

spider_embeddings = spider_df["question"].progress_apply(lambda question: ada_embedding(question))
spider_df["question_embedding"] = spider_embeddings

def insert_doc(row):
    return collection.insert_one(
        {
            "_id": row.name,
            "question": row["question"],
            "gold_sql": row["query"],
            "db_id": row["db_id"],
            "$vector": row["question_embedding"],
        }
    )

spider_df.progress_apply(insert_doc, axis=1)

def dynamic_few_shot_prompt_fn(db_id: str, question: str) -> str:
    emb = ada_embedding(question)
    docs = collection.find(sort={"$vector": emb}, limit=2)
    closest_q_ids = [
        doc["_id"] for doc in docs
    ]

    return few_shot_prompt(db_id, question, indices=closest_q_ids)

# Include the prompt in the output to show the model's context
sql1 = generate_sql(
    question="How many concerts occurred in 2014 or 2015?",
    db_id="concert_singer",
    prompt_fn=dynamic_few_shot_prompt_fn,
    debug_prompt=True,
)
print(sql1)

sql2 = generate_sql(
    question="Show the stadium names without any concert.",
    db_id="concert_singer",
    prompt_fn=dynamic_few_shot_prompt_fn,
)
print(sql2)

sql3 = generate_sql(
    question="What are the students' first names who have both cats and dogs as pets?",
    db_id="pets_1",
    prompt_fn=dynamic_few_shot_prompt_fn,
)
print(sql3)

sql4 = generate_sql(
    question="Find the id of students who do not have a cat pet.",
    db_id="pets_1",
    prompt_fn=dynamic_few_shot_prompt_fn,
)
print(sql4)

sql5 = generate_sql(
    question="Find the first name and age of students who have a pet.",
    db_id="pets_1",
    prompt_fn=dynamic_few_shot_prompt_fn,
)
print(sql5)

sqls = [sql1, sql2, sql3, sql4, sql5]
report = eval_generated_queries(sqls)
display(report)