import time
from datetime import datetime
from itertools import product

import pandas as pd
from snowflake.cortex import CompleteOptions, complete
from snowflake.snowpark import Session


from utils import generate_first_prompt, generate_prompt, llm_review, llm_verdict

if __name__ == "__main__":
    session = Session.builder.config("connection_name", "email_draft_poc").create()

    review_model = "llama3.1-70b"
    model_temperature = 0
    models = ["llama3.1-8b", "mistral-7b", "llama3.1-70b"]

    use_llm_for_draft = [False]
    use_llm_for_review = [True, False]

    artists = ["Miley Cyrus", "Addison Rae", "Tate McRae"]
    content_purpose = ["Promote new album", "Promote merch", "Promote tour", "Promote new single"]
    use_historical_messages = True
    release_name = "Music Makes Me Happy"

    for model, artist, purpose, draft, review in product(
        models, artists, content_purpose, use_llm_for_draft, use_llm_for_review
    ):
        # Start the timer
        start_time = time.time()
        print(model, artist, purpose, draft, review)

        suggested_first_prompt = generate_first_prompt(artist, purpose, "", use_historical_messages)

        llm_instructions = generate_prompt(
            suggested_first_prompt,
            "",
            artist,
            purpose,
            release_name,
            draft,
            session,
            use_historical_messages,
            model,
        )

        options: CompleteOptions = {
            "temperature": model_temperature,
            "guardrails": True,
        }
        initial_response = complete(
            model=model,
            prompt=llm_instructions,
            session=session,
            options=options,
        )

        response = llm_review(
            initial_response, review, artist, initial_response, session, review_model
        )

        final_verdict = llm_verdict(response, session)
        # End the timer
        end_time = time.time()
        elapsed_time = end_time - start_time

        verdict_artist = final_verdict["verdict_for_content"][0]["artist_name"]
        verdict_call_to_action = final_verdict["verdict_for_content"][0]["call_to_action"]
        verdict_score = final_verdict["verdict_for_content"][0]["email_score"]
        verdict_reason = final_verdict["verdict_for_content"][0]["score_reason"]

        experiment_df = pd.DataFrame(
            {
                "original_artist": [artist],
                "original_model": [model],
                "original_purpose": [purpose],
                "original_use_draft": [draft],
                "original_apply_review": [review],
                "verdict_artist": [verdict_artist],
                "verdict_call_to_action": [verdict_call_to_action],
                "verdict_score": [verdict_score],
                "verdict_reason": [verdict_reason],
                "first_prompt": [suggested_first_prompt],
                "revised_prompt": [llm_instructions],
                "email_draft_before_review": [initial_response],
                "final_email_draft": [response],
                "experiment_date": [datetime.now()],
                "experiment_duration_seconds": elapsed_time,
                "historical_messages": [use_historical_messages],
            }
        )

        experiment_sp = session.create_dataframe(data=experiment_df)
        experiment_sp.write.mode("append").save_as_table(
            "email_draft_experiment_results", table_type="transient"
        )