# ========= Copyright 2023-2026 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2026 @ CAMEL-AI.org. All Rights Reserved. =========
import json
import re
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

from camel.logger import get_logger
from camel.models import BaseModelBackend, ModelFactory
from camel.types import ModelPlatformType, ModelType

from .actions import ActionExecutor
from .browser_session import HybridBrowserSession

if TYPE_CHECKING:
    from camel.agents import ChatAgent

logger = get_logger(__name__)


class PlaywrightLLMAgent:
    r"""High-level orchestration layer connecting page snapshots, an LLM,
    and an action executor for browser automation.

    Args:
        user_data_dir (str, optional): Path to the browser user data
            directory for persistent sessions. (default: :obj:`None`)
        headless (bool, optional): Whether to run the browser in headless
            mode. (default: :obj:`False`)
        stealth (bool, optional): Whether to enable stealth mode to avoid
            bot detection. (default: :obj:`False`)
        model_backend (BaseModelBackend, optional): The model backend to
            use for LLM calls. If :obj:`None`, defaults to
            :obj:`ModelFactory` with :obj:`ModelPlatformType.DEFAULT` and
            :obj:`ModelType.DEFAULT`. (default: :obj:`None`)
        default_timeout (int, optional): Default timeout in milliseconds
            for browser actions. If :obj:`None`, uses the configured
            default. (default: :obj:`None`)
        short_timeout (int, optional): Short timeout in milliseconds for
            quick browser operations. If :obj:`None`, uses the configured
            default. (default: :obj:`None`)
    """

    # System prompt as class constant to avoid recreation
    SYSTEM_PROMPT = """
You are a web automation assistant.

" Analyse the page snapshot and create a short high-level plan, "
"then output the FIRST action to start with.\n\n"
"Return a JSON object in *exactly* this shape:\n"
"Action format json_object examples:\n"
"{\n  \"plan\": [\"Step 1\", \"Step 2\"],\n  \"action\": {\n    \"type\":
\"click\",\n    \"ref\": \"e1\"\n  }\n}\n\n"
"If task is already complete:\n"
"{\n  \"plan\": [],\n  \"action\": {\n    \"type\": \"finish\",
\n    \"ref\": null,\n    \"summary\": \"Task was already completed. Summary
of what was found...\"\n  }\n}"

Available action types:
- 'click': {"type": "click", "ref": "e1"} or {"type": "click", "text":
"Button Text"} or {"type": "click", "selector": "button"}
- 'type': {"type": "type", "ref": "e1", "text": "search text"} or {"type":
"type", "selector": "input", "text": "search text"}
- 'select': {"type": "select", "ref": "e1", "value": "option"} or {"type":
"select", "selector": "select", "value": "option"}
- 'wait': {"type": "wait", "timeout": 2000} or {"type": "wait", "selector":
"#element"}
- 'scroll': {"type": "scroll", "direction": "down", "amount": 300}
- 'enter': {"type": "enter", "ref": "e1"} or {"type": "enter", "selector":
"input[name=q]"} or {"type": "enter"}
- 'navigate': {"type": "navigate", "url": "https://example.com"}
- 'finish': {"type": "finish", "ref": null, "summary": "task completion
summary"}

IMPORTANT:
- For 'click': Use 'ref' from snapshot, or 'text' for visible text,
or 'selector' for CSS selectors
- For 'type'/'select': Use 'ref' from snapshot or 'selector' for CSS selectors
- Only use 'ref' values that exist in the snapshot (e.g., ref=e1, ref=e2, etc.)
- Use 'finish' when the task is completed successfully with a summary of
what was accomplished
- Use 'enter' to press the Enter key (optionally focus an element first)
- Use 'navigate' to open a new URL before interacting further
- click can choose radio, checkbox...
        """

    def __init__(
        self,
        *,
        user_data_dir: Optional[str] = None,
        headless: bool = False,
        stealth: bool = False,
        model_backend: Optional[BaseModelBackend] = None,
        default_timeout: Optional[int] = None,
        short_timeout: Optional[int] = None,
    ):
        self._session = HybridBrowserSession(
            headless=headless,
            user_data_dir=user_data_dir,
            stealth=stealth,
            default_timeout=default_timeout,
            short_timeout=short_timeout,
        )
        from camel.agents import ChatAgent

        # Populated lazily after first page load
        self.action_history: List[Dict[str, Any]] = []
        if model_backend is None:
            model_backend = ModelFactory.create(
                model_platform=ModelPlatformType.DEFAULT,
                model_type=ModelType.DEFAULT,
                model_config_dict={"temperature": 0, "top_p": 1},
            )
        self.model_backend = model_backend
        # Reuse ChatAgent instance to avoid recreation overhead
        self._chat_agent: Optional[ChatAgent] = None

    async def navigate(self, url: str) -> str:
        r"""Navigate to a URL and return the page snapshot.

        Args:
            url (str): The URL to navigate to.

        Returns:
            str: The page snapshot after navigation, or an error message
                if navigation fails.
        """
        try:
            # HybridBrowserSession handles waits internally
            logger.debug("Navigated to URL: %s", url)
            await self._session.visit(url)
            return await self._session.get_snapshot(force_refresh=True)
        except Exception as exc:
            error_msg = f"Error: could not navigate to {url} - {exc}"
            logger.error(error_msg)
            return error_msg

    def _get_chat_agent(self) -> "ChatAgent":
        r"""Get or create the :obj:`ChatAgent` instance.

        Returns:
            ChatAgent: The existing or newly created :obj:`ChatAgent`
                initialized with the system prompt and model backend.
        """
        from camel.agents import ChatAgent

        if self._chat_agent is None:
            self._chat_agent = ChatAgent(
                system_message=self.SYSTEM_PROMPT, model=self.model_backend
            )
        return self._chat_agent

    def _safe_parse_json(self, content: str) -> Dict[str, Any]:
        r"""Safely parse JSON from an LLM response with multiple fallback
        strategies.

        Attempts direct parsing first, then regex-based extraction, then
        line-by-line parsing. Falls back to a default structure if all
        strategies fail.

        Args:
            content (str): The raw LLM response string to parse.

        Returns:
            Dict[str, Any]: The parsed JSON as a dictionary, or a fallback
                response structure if parsing fails.
        """
        # First attempt: direct parsing
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Second attempt: extract JSON-like block using regex
        # Look for content between outermost braces
        json_pattern = re.compile(
            r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', re.DOTALL
        )
        json_matches = json_pattern.findall(content)

        for match in json_matches:
            try:
                return json.loads(match)
            except json.JSONDecodeError:
                continue

        # Third attempt: try to find and parse line by line
        lines = content.split('\n')
        json_lines = []
        in_json = False

        for line in lines:
            line = line.strip()
            if line.startswith('{'):
                in_json = True
                json_lines = [line]
            elif in_json:
                json_lines.append(line)
                if line.endswith('}'):
                    try:
                        json_text = '\n'.join(json_lines)
                        return json.loads(json_text)
                    except json.JSONDecodeError:
                        pass
                    in_json = False
                    json_lines = []

        # Fallback: return default structure
        logger.warning(
            "Could not parse JSON from LLM response: %s", content[:200]
        )
        return self._get_fallback_response("Parsing error")

    def _get_fallback_response(self, error_msg: str) -> Dict[str, Any]:
        r"""Generate a fallback response structure for parsing failures.

        Args:
            error_msg (str): A description of the error that caused the
                fallback.

        Returns:
            Dict[str, Any]: A default response dictionary with a
                ``finish`` action and the error message as summary.
        """
        return {
            "plan": [f"Could not parse response: {error_msg}"],
            "action": {
                "type": "finish",
                "ref": None,
                "summary": f"Parsing error: {error_msg}",
            },
        }

    def _llm_call(
        self,
        prompt: str,
        snapshot: str,
        is_initial: bool,
        history: Optional[List[Dict[str, Any]]] = None,
    ) -> Dict[str, Any]:
        r"""Call the LLM via :obj:`ChatAgent` to get a plan and next action.

        Args:
            prompt (str): The user task description.
            snapshot (str): The current page snapshot text.
            is_initial (bool): Whether this is the first LLM call for the
                task. If :obj:`True`, no action history is included.
            history (List[Dict[str, Any]], optional): List of previous
                actions and their results. Only used when ``is_initial``
                is :obj:`False`. (default: :obj:`None`)

        Returns:
            Dict[str, Any]: The parsed LLM response containing ``plan``
                (List[str]) and ``action`` (Dict[str, Any]).
        """
        # Build user message
        if is_initial:
            user_content = f"Snapshot:\n{snapshot}\n\nTask: {prompt}"
        else:
            hist_lines = [
                (
                    f"{i + 1}. {'✅' if h['success'] else '❌'} "
                    f"{h['action']['type']} -> {h['result']}"
                )
                for i, h in enumerate(history or [])
            ]
            user_content = (
                f"Snapshot:\n{snapshot}\n\nHistory:\n"
                + "\n".join(hist_lines)
                + f"\n\nTask: {prompt}"
            )

        # Run ChatAgent
        chat_agent = self._get_chat_agent()
        response = chat_agent.step(user_content)
        content = response.msgs[0].content if response.msgs else "{}"

        # Safely parse JSON response
        return self._safe_parse_json(content)

    async def process_command(self, prompt: str, max_steps: int = 15):
        r"""Process a natural language command using LLM-guided browser
        automation.

        Iteratively captures page snapshots, calls the LLM for the next
        action, executes it, and repeats until the task is complete or
        ``max_steps`` is reached.

        Args:
            prompt (str): The natural language task to execute in the
                browser.
            max_steps (int, optional): Maximum number of action steps
                before stopping. (default: :obj:`15`)
        """
        # initial full snapshot
        full_snapshot = await self._session.get_snapshot()
        assert self._session.snapshot is not None
        meta = self._session.snapshot.last_info
        logger.info("Initial snapshot priorities=%s", meta["priorities"])
        logger.debug("Full snapshot:\n%s", full_snapshot)

        plan_resp = self._llm_call(
            prompt, full_snapshot or "", is_initial=True
        )
        plan = plan_resp.get("plan", [])
        action = plan_resp.get("action")

        logger.info("Plan generated: %s", json.dumps(plan, ensure_ascii=False))

        steps = 0
        while action and steps < max_steps:
            if action.get("type") == "finish":
                logger.info("Task finished: %s", action.get("summary", "Done"))
                break

            result = await self._run_action(action)
            logger.debug("Executed action: %s | Result: %s", action, result)

            success = False
            result_for_history = ""

            if isinstance(result, str):
                success = "Error" not in result
                result_for_history = result
            elif isinstance(result, dict):
                success = result.get('success', False)
                result_for_history = result.get('message', str(result))
            else:
                # Fallback case
                success = False
                result_for_history = str(result)

            self.action_history.append(
                {
                    "action": action,
                    "result": result_for_history,
                    "success": success,
                }
            )

            diff_snapshot = await self._session.get_snapshot(
                force_refresh=ActionExecutor.should_update_snapshot(action),
                diff_only=True,
            )
            assert self._session.snapshot is not None
            meta = self._session.snapshot.last_info
            logger.debug(
                "Snapshot after action (diff=%s):\n%s",
                meta["is_diff"],
                diff_snapshot,
            )

            # Update full snapshot if page changed
            if meta["is_diff"] and not diff_snapshot.startswith(
                "- Page Snapshot (no structural changes)"
            ):
                assert self._session.snapshot is not None
                full_snapshot = self._session.snapshot.snapshot_data or ""

            action = self._llm_call(
                prompt,
                full_snapshot or "",
                is_initial=False,
                history=self.action_history,
            ).get("action")
            steps += 1

        logger.info("Process completed with %d steps", steps)

    async def _run_action(
        self, action: Dict[str, Any]
    ) -> Union[str, Dict[str, Any]]:
        r"""Execute a single browser action and return the result.

        Args:
            action (Dict[str, Any]): The action dictionary to execute.
                Navigate actions are handled directly; all others are
                delegated to the session executor.

        Returns:
            Union[str, Dict[str, Any]]: The action result as a string
                for navigate actions, or a dictionary for all other
                action types.
        """
        if action.get("type") == "navigate":
            return await self.navigate(action.get("url", ""))
        return await self._session.exec_action(action)

    async def close(self):
        r"""Clean up the browser session and release all resources."""
        await self._session.close()
