# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_navigator

import os
from typing import Any, Dict, List, Literal, Optional, Union

from coreason_identity.models import UserContext
from fastmcp import FastMCP
from pydantic import BaseModel, Field

from coreason_navigator.driver import PlaywrightNavigatorAsync
from coreason_navigator.types import (
    BrowserAction,
    ClickAction,
    GotoAction,
    ScreenshotAction,
    ScrollAction,
    TypeAction,
    WaitAction,
)
from coreason_navigator.utils.logger import logger

# Initialize FastMCP server
mcp = FastMCP("coreason-navigator")

# Singleton Navigator Instance
# This is a module-level singleton to persist state across MCP tool calls.
# In a multi-user environment, this would need to be a session-managed store.
_navigator: PlaywrightNavigatorAsync | None = None
_server_context: UserContext | None = None


def get_navigator() -> PlaywrightNavigatorAsync:
    """
    Retrieves the existing singleton navigator instance or creates a new one.

    Reads configuration (rate limit, allowed domains) from environment variables.

    Returns:
        PlaywrightNavigatorAsync: The active navigator instance.
    """
    global _navigator
    if _navigator is None:
        logger.info("Initializing new PlaywrightNavigatorAsync instance")
        # Read rate limit from environment, default to 2.0
        rate_limit = float(os.getenv("NAVIGATOR_RATE_LIMIT", "2.0"))

        # AUC-10: Read allowed domains
        allowed_domains_env = os.getenv("NAVIGATOR_ALLOWED_DOMAINS")
        allowed_domains = [d.strip() for d in allowed_domains_env.split(",")] if allowed_domains_env else None

        # We start with headless=True for production efficiency,
        # but this could be configurable via env vars if needed.
        _navigator = PlaywrightNavigatorAsync(
            headless=True,
            rate_limit=rate_limit,
            allowed_domains=allowed_domains,
        )
    return _navigator


def set_server_context(context: UserContext) -> None:
    """Sets the global server context."""
    global _server_context
    _server_context = context


def get_server_context() -> UserContext:
    """Retrieves the global server context or creates a default one."""
    global _server_context
    if _server_context is None:
        # Create a default context for MCP server usage when not initialized via main.py
        _server_context = UserContext(
            sub="mcp-server-default",
            email="mcp-server-default@example.com",
            permissions=["system"],
        )
    return _server_context


class BrowserStateResponse(BaseModel):
    """
    Standard response model for browser actions.
    Returns the visual state (screenshot, SoM) and context.
    """

    url: str
    title: str
    screenshot_base64: str = Field(..., description="Base64 encoded JPEG screenshot of the current viewport")
    som_screenshot_base64: str | None = Field(
        None, description="Screenshot with Set-of-Marks (bounding boxes + IDs) overlaid"
    )
    interactive_elements_count: int = Field(..., description="Number of interactive elements detected")


async def handle_request(request: Dict[str, Any], context: UserContext) -> str:
    """
    Handles a request dictionary and executes the corresponding action using the driver.

    Args:
        request: A dictionary containing the command and arguments.
                 Supported commands: 'navigate', 'interact', 'extract'.
        context: The user context for the operation.

    Returns:
        str: The result of the operation.
    """
    nav = get_navigator()
    command = request.get("command")

    if command == "navigate":
        url = request.get("url")
        if not url:
            return "Error: URL required for navigation"
        # Ensure browser is launched
        await nav.launch()
        state = await nav.navigate(url, context=context)
        return f"Navigated to {state.url}"

    elif command == "interact":
        action_type = request.get("action")
        browser_action: BrowserAction

        # Map request to BrowserAction ... this logic duplicates browser_interact slightly
        # but creates a direct path for CLI.
        if action_type == "click":
            target_id = request.get("target_id")
            coords = request.get("coordinate")
            x, y = None, None
            if coords and len(coords) == 2:
                x, y = float(coords[0]), float(coords[1])
            browser_action = ClickAction(target_id=target_id, x=x, y=y)
        elif action_type == "type":
            text = request.get("text")
            if not text:
                return "Error: Text required for type action"
            browser_action = TypeAction(text=text)
        elif action_type == "scroll":
            direction = request.get("direction", "down")
            amount = request.get("amount", 500)
            browser_action = ScrollAction(direction=direction, amount=amount)
        elif action_type == "wait":
            browser_action = WaitAction()
        elif action_type == "goto":
            url = request.get("url") or request.get("text")
            if not url:
                return "Error: URL required for goto"
            browser_action = GotoAction(url=url)
        elif action_type == "screenshot":
            browser_action = ScreenshotAction()
        else:
            return f"Error: Unknown action type {action_type}"

        state = await nav.perform_action(browser_action, context=context)
        return f"Interaction complete. New State: {state.url}"

    elif command == "extract":
        fmt = request.get("format", "markdown")
        content = await nav.extract_content(format=fmt, context=context)
        return content

    else:
        return f"Error: Unknown command {command}"


