import json
import os
import time
import google.generativeai as genai

from dotenv import load_dotenv
load_dotenv()
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

MAX_RETRIES = 5
RETRY_DELAY_SECONDS = 15  

COMMON_SYSTEM_PROMPT = """You are an expert ml researcher having previous background in SEO and search engines in general. You are working on novel research ideas for next generation of products. These products will have language models augmented with search engines, with the task of answering questions based on sources backed by the search engine. This new set of systems will be collectively called language engines (generative search engines). This will require websites to update their SEO techniques to rank higher in the llm generated answer. Specifically they will use GEO (Generative Engine Optimization) techniques to boost their visibility in the final text answer outputted by the Language Engine.  
"""

def call_gemini(user_prompt, system_prompt=COMMON_SYSTEM_PROMPT, model_name="gemini-2.5-pro", temperature=0.7):

    messages = f"system_prompt: {system_prompt}\n\n user_prompt:{user_prompt}"
    
    for attempt in range(MAX_RETRIES):
        try:
            print(f"Running Gemini Model (call_gemini) - Attempt {attempt + 1}/{MAX_RETRIES}")
            model = genai.GenerativeModel(model_name)
            response = model.generate_content(messages)
            print("Response Done")
            return response.text
        except Exception as e:
            print(f"Error in calling Gemini API (call_gemini): {e}")
            if attempt < MAX_RETRIES - 1:
                print(f"Retrying in {RETRY_DELAY_SECONDS} seconds...")
                time.sleep(RETRY_DELAY_SECONDS)
            else:
                print(f"API call failed after {MAX_RETRIES} attempts.")
                raise  

def new_autorule_mine(target_document, ruleset):
    
    geobench_rules_list = ruleset

    rules_string = "- " + "\n- ".join(geobench_rules_list)


    user_prompt = f"""
    Here is the source:
    {target_document}

    You are given a website document as a source. This source, along with other sources, will be used by a language model (LLM) to generate answers to user questions, with each line in the generated answer being cited with its original source. Your task, as the owner of the source, is to **rewrite your document in a way that maximizes its visibility and impact in the LLM's final answer, ensuring your source is more likely to be quoted and cited**.

    Your can regenerate the provided source so that it strictly adheres to the "Quality Guidelines", and you can also apply any other methods or techniques, as long as they help your rewritten source text rank higher in terms of relevance, authority, and impact in the LLM's generated answers.

    ## Quality Guidelines to Follow:

    {rules_string}

    """.strip()
    
    return call_gemini(user_prompt)

if __name__ == "__main__":
    with open("sample.txt", "r", encoding="utf-8") as f:
        target_document = f.read()

    with open("Prompts/gemini_RG.json", "r", encoding="utf-8") as f:
        ruleset = json.load(f)["rule"]

    rewritten_doc = new_autorule_mine(target_document, ruleset)
    print("Rewritten Document:\n", rewritten_doc)
