import transformers
import torch
import random
from datasets import load_dataset
import requests
import json
import time
import gc
import os

# Configuration variables
INPUT_FILE = "refusal_datasets/arditi_harmful_full.json"
OUTPUT_FILE = "refusal_responses/llama3b_ppo_web/llama_refusal_full_search_prefill_1_loop_answer_web.json"

# SerpAPI Configuration
SERPAPI_KEY = os.getenv("SERPAPI_KEY")  # Set your SerpAPI key as environment variable
if not SERPAPI_KEY:
    raise ValueError("SERPAPI_KEY environment variable is required. Please set it before running the script.")
SERPAPI_URL = "https://serpapi.com/search"

# Rate limiting
class RateLimitTracker:
    def __init__(self):
        self.query_count = 0
        self.start_time = time.time()
    
    def can_make_query(self):
        # Basic rate limiting - adjust based on your SerpAPI plan
        return True  # SerpAPI handles rate limiting on their end
    
    def record_query(self):
        self.query_count += 1
        if self.query_count % 10 == 0:
            elapsed = time.time() - self.start_time
            print(f"Made {self.query_count} queries in {elapsed:.1f} seconds")

# Global rate tracker
rate_tracker = RateLimitTracker()


# Model ID and device setup
model_id = "your-username/your-rl-trained-search-model-llama3.2-3b"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

curr_eos = [128001, 128009] # for Llama series models
curr_search_template = '\n\n{output_text}<information>{search_results}</information>\n\n'

# Initialize the tokenizer and model
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
# model = transformers.AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    device_map="auto",
    attn_implementation="eager"  # Disable flash attention
)

# Define the custom stopping criterion
class StopOnSequence(transformers.StoppingCriteria):
    def __init__(self, target_sequences, tokenizer):
        # Encode the string so we have the exact token-IDs pattern
        self.target_ids = [tokenizer.encode(target_sequence, add_special_tokens=False) for target_sequence in target_sequences]
        self.target_lengths = [len(target_id) for target_id in self.target_ids]
        self._tokenizer = tokenizer

    def __call__(self, input_ids, scores, **kwargs):
        # Make sure the target IDs are on the same device
        targets = [torch.as_tensor(target_id, device=input_ids.device) for target_id in self.target_ids]

        if input_ids.shape[1] < min(self.target_lengths):
            return False

        # Compare the tail of input_ids with our target_ids
        for i, target in enumerate(targets):
            if torch.equal(input_ids[0, -self.target_lengths[i]:], target):
                return True

        return False

def get_query(text):
    import re
    pattern = re.compile(r"<search>(.*?)</search>", re.DOTALL)
    matches = pattern.findall(text)
    if matches:
        return matches[-1]
    else:
        return None

def search(query: str):
    """Perform web search using SerpAPI with Google search"""
    try:
        # Check rate limits
        if not rate_tracker.can_make_query():
            return "Rate limit exceeded, please wait"
        
        # SerpAPI parameters for Google search
        params = {
            "engine": "google",
            "q": query,
            "api_key": SERPAPI_KEY,
            "num": 3,  # Number of results to return
            "safe": "active"  # Safe search
        }
        
        # Make request to SerpAPI
        response = requests.get(SERPAPI_URL, params=params)
        response.raise_for_status()
        
        # Record the query
        rate_tracker.record_query()
        
        data = response.json()
        
        # Extract search results
        results = []
        
        # Get answer box if available (direct answers)
        answer_box = data.get('answer_box', {})
        if answer_box:
            title = answer_box.get('title', 'Direct Answer')
            snippet = answer_box.get('snippet', answer_box.get('answer', 'No snippet available'))
            results.append({
                'title': title,
                'content': snippet
            })
        
        # Get organic search results
        organic_results = data.get('organic_results', [])
        for result in organic_results[:3]:  # Top 3 results
            title = result.get('title', 'No title')
            snippet = result.get('snippet', 'No snippet available')
            results.append({
                'title': title,
                'content': snippet
            })
        
        # Format results similar to local search
        format_reference = ''
        for idx, result in enumerate(results):
            format_reference += f"Doc {idx+1}(Title: {result['title']}) {result['content']}\n"
        
        return format_reference
        
    except requests.exceptions.RequestException as e:
        print(f"SerpAPI request error: {e}")
        return f"Search error: Unable to fetch results for '{query}'"
    except Exception as e:
        print(f"Search error: {e}")
        return f"Search error: {str(e)}"

