"""
This script enables creating preference datasets using previous conversation logs
(found under eval/runs/). It will load the chosen conversation log, and for each
turn will:
    - extract the original agent response from the log file
    - generate an alternative candidate response
    - ask the human reviewer to select the preferred response
    - store the preferences in a .jsonl file

How to Run
----------
Instructions for running the simulation are linked here:
https://github.com/Versebuilding/WalkXR-AI/tree/develop/eval
"""

import sys
from pathlib import Path

import json
import argparse
import requests
import random

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--conversation_log_path",
        type=str,
        required=True,
        help="Path to a raw conversation log .jsonl file under eval/runs/<EVAL_ID>/raw/",
    )
    parser.add_argument("--temperature", type=float, default=1.2)
    parser.add_argument("--timeout", type=int, default=120)
    parser.add_argument("--overwrite", action='store_true')
    args = parser.parse_args()

    # Create file to store results
    before, after = args.conversation_log_path.split("/raw/", maxsplit=1)
    results_path = before + "/preferences"
    results_path = Path(results_path)
    results_path.mkdir(parents=True, exist_ok=True)
    results_path = results_path / after
    
    if results_path.exists():
        if args.overwrite:
            # Clear the file
            results_path.write_text("")
        else:
             # Abort if file already exists and --overwrite was not set
            print("""\033[1;31mA preference dataset was already created for this conversation log. Please use the \"--overwrite\" flag if you wish to create the preference dataset again.\033[0m""")
            return

    # Extract json objects from file
    turn = 1
    with open(args.conversation_log_path, 'r') as log_file, open(results_path, "a", encoding="utf-8") as results_file:
        for line in log_file:
            obj = json.loads(line)

            # Print current turn
            print(f"\033[1;32m======== Turn {turn} ========\033[0m")
            print()
            turn += 1

            # Print relevant info to give the human reviewer context
            print("\033[1;36mScenario Prompt:\033[0m")
            print(obj["scenario_prompt"])
            print()
            print("\033[1;36mInitial Emotion:\033[0m")
            print(obj["initial_emotion"])
            print()
            print("\033[1;36mDesired AI Behavior:\033[0m")
            print(obj["desired_ai_behavior"])
            print()
            print("\033[1;36mDesired Tone:\033[0m")
            print(obj["desired_tone"])
            print()
            print("\033[1;36mUser Input:\033[0m")
            user_input = obj["user_input"]
            print(user_input)
            print()

            # Store other necessary variables
            scenario_uid = obj["scenario_id"]
            rep = obj["repeat"]
            run_id = obj["run_id"]
            
            # Store original agent response
            OG_agent_response = obj["agent_response"]

            # Generate alternative agent response
            ALT_agent_response = ""
            response_data = {}
            conversation_session_id = f"{scenario_uid}_rep{rep}_{run_id}"
            api_success = False
            
            for attempt in range(2):
                try:
                    res = requests.post(
                        obj["endpoint"],
                        json={
                            "user_id": "sim_user",
                            "session_id": conversation_session_id,
                            "stage": "demo",
                            "message": user_input,
                            "history": [],
                            "test": {"temperature": args.temperature}
                        },
                        timeout=args.timeout,
                    )
                    res.raise_for_status()
                    response_data = res.json()
                    ALT_agent_response = response_data.get("response_text", "")
                    api_success = True
                    break

                except requests.exceptions.Timeout as e:
                    if attempt == 1:
                        ALT_agent_response = f"API_ERROR: {str(e)}"
                        print(f"API Timeout: {e}")
                    else:
                        print(f"Timeout on attempt {attempt + 1}; retrying once...")
                
                except Exception as e:
                    ALT_agent_response = f"API_ERROR: {str(e)}"
                    print(f"API Error: {e}")
                    break

            # Present both responses in a random order to prevent biased review
            agent_responses = [OG_agent_response, ALT_agent_response]
            # TODO: "None" indicates the default temperature, for future sim log entries it may be beneficial to store the temperature used
            temperature = [None, args.temperature]
            option_1 = random.randint(0, 1)
            option_2 = (option_1 + 1) % 2
            present_order = [agent_responses[option_1], agent_responses[option_2]]
            present_order_temperature = [temperature[option_1], temperature[option_2]]
            print("\033[1;31mAgent Response 1:\033[0m")
            print(present_order[0])
            print()
            print("\033[1;31mAgent Response 2:\033[0m")
            print(present_order[1])
            print()

            # Prompt reviewer to select 1 or 2 depending on which agent_response is better
            while True:
                try:
                    best = int(input("\033[1;35mType 1 or 2 to select the best agent response: \033[0m"))
                    print()
                    if best == 1 or best == 2:
                        break
                except ValueError:
                    print()
                    continue
            
            # Store result as (input, chosen_response, chosen_temperature, rejected_response, rejected_temperature, api_success)
            result = { 
                "user_input": user_input, 
                "chosen_response": present_order[best - 1], 
                "chosen_temperature": present_order_temperature[best - 1],
                "rejected_response": present_order[best % 2],
                "rejected_temperature": present_order_temperature[best % 2],
                "api_success": api_success
            }
            results_file.write(json.dumps(result, ensure_ascii=False) + "\n")
            results_file.flush()

if __name__ == "__main__":
    main()