import asyncio
import os
import json
import logging
import sys
from typing import Any, Optional
from dotenv import load_dotenv
from pathlib import Path

# Ensure we can import nanobot if installed in venv
from playwright.async_api import async_playwright, BrowserContext, Page, Browser

try:
    from nanobot.agent.loop import AgentLoop
    from nanobot.agent.tools.base import Tool
    from nanobot.bus.queue import MessageBus
    from nanobot.providers.litellm_provider import LiteLLMProvider
    from nanobot.config.schema import ExecToolConfig
except ImportError:
    print("Error: 'nanobot' not found. Please install 'nanobot-ai'.")
    # We don't exit here so it can be imported safely, but it will likely fail on run
    pass

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("nps_nanobot")

class PlaywrightTool(Tool):
    """
    A tool to control a web browser using Playwright.
    Supports navigation, clicking, typing, and extracting text.
    """
    name = "browser"
    description = (
        "Control a web browser to navigate, interact, and extract information. "
        "Use this tool to visit websites, click elements, fill forms, and read content."
    )
    parameters = {
        "type": "object",
        "properties": {
            "action": {
                "type": "string",
                "enum": ["navigate", "click", "type", "press", "scroll", "get_text", "get_links", "get_dom", "screenshot"],
                "description": "The action to perform."
            },
            "url": {"type": "string", "description": "URL to navigate to (required for 'navigate')"},
            "selector": {"type": "string", "description": "CSS selector or text to interact with (required for 'click', 'type', 'get_text')"},
            "text": {"type": "string", "description": "Text to type (required for 'type')"},
            "key": {"type": "string", "description": "Key to press (required for 'press', e.g., 'Enter')"},
            "direction": {"type": "string", "enum": ["up", "down"], "description": "Scroll direction (required for 'scroll')"},
            "amount": {"type": "integer", "description": "Scroll amount in pixels (optional for 'scroll')"},
        },
        "required": ["action"]
    }

    def __init__(self):
        self._playwright = None
        self._browser = None
        self._context: Optional[BrowserContext] = None
        self._page: Optional[Page] = None
        self._initialized = False

    async def _ensure_initialized(self):
        if not self._initialized:
            self._playwright = await async_playwright().start()
            
            # Auto-detect HF Spaces or allow HEADLESS env var override
            is_spaces = "SPACE_ID" in os.environ
            default_headless = "true" if is_spaces else "false"
            is_headless = os.getenv("HEADLESS", default_headless).lower() == "true"

            # Launch with anti-detection args
            self._browser = await self._playwright.chromium.launch(
                headless=is_headless,
                args=[
                    "--no-sandbox",
                    "--disable-blink-features=AutomationControlled",
                    "--disable-infobars"
                ],
                ignore_default_args=["--enable-automation"]
            )
            self._context = await self._browser.new_context(
                viewport={"width": 1280, "height": 720},
                user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
            )
            self._page = await self._context.new_page()
            self._page.set_default_timeout(15000) # 15s timeout
            
            # Auto-handle native dialogs
            self._page.on("dialog", lambda dialog: dialog.accept())
            
            self._initialized = True
            logger.info("Playwright browser initialized")

    async def _handle_popups(self):
        """Try to dismiss common cookie banners and popups."""
        try:
            # Common patterns for dismissal
            patterns = [
                "text=Accept", "text=Agree", "text=Allow", 
                "text=Reject", "text=Close", "text=No thanks", 
                "text=Not now", "text=OK", "text=I understand",
                "button[aria-label='Close']", ".close-button",
                "#dismiss-button", "div[aria-label='Close ad']",
                ".ri-close-line", ".close-icon"
            ]
            for selector in patterns:
                try:
                    # Check if visible and click. fast timeout.
                    elem = await self._page.query_selector(selector)
                    if elem and await elem.is_visible():
                        await elem.click(timeout=500)
                        logger.info(f"Dismissed popup using {selector}")
                        await asyncio.sleep(0.5) # Wait for animation
                except Exception:
                    pass
        except Exception as e:
            logger.warning(f"Popup dismissal error: {e}")

    async def execute(self, action: str, **kwargs: Any) -> str:
        try:
            await self._ensure_initialized()
            page = self._page
            
            if action == "navigate":
                url = kwargs.get("url")
                if not url: return "Error: 'url' is required for navigate"
                if not url.startswith("http"): url = "https://" + url
                try:
                    await page.goto(url, timeout=20000)
                    await page.wait_for_load_state("domcontentloaded", timeout=10000)
                    # Attempt to clean up view
                    await self._handle_popups()
                except Exception as e:
                    # Ignore timeout and assume page loaded partially
                    logger.warning(f"Navigation to {url} timed out: {e}")
                    await self._handle_popups() # Try finding popups even if timed out
                    return f"Navigated to {url}, but timed out waiting for full load. Content might be available."
                
                title = await page.title()
                return f"Navigated to {url}. Page title: {title}"

            elif action == "click":
                selector = kwargs.get("selector")
                if not selector: return "Error: 'selector' is required for click"
                # specialized selector handling: try text if loose selector
                try:
                    # If it looks like "link: <text>", extract
                    if selector.startswith("link: "):
                        selector = f"text={selector[6:]}"
                    await page.click(selector, timeout=5000)
                except Exception:
                    # Fallback to text match
                    await page.click(f"text={selector}")
                return f"Clicked element: {selector}"

            elif action == "type":
                selector = kwargs.get("selector")
                text = kwargs.get("text")
                if not selector or text is None: return "Error: 'selector' and 'text' are required for type"
                await page.fill(selector, text)
                return f"Typed '{text}' into {selector}"

            elif action == "press":
                key = kwargs.get("key")
                if not key: return "Error: 'key' is required for press"
                await page.keyboard.press(key)
                return f"Pressed key: {key}"

            elif action == "scroll":
                direction = kwargs.get("direction", "down")
                amount = kwargs.get("amount", 500)
                if direction == "down":
                    await page.evaluate(f"window.scrollBy(0, {amount})")
                else:
                    await page.evaluate(f"window.scrollBy(0, -{amount})")
                return f"Scrolled {direction} by {amount} pixels"

            elif action == "get_text":
                selector = kwargs.get("selector", "body")
                if selector == "body":
                    # Smart extraction: get text but also hints of links
                    # We'll stick to text for this action to keep it simple, 
                    # but maybe use Readability if we had it. For now, innerText is fine.
                    content = await page.evaluate("document.body.innerText")
                    return content[:4000] + "... (truncated)"
                else:
                    element = await page.query_selector(selector)
                    if element:
                        text = await element.inner_text()
                        return text
                    return f"Element not found: {selector}"
            
            elif action == "get_links":
                # new action to find clickable links
                links = await page.evaluate("""() => {
                    const params = [];
                    document.querySelectorAll('a').forEach(a => {
                        if (a.innerText.trim().length > 3 && a.offsetParent !== null) {
                            params.push({text: a.innerText.trim().slice(0, 50), href: a.href});
                        }
                    });
                    return params.slice(0, 20); // Top 20 relevant links
                }""")
                return json.dumps(links, indent=2)

            elif action == "get_dom":
                # Simplified interactive DOM tree for buttons and inputs
                dom = await page.evaluate("""() => {
                    const params = [];
                    const elements = document.querySelectorAll('button, input, select, textarea, [role="button"]');
                    elements.forEach(el => {
                        if (el.offsetParent !== null) { // visible
                            let label = el.innerText || el.placeholder || el.getAttribute('aria-label') || el.name || el.id || el.value;
                            label = label ? label.trim().slice(0, 50) : '[No Label]';
                            const tag = el.tagName.toLowerCase();
                            params.push(`${tag}: ${label}`);
                        }
                    });
                    return params.slice(0, 50).join('\\n'); // Top 50 interactive elements
                }""")
                return dom if dom else "No interactive elements found."

            elif action == "screenshot":
                path = "screenshot.png"
                await page.screenshot(path=path)
                return f"Screenshot saved to {path}"

            else:
                return f"Unknown action: {action}"

        except Exception as e:
            logger.error(f"Playwright error: {e}")
            return f"Error executing {action}: {str(e)}"

    async def close(self):
        if self._browser:
            await self._browser.close()
        if self._playwright:
            await self._playwright.stop()