@mcp.tool
async def browser_open(url: str) -> str:
    """
    Opens a browser session and navigates to the specified URL.

    Args:
        url: The URL to navigate to.

    Returns:
        str: A JSON string representing the BrowserStateResponse (screenshot, SoM, context).
    """
    logger.info(f"MCP Tool: browser_open called with url={url}")
    nav = get_navigator()
    context = get_server_context()

    # Launch if not already running (driver handles idempotency)
    await nav.launch()

    state = await nav.navigate(url, context=context)

    # Convert to response model for cleaner output, but return as string
    # because MCP tools typically return text/json strings.
    response = BrowserStateResponse(
        url=state.url,
        title=state.title,
        screenshot_base64=state.screenshot_base64,
        som_screenshot_base64=state.som_screenshot_base64,
        interactive_elements_count=len(state.interactive_elements),
    )

    return str(response.model_dump_json())


@mcp.tool
async def browser_screenshot() -> str:
    """
    Captures the current state of the browser without performing any action.
    Useful for re-orienting if the agent loses context.

    Returns:
        str: A JSON string representing the BrowserStateResponse, or an error message.
    """
    logger.info("MCP Tool: browser_screenshot called")
    nav = get_navigator()
    context = get_server_context()

    # We must ensure the browser is open.
    # If not, we can't really screenshot anything.
    if not nav.page:
        return "Error: Browser session not started. Call browser_open first."

    state = await nav.perform_action(ScreenshotAction(), context=context)

    response = BrowserStateResponse(
        url=state.url,
        title=state.title,
        screenshot_base64=state.screenshot_base64,
        som_screenshot_base64=state.som_screenshot_base64,
        interactive_elements_count=len(state.interactive_elements),
    )

    return str(response.model_dump_json())


@mcp.tool
async def browser_interact(
    action: Literal["click", "type", "scroll", "wait", "goto", "screenshot"],
    target_id: Optional[str] = None,
    text: Optional[str] = None,
    coordinate: Optional[List[int]] = None,
    scroll_direction: Optional[str] = None,
    scroll_amount: Optional[int] = None,
) -> str:
    """
    Interacts with the browser instance (click, type, scroll, wait, goto).

    Args:
        action: The type of action to perform (click, type, scroll, wait, goto, screenshot).
        target_id: The ID of the element to interact with (for clicks). Corresponds to SoM label.
        text: The text to type (for type action) or URL (for goto action if needed).
        coordinate: Optional [x, y] coordinates to override target_id for clicks.
        scroll_direction: Direction to scroll ("up", "down", "left", "right"). Defaults to "down".
        scroll_amount: Amount to scroll in pixels. Defaults to 500.

    Returns:
        str: A JSON string representing the updated BrowserStateResponse, or an error message.
    """
    logger.info(f"MCP Tool: browser_interact called with action={action}")
    nav = get_navigator()
    context = get_server_context()

    if not nav.page:
        return "Error: Browser session not started. Call browser_open first."

    try:
        browser_action: Union[BrowserAction, None] = None

        if action == "click":
            x, y = None, None
            if coordinate and len(coordinate) == 2:
                x, y = float(coordinate[0]), float(coordinate[1])

            # Require either target_id or coordinates
            if not target_id and (x is None or y is None):
                return "Error: Click action requires target_id or coordinate=[x, y]."

            browser_action = ClickAction(target_id=target_id, x=x, y=y)

        elif action == "type":
            if not text:
                return "Error: Type action requires 'text' argument."
            browser_action = TypeAction(text=text)

        elif action == "scroll":
            direction: Literal["up", "down", "left", "right"] = "down"
            if scroll_direction and scroll_direction in ["up", "down", "left", "right"]:
                direction = scroll_direction  # type: ignore

            amount = scroll_amount if scroll_amount is not None else 500
            browser_action = ScrollAction(direction=direction, amount=amount)

        elif action == "wait":
            browser_action = WaitAction()

        elif action == "goto":
            if text:
                browser_action = GotoAction(url=text)
            else:
                return "Error: Goto action requires a URL (pass in 'text' field)."

        elif action == "screenshot":
            browser_action = ScreenshotAction()

        if browser_action:
            state = await nav.perform_action(browser_action, context=context)
            response = BrowserStateResponse(
                url=state.url,
                title=state.title,
                screenshot_base64=state.screenshot_base64,
                som_screenshot_base64=state.som_screenshot_base64,
                interactive_elements_count=len(state.interactive_elements),
            )
            return str(response.model_dump_json())

        return "Error: Unknown action or failed to construct action."

    except Exception as e:
        logger.exception(f"Action failed: {e}")
        return f"Error: {str(e)}"


@mcp.tool
async def browser_read(format: Literal["markdown", "html"] = "markdown") -> str:
    """
    Extracts the main content of the current page in Markdown or HTML format.
    Useful for reading long-form text without visual noise.

    Args:
        format: The desired output format ("markdown" or "html"). Defaults to "markdown".

    Returns:
        str: The extracted content as a string, or an error message.
    """
    logger.info(f"MCP Tool: browser_read called with format={format}")
    nav = get_navigator()
    context = get_server_context()

    if not nav.page:
        return "Error: Browser session not started. Call browser_open first."

    try:
        content = await nav.extract_content(format=format, context=context)
        return content
    except Exception as e:
        logger.exception(f"Content extraction failed: {e}")
        return f"Error: {str(e)}"