# Initialize the stopping criteria
target_sequences = ["</search>", " </search>", "</search>\n", " </search>\n", "</search>\n\n", " </search>\n\n"]
stopping_criteria = transformers.StoppingCriteriaList([StopOnSequence(target_sequences, tokenizer)])

def prepare_prompt(question_text):
    """Prepare a single prompt for a question"""
    question = question_text.strip()
    if question[-1] != '?':
        question += '?'
    
    # Prepare the message - exactly the same as infer_search.py
    prompt = f"""Answer the given question. \
You must conduct reasoning inside <think> and </think> every time you get new information. \
After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. \
You can search as many times as your want. \
If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. Question: {question}\n"""

    if tokenizer.chat_template:
        prompt = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False)
    
    return prompt

def process_single_question(question_text):
    """Process a single question using the same logic as infer_search.py"""
    
    prompt = prepare_prompt(question_text)
    
    print(f'\n\n################# [Start Reasoning + Searching] ##################\n\n')
    print(f"Question: {question_text}")
    
    # Initialize variables exactly like infer_search.py
    cnt = 0
    full_response = ""
    current_prompt = prompt
    search_information = []  # Store all search queries and results
    
    # Process the question with potential search iterations - same logic as infer_search.py
    max_searches = 10  # Maximum number of searches per question
    while cnt < max_searches:
        # Prefill with "<search>" to force it as the first token
        prefilled_prompt = current_prompt + "<search>"
        input_ids = tokenizer.encode(prefilled_prompt, return_tensors='pt').to(device)
        attention_mask = torch.ones_like(input_ids)
        
        # Generate text with the stopping criteria (starting after the prefilled "<search>")
        outputs = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=4096*4,
            stopping_criteria=stopping_criteria,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=False, # Greedy decoding (temperature=0.0)
            use_cache=True  # Enable KV caching for faster generation
        )

        if outputs[0][-1].item() in curr_eos:
            generated_tokens = outputs[0][input_ids.shape[1]:]
            output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
            # Add the prefilled "<search>" to the response since it's not in generated_tokens
            full_response += "<search>" + output_text
            print("<search>" + output_text)
            break

        generated_tokens = outputs[0][input_ids.shape[1]:]
        output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
        # Add the prefilled "<search>" to the response since it's not in generated_tokens
        full_response += "<search>" + output_text
        
        tmp_query = get_query(tokenizer.decode(outputs[0], skip_special_tokens=True))
        if tmp_query:
            # print(f'searching "{tmp_query}"...')
            search_results = search(tmp_query)
            # Store the search query 
            search_information.append({
                "query": tmp_query
            })
        else:
            search_results = ''

        search_text = curr_search_template.format(output_text="<search>" + output_text, search_results=search_results)
        current_prompt += search_text
        cnt += 1
        print(f"Search {cnt}/{max_searches}: {search_text}")
    
    # Check if we hit the search limit
    if cnt >= max_searches:
        print(f"⚠️ Reached maximum searches ({max_searches}), forcing answer generation")
        # Prefill with "<answer>" to force answer generation
        prefilled_prompt = current_prompt + "<answer>"
        input_ids = tokenizer.encode(prefilled_prompt, return_tensors='pt').to(device)
        attention_mask = torch.ones_like(input_ids)
        
        # Generate final answer
        outputs = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=4096*4,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=False,
            use_cache=True
        )
        
        generated_tokens = outputs[0][input_ids.shape[1]:]
        output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
        full_response += "<answer>" + output_text
        print("<answer>" + output_text)
    
    # Clear GPU memory after processing
    torch.cuda.empty_cache()
    gc.collect()
    
    return full_response, search_information

