import os
import json
import base64
import io
import argparse
import logging
import gradio as gr
import openai
import gymnasium as gym
import browsergym.core
from PIL import Image
import numpy as np
from browsergym.core.action.highlevel import HighLevelActionSet
from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, prune_html
from browsergym.experiments import Agent
from dotenv import load_dotenv
import cv2

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('browser_agent.log')
    ]
)
logger = logging.getLogger(__name__)

load_dotenv()

# Set your OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

# Example instructions to display
EXAMPLES = [
    "Search for the latest AI news on Google",
    "Go to Wikipedia and find the population of Seoul",
    "Open YouTube and play the top trending video",
]

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise argparse.ArgumentTypeError("Boolean value expected.")

def parse_args():
    parser = argparse.ArgumentParser(description="Run BrowserGym web agent.")
    parser.add_argument(
        "--model_name",
        type=str,
        default="gpt-4o",
        help="OpenAI model name.",
    )
    parser.add_argument(
        "--start_url",
        type=str,
        default="https://www.duckduckgo.com",
        help="Starting URL for the openended task.",
    )
    parser.add_argument(
        "--visual_effects",
        type=str2bool,
        default=True,
        help="Add visual effects when the agent performs actions.",
    )
    parser.add_argument(
        "--use_html",
        type=str2bool,
        default=False,
        help="Use HTML in the agent's observation space.",
    )
    parser.add_argument(
        "--use_axtree",
        type=str2bool,
        default=True,
        help="Use AXTree in the agent's observation space.",
    )
    parser.add_argument(
        "--use_screenshot",
        type=str2bool,
        default=False,
        help="Use screenshot in the agent's observation space.",
    )
    parser.add_argument(
        "--log_level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Set the logging level.",
    )
    return parser.parse_args()

def image_to_jpg_base64_url(image: np.ndarray | Image.Image):
    """Convert a numpy array to a base64 encoded image url."""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if image.mode in ("RGBA", "LA"):
        image = image.convert("RGB")

    with io.BytesIO() as buffer:
        image.save(buffer, format="JPEG")
        image_base64 = base64.b64encode(buffer.getvalue()).decode()

    return f"data:image/jpeg;base64,{image_base64}"

