import re
import os
import json
from pathlib import Path
import time
from dotenv import load_dotenv
import openai

# --- Configuration ---
# Load environment variables from .env file for API keys.
# Create a .env file in the root directory and add your key:
# OPENROUTER_API_KEY="your-key-here"
load_dotenv()

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

# A list of the models we plan to test, using OpenRouter identifiers.
# A quick search on OpenRouter's site is needed to confirm PanGu's availability.
# If PanGu is not on OpenRouter, we may need a separate client for it.
TARGET_MODELS = [
    # "openai/gpt-4o",
    # "openai/gpt-5",
    # "meta-llama/llama-3.1-405b-instruct",
    # "meta-llama/llama-3.1-405b",
    # "anthropic/claude-opus-4.1",
    # "google/gemini-2.5-pro",
    # "x-ai/grok-4",
    # "deepseek/deepseek-r1-0528:free"
    # "qwen/qwen3-235b-a22b",
    # "openai/gpt-oss-20b",
    # "qwen/qwen-2.5-14b",
    # "qwen/qwen3-30b-a3b",
    # "meta-llama/llama-3.3-70b-instruct",
    # "deepseek/deepseek-r1-distill-qwen-14b",
    # "deepseek/deepseek-r1-distill-llama-70b",
    # "z-ai/glm-4-32b"
    # "mistralai/mistral-small-3.2-24b-instruct",
    # "pangu/pangu-model-name", # Placeholder for PanGu - needs verification
]

# Initialize the OpenAI client to point to the OpenRouter API
if OPENROUTER_API_KEY:
    client = openai.OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=OPENROUTER_API_KEY,
    )
else:
    client = None
    print("Warning: OPENROUTER_API_KEY not found. The script will run in simulation mode.")

def parse_tex_file(file_path):
    """
    Parses a LaTeX file to extract prompts and their IDs.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except FileNotFoundError:
        print(f"Error: The file at {file_path} was not found.")
        return []
    prompt_regex = re.compile(
        r"\\item\[Prompt\s+([\d\.]+).*?\]\s*``(.*?)''",
        re.DOTALL
    )
    prompts = []
    matches = prompt_regex.finditer(content)
    for match in matches:
        prompt_id = match.group(1).strip()
        prompt_text = ' '.join(match.group(2).strip().split())
        prompts.append({'id': prompt_id, 'text': prompt_text})
    return prompts

def get_model_response(model_name, prompt_text):
    """
    Gets a response from a specified model via the OpenRouter API.
    """
    if not client:
        time.sleep(0.5)
        return f"This is a simulated response from {model_name} because no API key was provided."

    try:
        print(f"--- Sending to {model_name} ---")
        completion = client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt_text}],

        )
        return completion.choices[0].message.content
    except openai.APIError as e:
        return f"Error: API call failed for {model_name}. Details: {e}"
    except Exception as e:
        return f"An unexpected error occurred for {model_name}: {e}"

def main():
    """
    Main function to execute the script.
    """
    project_root = Path(__file__).parent.parent
    comm_records_dir = project_root / 'AI-comm-records'
    tex_file_path = comm_records_dir / 'prompt_suite.tex'
    prompts_json_path = comm_records_dir / 'prompts.json'
    results_dir = project_root / 'results'

    print("Step 1: Loading prompts...")
    if prompts_json_path.exists():
        print(f"Found cached prompts file at {prompts_json_path}. Loading from JSON.")
        with open(prompts_json_path, 'r', encoding='utf-8') as f:
            extracted_prompts = json.load(f)
    else:
        print(f"No cached prompts file found. Parsing from {tex_file_path}.")
        extracted_prompts = parse_tex_file(tex_file_path)
        if extracted_prompts:
            with open(prompts_json_path, 'w', encoding='utf-8') as f:
                json.dump(extracted_prompts, f, indent=4)
            print(f"Saved extracted prompts to {prompts_json_path}.")

    if not extracted_prompts:
        print("No prompts found. Exiting.")
        return
    print(f"Loaded {len(extracted_prompts)} prompts.\n")

    print("Step 2: Iterating through models and prompts to get responses...")
    for model_name in TARGET_MODELS:
        model_results_dir = results_dir / model_name
        model_results_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"\nProcessing model: {model_name}")
        
        for prompt in extracted_prompts:
            prompt_id = prompt['id']
            prompt_text = prompt['text']
            
            output_file_path = model_results_dir / f"{prompt_id}.txt"

            if output_file_path.exists():
                print(f"Skipping prompt {prompt_id} for {model_name}, already completed.")
                continue

            response = get_model_response(model_name, prompt_text)
            
            with open(output_file_path, 'w', encoding='utf-8') as f:
                f.write(response)
            
            print(f"Saved response for prompt {prompt_id} from {model_name}.")
            time.sleep(1)

    print("\nExperiment complete.")

if __name__ == "__main__":
    main()
