from locust import HttpUser, task, between
from gradio_client import Client, handle_file


class GradioAppUser(HttpUser):
    wait_time = between(1, 3)

    def on_start(self):
        self.client = Client("http://sd.demo.polygraf.ai:7890/")

    @task
    def test_generate_text_with_google_search_and_pdf(self):
        generated_text = self.client.predict(
            input_role="Student",
            topic="Low Resource Large Language Models",
            keywords="",
            article_length=350,
            format="Article",
            writing_style="Formal",
            tone="Professional",
            user_category="General Public",
            depth_of_content="Moderate analysis",
            structure="Introduction, Body, Conclusion",
            references="News outlets",
            num_examples="1-2",
            conclusion_type="Call to Action",
            # ai_model="LLaMA 3",
            ai_model="OpenAI GPT 4o Mini",
            google_search_check=False,
            year_from="2000",
            month_from="January",
            day_from="01",
            year_to="2024",
            month_to="August",
            day_to="08",
            domains_to_include=["com", "org", "net", "int", "edu", "gov", "mil"],
            include_sites="",
            exclude_sites="",
            pdf_file_input=None,
            api_name="/generate_and_format",
        )

        self.client.predict(text="Polygraf AI (Advanced Model)", api_name="/highlight_visible_1")
        detect_generated = self.client.predict(
            text=generated_text,
            option="Polygraf AI (Advanced Model)",
            api_name="/ai_check",
        )

        humanized_text = self.client.predict(
            text=generated_text,
            model="XL Model",
            temperature=1.2,
            repetition_penalty=1,
            top_k=50,
            length_penalty=1,
            api_name="/humanize",
        )

        self.client.predict(text="Polygraf AI (Advanced Model)", api_name="/highlight_visible_1")
        detect_humanized = self.client.predict(
            text=humanized_text,
            option="Polygraf AI (Advanced Model)",
            api_name="/ai_check",
        )

        print(detect_generated)
        print(detect_humanized)

    @task
    def test_generate_text_without_google_search_and_pdf(self):
        generated_text = self.client.predict(
            input_role="Student",
            topic="Low Resource Large Language Models",
            keywords="",
            article_length=400,
            format="Article",
            writing_style="Formal",
            tone="Professional",
            user_category="General Public",
            depth_of_content="Moderate analysis",
            structure="Introduction, Body, Conclusion",
            references="News outlets",
            num_examples="1-2",
            conclusion_type="Call to Action",
            ai_model="LLaMA 3",
            google_search_check=False,
            year_from="2000",
            month_from="January",
            day_from="01",
            year_to="2024",
            month_to="August",
            day_to="08",
            domains_to_include=["com", "org", "net", "int", "edu", "gov", "mil"],
            include_sites="",
            exclude_sites="",
            pdf_file_input=[handle_file("/home/eljan/article_writer/Abstract.pdf")],
            api_name="/generate_and_format",
        )

        self.client.predict(text="Polygraf AI (Advanced Model)", api_name="/highlight_visible_1")
        detect_generated = self.client.predict(
            text=generated_text,
            option="Polygraf AI (Advanced Model)",
            api_name="/ai_check",
        )

        humanized_text = self.client.predict(
            text=generated_text,
            model="XL Model",
            temperature=1.2,
            repetition_penalty=1,
            top_k=50,
            length_penalty=1,
            api_name="/humanize",
        )

        self.client.predict(text="Polygraf AI (Advanced Model)", api_name="/highlight_visible_1")
        detect_humanized = self.client.predict(
            text=humanized_text,
            option="Polygraf AI (Advanced Model)",
            api_name="/ai_check",
        )

        print(detect_generated[0]["confidences"])
        print(detect_humanized[0]["confidences"])