class BrowserAgent(Agent):
    def obs_preprocessor(self, obs: dict) -> dict:
        return {
            "chat_messages": obs["chat_messages"],
            "screenshot": obs["screenshot"],
            "goal_object": obs["goal_object"],
            "last_action": obs["last_action"],
            "last_action_error": obs["last_action_error"],
            "open_pages_urls": obs["open_pages_urls"],
            "open_pages_titles": obs["open_pages_titles"],
            "active_page_index": obs["active_page_index"],
            "axtree_txt": flatten_axtree_to_str(obs["axtree_object"], filter_visible_only=True, extra_properties=obs['extra_element_properties']),
            "pruned_html": prune_html(flatten_dom_to_str(obs["dom_object"])),
        }

    def __init__(self, model_name: str = "gpt-4o", use_html: bool = False, use_axtree: bool = True, use_screenshot: bool = False):
        super().__init__()
        logger.info(f"Initializing BrowserAgent with model: {model_name}")
        logger.info(f"Observation space: HTML={use_html}, AXTree={use_axtree}, Screenshot={use_screenshot}")
        
        self.model_name = model_name
        self.use_html = use_html
        self.use_axtree = use_axtree
        self.use_screenshot = use_screenshot
        
        if not (use_html or use_axtree):
            raise ValueError("Either use_html or use_axtree must be set to True.")
        
        self.openai_client = openai.OpenAI()
        
        self.action_set = HighLevelActionSet(
            subsets=["chat", "tab", "nav", "bid", "infeas"],
            strict=False,
            multiaction=False,
            demo_mode="default"
        )
        self.action_history = []

    def get_action(self, obs: dict) -> tuple[str, dict]:
        logger.debug("Preparing action request")
        
        system_msgs = [{
            "type": "text",
            "text": """\
# Instructions

You are a UI Assistant, your goal is to help the user perform tasks using a web browser. You can
communicate with the user via a chat, to which the user gives you instructions and to which you
can send back messages. You have access to a web browser that both you and the user can see,
and with which only you can interact via specific commands.

Review the instructions from the user, the current state of the page and all other information
to find the best possible next action to accomplish your goal. Your answer will be interpreted
and executed by a program, make sure to follow the formatting instructions.
"""
        }]

        user_msgs = []
        
        # Add chat messages
        user_msgs.append({
            "type": "text",
            "text": "# Chat Messages\n"
        })
        for msg in obs["chat_messages"]:
            if msg["role"] in ("user", "assistant", "infeasible"):
                user_msgs.append({
                    "type": "text",
                    "text": f"- [{msg['role']}] {msg['message']}\n"
                })
                logger.debug(f"Added chat message: [{msg['role']}] {msg['message']}")
            elif msg["role"] == "user_image":
                user_msgs.append({"type": "image_url", "image_url": msg["message"]})
                logger.debug("Added user image message")

        # Add open tabs info
        user_msgs.append({
            "type": "text",
            "text": "# Currently open tabs\n"
        })
        for page_index, (page_url, page_title) in enumerate(
            zip(obs["open_pages_urls"], obs["open_pages_titles"])
        ):
            user_msgs.append({
                "type": "text",
                "text": f"""\
Tab {page_index}{" (active tab)" if page_index == obs["active_page_index"] else ""}
  Title: {page_title}
  URL: {page_url}
"""
            })
            logger.debug(f"Added tab info: {page_title} ({page_url})")

        # Add accessibility tree if enabled
        if self.use_axtree:
            user_msgs.append({
                "type": "text",
                "text": f"""\
# Current page Accessibility Tree

{obs["axtree_txt"]}

"""
            })
            logger.debug("Added accessibility tree")

        # Add HTML if enabled
        if self.use_html:
            user_msgs.append({
                "type": "text",
                "text": f"""\
# Current page DOM

{obs["pruned_html"]}

"""
            })
            logger.debug("Added HTML DOM")

        # Add screenshot if enabled
        if self.use_screenshot:
            user_msgs.append({
                "type": "text",
                "text": "# Current page Screenshot\n"
            })
            user_msgs.append({
                "type": "image_url",
                "image_url": {
                    "url": image_to_jpg_base64_url(obs["screenshot"]),
                    "detail": "auto"
                }
            })
            logger.debug("Added screenshot")

        # Add action space description
        user_msgs.append({
            "type": "text",
            "text": f"""\
# Action Space

{self.action_set.describe(with_long_description=False, with_examples=True)}

Here are examples of actions with chain-of-thought reasoning:

I now need to click on the Submit button to send the form. I will use the click action on the button, which has bid 12.
```click("12")```

I found the information requested by the user, I will send it to the chat.
```send_msg_to_user("The price for a 15\\" laptop is 1499 USD.")```

"""
        })

        # Add action history and errors
        if self.action_history:
            user_msgs.append({
                "type": "text",
                "text": "# History of past actions\n"
            })
            for action in self.action_history:
                user_msgs.append({
                    "type": "text",
                    "text": f"\n{action}\n"
                })
                logger.debug(f"Added past action: {action}")

            if obs["last_action_error"]:
                user_msgs.append({
                    "type": "text",
                    "text": f"""\
# Error message from last action

{obs["last_action_error"]}

"""
                })
                logger.warning(f"Last action error: {obs['last_action_error']}")

        # Ask for next action
        user_msgs.append({
            "type": "text",
            "text": """\
# Next action

You will now think step by step and produce your next best action. Reflect on your past actions, any resulting error message, and the current state of the page before deciding on your next action.
"""
        })

        # Log the full prompt for debugging
        prompt_text_strings = []
        for message in system_msgs + user_msgs:
            match message["type"]:
                case "text":
                    prompt_text_strings.append(message["text"])
                case "image_url":
                    image_url = message["image_url"]
                    if isinstance(message["image_url"], dict):
                        image_url = image_url["url"]
                    if image_url.startswith("data:image"):
                        prompt_text_strings.append(
                            "image_url: " + image_url[:30] + "... (truncated)"
                        )
                    else:
                        prompt_text_strings.append("image_url: " + image_url)
                case _:
                    raise ValueError(
                        f"Unknown message type {repr(message['type'])} in the task goal."
                    )
        full_prompt_txt = "\n".join(prompt_text_strings)
        logger.debug(full_prompt_txt)

        # Query OpenAI model
        logger.info("Sending request to OpenAI")
        response = self.openai_client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "system", "content": system_msgs},
                {"role": "user", "content": user_msgs}
            ]
        )
        action = response.choices[0].message.content
        logger.info(f"Received action from OpenAI: {action}")
        self.action_history.append(action)
        return action, {}

