import os
import json
from openai import OpenAI

# Mock classes to match BrowserGym official sample structure
class BrowserGymAction:
    def __init__(self, action_str):
        self.action_str = action_str

def parse_model_action(text):
    # Dummmy parser for official sample
    return text.strip()

# Constants for official sample
SYSTEM_PROMPT = "You are a web browsing agent."
MODEL_NAME = os.getenv("MODEL_NAME", "step-fun-3.5-flash") # Use env var if available
TEMPERATURE = 0.0
MAX_TOKENS = 512
FALLBACK_ACTION = "stop"
MAX_STEPS = 10

def main():
    # Prefer injected API_KEY and API_BASE_URL
    api_key = os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY", "EMPTY")
    base_url = os.getenv("API_BASE_URL") or "https://api.openai.com/v1"
    
    client = OpenAI(
        api_key=api_key, 
        base_url=base_url,
        default_headers={
            "HTTP-Referer": "https://huggingface.co/spaces/armaan020/AegisOpenEnv",
            "X-Title": "AegisOpenEnv Official Sample check"
        } if "openrouter" in base_url.lower() else None
    )
    
    print(f"Official Sample Inference Logic Initialized (Model: {MODEL_NAME})...")
    
    # --- START OF USER PROVIDED SNIPPET ---
    history = []
    task_name = "official_sample_task"
    print(f"[START] task={task_name}", flush=True)
    step_count = 0
    score = 0.5
    try:
        # Dummy loop to represent the user snippet's context
        for step in range(MAX_STEPS):
            # user_content would normally be defined here with AXTree/Accessibility logs
            user_content = [{"type": "text", "text": "Task: Navigate to example.com"}]
            
            messages = [
                {
                    "role": "system",
                    "content": [{"type": "text", "text": SYSTEM_PROMPT}],
                },
                {
                    "role": "user",
                    "content": user_content,
                },
            ]

            try:
                completion = client.chat.completions.create(
                    model=MODEL_NAME,
                    messages=messages,
                    temperature=TEMPERATURE,
                    max_tokens=MAX_TOKENS,
                    stream=False,
                )
                response_text = completion.choices[0].message.content or ""
            except Exception as exc: 
                failure_msg = f"Model request failed ({exc}). Using fallback action."
                print(failure_msg, flush=True)
                response_text = FALLBACK_ACTION

            action_str = parse_model_action(response_text)
            print(f"Step {step}: model suggested -> {action_str}", flush=True)

            step_count += 1
            reward = 0.5  # dummy reward
            score += reward
            print(f"[STEP] step={step_count} reward={reward}", flush=True)

            break # Break here for safety in dummy script

        else:
            print(f"Reached max steps ({MAX_STEPS}).", flush=True)

    finally:
        print(f"[END] task={task_name} score={score} steps={step_count}", flush=True)
    # --- END OF USER PROVIDED SNIPPET ---

if __name__ == "__main__":
    main()