import sys
import os
import pyperclip

class GrammarCorrector:
    def __init__(self, api_key):
        self.api_key = api_key

    def correct_grammar(self, text):
        raise NotImplementedError("This method should be implemented by subclasses.")

    def rephrase_text(self, text):
        raise NotImplementedError("This method should be implemented by subclasses.")

class GroqCorrector(GrammarCorrector):
    def __init__(self, api_key):
        super().__init__(api_key)
        from groq import Groq
        self.client = Groq(api_key=api_key)

    def correct_grammar(self, text):
        response = self.client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": f"Correct the grammar, do not explain: '{text}'",
                }
            ],
            model="mixtral-8x7b-32768",
        )
        corrected_text = response.choices[0].message.content
        return corrected_text

    def rephrase_text(self, text):
        response = self.client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": f"Rephrase the sentence: '{text}'",
                }
            ],
            model="mixtral-8x7b-32768",
        )
        rephrased_text = response.choices[0].message.content
        return rephrased_text
        

class OpenAICorrector(GrammarCorrector):
    def __init__(self, api_key):
        super().__init__(api_key)
        import openai
        openai.api_key = api_key
        self.client = openai

    def correct_grammar(self, text):
        response = self.client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "role": "system",
                    "content": "You are a highly skilled assistant tasked with correcting grammar mistakes."
                },
                {
                    "role": "user",
                    "content": text,
                }
            ],
        )
        corrected_text = response.choices[0].message.content
        return corrected_text
        
    def rephrase_text(self, text):
        # Assuming a similar process for rephrasing
        response = self.client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "role": "system",
                    "content": "You are a highly skilled assistant tasked with rephrasing sentences for clarity and conciseness."
                },
                {
                    "role": "user",
                    "content": text,
                }
            ],
        )
        rephrased_text = response.choices[0].message.content
        return rephrased_text

if __name__ == "__main__":
    api_key = os.environ.get("GROQ_API_KEY")
    corrector_type = "groq"
    # api_key = os.environ.get("OPENAI_API_KEY")
    # corrector_type = "openai"

    if corrector_type == "groq":
        corrector = GroqCorrector(api_key)
    elif corrector_type == "openai":
        corrector = OpenAICorrector(api_key)
    else:
        raise ValueError("Unsupported corrector type")

    if len(sys.argv) > 1:
        text_to_correct = " ".join(sys.argv[1:])
        corrected_text = corrector.correct_grammar(text_to_correct)
        print("Corrected:", corrected_text)
        
        print("---------------------------------")
        rephrased_text = corrector.rephrase_text(text_to_correct)
        print("Rephrased:", rephrased_text)
        
        pyperclip.copy(corrected_text)
    else:
        print("Usage: c 'text to correct'")