def run_agent(instruction: str, model_name: str = "gpt-4o", start_url: str = "https://www.duckduckgo.com",
              use_html: bool = False, use_axtree: bool = True, use_screenshot: bool = False):
    logger.info(f"Starting agent with instruction: {instruction}")
    logger.info(f"Configuration: model={model_name}, start_url={start_url}")
    
    trajectory = []
    agent = BrowserAgent(
        model_name=model_name,
        use_html=use_html,
        use_axtree=use_axtree,
        use_screenshot=use_screenshot
    )

    # Initialize BrowserGym environment
    logger.info("Initializing BrowserGym environment")
    env = gym.make(
        "browsergym/openended",
        task_kwargs={
            "start_url": start_url,
            "task": "openended",  # Required task parameter
            "goal": instruction,
        },
        wait_for_user_message=True
    )
    obs, info = env.reset()
    logger.info("Environment initialized")

    # Send user instruction to the environment
    logger.info("Sending user instruction to environment")
    obs, reward, terminated, truncated, info = env.step({
        "type": "send_msg_to_user",
        "message": instruction
    })
    processed_obs = agent.obs_preprocessor(obs)
    logger.info(f"Obs: {processed_obs.keys()}")
    logger.info(f"axtree_txt: {processed_obs['axtree_txt']}")

    # 초기 상태 yield
    trajectory.append((obs['screenshot'], "Initial state"))
    yield obs['screenshot'], trajectory.copy()

    try:
        step_count = 0
        while True:
            logger.info(f"Step {step_count}: Getting next action")
            # Get next action from agent
            action, _ = agent.get_action(processed_obs)
            
            # Execute action
            logger.info(f"Step {step_count}: Executing action: {action}")
            obs, reward, terminated, truncated, info = env.step(action)
            processed_obs = agent.obs_preprocessor(obs)
            
            # trajectory에 numpy array 직접 저장
            trajectory.append((obs['screenshot'], action))
            logger.info(f"Step {step_count}: Saved screenshot and updated trajectory")
            step_count += 1

            # 매 step마다 yield
            yield obs['screenshot'], trajectory.copy()

            if terminated or truncated:
                logger.info(f"Episode ended: terminated={terminated}, truncated={truncated}")
                break

    finally:
        logger.info("Closing environment")
        env.close()

def main():
    args = parse_args()
    
    # Set logging level from command line argument
    logger.setLevel(getattr(logging, args.log_level))
    logger.info("Starting BrowserGym web agent")
    logger.info(f"Arguments: {args}")
    
    with gr.Blocks(title="🎯 Web Agent Demo with BrowserGym & OpenAI") as demo:
        gr.Markdown("# Web Agent Demo (BrowserGym + OpenAI)")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## Examples")
                gr.Examples(
                    examples=[[e] for e in EXAMPLES],
                    inputs=[gr.Textbox(label="Instruction")],
                    cache_examples=False,
                )
            with gr.Column(scale=2):
                instruction = gr.Textbox(
                    label="Enter your instruction here",
                    placeholder="E.g., 'Search for AI then click #result-stats'",
                    lines=2,
                )
                model_name = gr.Dropdown(
                    label="Model",
                    choices=["gpt-4o", "gpt-4o-mini"],
                    value=args.model_name
                )
                run_btn = gr.Button("Run Agent")
                browser_view = gr.Image(label="Browser View")
            with gr.Column(scale=2):
                gr.Markdown("## Trajectory History")
                trajectory_gallery = gr.Gallery(label="Action & State", columns=2)

        run_btn.click(
            fn=run_agent,
            inputs=[instruction, model_name],
            outputs=[browser_view, trajectory_gallery],
            api_name="run_agent",
            show_progress=True,
            concurrency_limit=1
        )

    logger.info("Launching Gradio interface")
    demo.launch()

if __name__ == "__main__":
    main()