async def run_nps_task(task: str) -> str:
    """
    Executes a task using the NPS Nanobot Agent.
    
    Args:
        task: The natural language task description.
        
    Returns:
        The result of the task execution.
    """
    print("\n" + "="*60)
    print("🐈  NPS OPERATIONAL AGENT (Nanobot + Playwright + Ollama)")
    print("="*60)

    # Initialize bus
    bus = MessageBus()
    
    # Configure LiteLLM for Mistral API (same key as the main app)
    model_name = "mistral/mistral-small-latest"
    
    mistral_api_key = os.getenv("MISTRAL_API_KEY")
    if not mistral_api_key:
        logger.warning("MISTRAL_API_KEY not found in environment variables. Agent may fail.")

    # Initialize LiteLLM provider for Mistral
    class MistralProvider(LiteLLMProvider):
        async def chat(self, messages, tools=None, model=None, max_tokens=4096, temperature=0.7):
            model = model or self.default_model
            
            kwargs = {
                "model": model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "api_key": os.getenv("MISTRAL_API_KEY", ""),
            }
            if tools:
                kwargs["tools"] = tools
                kwargs["tool_choice"] = "auto"
                
            max_retries = 5
            for attempt in range(max_retries):
                try:
                    from litellm import acompletion
                    from nanobot.providers.base import LLMResponse, ToolCallRequest
                    import json
                    
                    response = await acompletion(**kwargs)
                    llm_response = self._parse_response(response)
                    
                    # Fallback: if no tool calls but content looks like JSON tool call, parse it
                    if not llm_response.tool_calls and llm_response.content:
                        import re
                        content = llm_response.content.strip()
                        json_match = re.search(r'(\[.*\]|\{.*\})', content, re.DOTALL)
                        
                        if json_match:
                            try:
                                json_str = json_match.group(1)
                                try:
                                    data = json.loads(json_str)
                                except json.JSONDecodeError:
                                    data = []
                                    for line in json_str.splitlines():
                                        line = line.strip()
                                        if line:
                                            try:
                                                obj = json.loads(line)
                                                data.append(obj)
                                            except json.JSONDecodeError:
                                                pass
                                    if not data: raise
                                
                                tool_calls = []
                                if isinstance(data, list):
                                    for item in data:
                                        if isinstance(item, dict) and "name" in item and "arguments" in item:
                                            args = item["arguments"]
                                            if isinstance(args, str):
                                                try: args = json.loads(args)
                                                except: pass
                                            import uuid
                                            tool_calls.append(ToolCallRequest(
                                                id=f"call_{uuid.uuid4().hex[:8]}",
                                                name=item["name"],
                                                arguments=args
                                            ))
                                        elif isinstance(item, dict) and "type" in item and item["type"] == "function":
                                            func = item["function"]
                                            args = func["arguments"]
                                            if isinstance(args, str):
                                                try: args = json.loads(args)
                                                except: pass
                                            import uuid
                                            tool_calls.append(ToolCallRequest(
                                                id=item.get("id", f"call_{uuid.uuid4().hex[:8]}"),
                                                name=func["name"],
                                                arguments=args
                                            ))
                                elif isinstance(data, dict):
                                    if "name" in data and "arguments" in data:
                                        args = data["arguments"]
                                        if isinstance(args, str):
                                            try: args = json.loads(args)
                                            except: pass
                                        import uuid
                                        tool_calls.append(ToolCallRequest(
                                            id=f"call_{uuid.uuid4().hex[:8]}",
                                            name=data["name"],
                                            arguments=args
                                        ))

                                if tool_calls:
                                    llm_response.tool_calls = tool_calls
                                    llm_response.content = None
                                    llm_response.finish_reason = "tool_calls"
                                    
                            except Exception as e:
                                logger.warning(f"Failed to parse JSON tool call fallback: {e}")
                    
                    return llm_response

                except Exception as e:
                    if attempt == max_retries - 1:
                        from nanobot.providers.base import LLMResponse
                        return LLMResponse(content=f"Error calling LLM after {max_retries} attempts: {str(e)}", finish_reason="error")
                    logger.warning(f"LLM call failed (attempt {attempt+1}/{max_retries}): {e}")
            
            return LLMResponse(content="Error: Max retries exceeded", finish_reason="error")

    provider = MistralProvider(
        default_model=model_name,
        api_key=mistral_api_key
    )

    # Initialize Agent Loop
    # We use current directory as workspace
    workspace_path = Path(os.getcwd())

    # Clear memory to ensure fresh start
    memory_path = workspace_path / "memory"
    if memory_path.exists():
        import shutil
        try:
            shutil.rmtree(memory_path)
            print("INFO: Cleared previous agent memory.")
        except Exception as e:
            print(f"WARN: Could not clear memory: {e}")

    agent = AgentLoop(
        bus=bus,
        provider=provider,
        workspace=workspace_path,
        model=model_name,
        max_iterations=15,
    )

    # Register our custom Playwright tool
    playwright_tool = PlaywrightTool()
    agent.tools.register(playwright_tool)

    print(f"\nINFO: Agent initialized with model: {model_name}")
    print("INFO: Browser tool (Playwright) registered")

    try:
        if not task:
            return "❌ No task provided."

        print(f"\n🚀 Starting task: {task}")
        
        # Append strong instruction to the task
        system_instruction = (
             "\n\nIMPORTANT INSTRUCTIONS:"
             "\n1. You are an autonomous browser agent. You MUST USE TOOLS to perform actions."
             "\n2. DO NOT narrate what you are going to do (e.g., 'I will click...'). JUST CALL THE TOOL."
             "\n3. If the task is not complete, your response MUST contain a tool call."
             "\n4. Only output text when you have fully completed the user's request and have the answer."
             "\n5. Start by navigating: browser(action='navigate', ...)."
             "\n6. Use 'get_links' to find clickable items."
             "\n7. Use 'get_dom' to find buttons and inputs on interactive pages."
             "\n8. If you see popups, 'OK' dialogs, or cookie banners that block you, click them to dismiss."
        )
        full_task = task + system_instruction
        
        # Use a unique session ID for each task execution to prevent history leakage
        import uuid
        session_id = f"task_{uuid.uuid4().hex[:8]}"
        
        # Run the agent with a unique session key
        response = await agent.process_direct(
            full_task,
            session_key=f"cli:{session_id}",
            channel="cli",
            chat_id=session_id
        )
        
        print("\n" + "="*60)
        print("✅ TASK COMPLETED")
        print("="*60)
        return response

    except KeyboardInterrupt:
        return "Task interrupted by user."
    finally:
        # Cleanup
        await playwright_tool.close()
        agent.stop()

if __name__ == "__main__":
    # Test block
    task = input("Enter test task: ")
    asyncio.run(run_nps_task(task))