def process_questions_sequential(questions, questions_data, output_file, save_interval=10):
    """Process questions sequentially with periodic saving"""
    
    results = []
    
    # Process questions individually (search requires sequential processing)
    for i, question in enumerate(questions):
        # Hard condition: end at question index 199
        if i > 198:
            continue
        print(f'\n\n################# [Processing Question {i+1}/{len(questions)}] ##################\n\n')
        
        try:
            # Process the question using the search logic
            response, search_info = process_single_question(question)
            
            # Create result entry
            result_entry = {
                "question": question,
                "response": response,
                "search_information": search_info,
                "question_index": i,
            }
            results.append(result_entry)
            
            print(f"Question {i+1}: {question[:100]}...")
            print(f"Response: {response[:200]}...")
            print(f"Search queries: {len(search_info)}")
            print("-" * 50)
            
            # Save progress every 10 questions
            if (i + 1) % 10 == 0 or (i + 1) == len(questions):
                print(f"\nSaving progress... ({i+1}/{len(questions)} questions)")
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(results, f, indent=2, ensure_ascii=False)
            
            # Small delay to prevent overwhelming the search service and respect rate limits
            time.sleep(2)  # Increased delay for SerpAPI rate limits
            
        except Exception as e:
            print(f"Error processing question {i+1}: {e}")
            result_entry = {
                "question": question,
                "response": f"ERROR: {str(e)}",
                "search_information": [],
                "question_index": i,
            }
            results.append(result_entry)
    
    return [r["response"] for r in results]

def main():
    # Load the questions from the JSON file
    print(f"Loading questions from {INPUT_FILE}...")
    with open(INPUT_FILE, 'r', encoding='utf-8') as f:
        questions_data = json.load(f)
    
    print(f"Found {len(questions_data)} questions to process")
    
    # Extract questions
    questions = [item.get("instruction", "") for item in questions_data if item.get("instruction", "")]
    
    print(f"Processing {len(questions)} valid questions sequentially...")
    

    
    try:
        # Process all questions sequentially with periodic saving
        all_responses = process_questions_sequential(questions, questions_data, OUTPUT_FILE, save_interval=10)
        
        print(f"Processing complete! Results saved to {OUTPUT_FILE}")
        print(f"Successfully processed {len(all_responses)} questions")
        
    except Exception as e:
        print(f"Error during sequential processing: {e}")
        print("Falling back to individual processing...")
        
        # Fallback to individual processing
        results = []
        for i, item in enumerate(questions_data):
            # Hard condition: end at question index 199
            if i > 198:
                continue 
            # Try both field names to be safe
            question = item.get("instruction", "") or item.get("question", "")
            if not question:
                continue
                
            print(f"Processing question {i+1}/{len(questions_data)}")
            
            try:
                # Process single question
                response, search_info = process_single_question(question)
                
                result_entry = {
                    "question": question,
                    "response": response,
                    "search_information": search_info,
                    "question_index": i,
                }
                results.append(result_entry)
                
                # Save progress every 10 questions
                if (i + 1) % 10 == 0:
                    print(f"Saving progress... ({i+1}/{len(questions_data)})")
                    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
                        json.dump(results, f, indent=2, ensure_ascii=False)
                
            except Exception as individual_error:
                print(f"Error processing question {i+1}: {individual_error}")
                result_entry = {
                    "question": question,
                    "response": f"ERROR: {str(individual_error)}",
                    "search_information": [],
                    "question_index": i,
                }
                results.append(result_entry)
        
        # Save final results from fallback
        print(f"\nSaving final results to {OUTPUT_FILE}...")
        with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"Fallback processing complete! Results saved to {OUTPUT_FILE}")
        print(f"Successfully processed {len(results)} questions")

if __name__ == "__main__":
    main() 