import asyncio
import json
import os
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime, timezone

from playwright.async_api import async_playwright

from src.connectors.base import BaseConnector
from src.utils.logging import get_logger

logger = get_logger(__name__)

COOKIE_DIR = Path(".penbot/sessions")


class PlaywrightConnector(BaseConnector):
    """
    Universal browser-automation connector for web-based chatbots.

    Supports:
    - Cookie/session persistence across runs
    - Form-based login (username/password)
    - Manual auth (pause for human login, then persist)
    - Cookie injection from a pre-exported JSON file
    - Iframe-aware element location
    - Streaming response detection via text-stability polling
    """

    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.url = self.connection_config.get("endpoint")
        self.selectors = config.get("selectors", {})
        self.browser_config = config.get("browser", {})
        self.auth_config = self.connection_config.get("auth", {})

        # State
        self.playwright = None
        self.browser = None
        self.context = None
        self.page = None
        self.last_message_count = 0

        # Cookie persistence
        self._cookie_path = self._resolve_cookie_path()

    # ------------------------------------------------------------------
    # Cookie persistence helpers
    # ------------------------------------------------------------------

    def _resolve_cookie_path(self) -> Path:
        """Determine where to store cookies for this target."""
        explicit = self.auth_config.get("cookie_storage_path")
        if explicit:
            return Path(explicit)
        safe_name = (self.name or "default").lower().replace(" ", "_")
        return COOKIE_DIR / f"{safe_name}_cookies.json"

    async def _save_cookies(self) -> None:
        """Persist browser context cookies to disk."""
        if not self.context:
            return
        try:
            cookies = await self.context.cookies()
            self._cookie_path.parent.mkdir(parents=True, exist_ok=True)
            self._cookie_path.write_text(
                json.dumps(cookies, indent=2, default=str), encoding="utf-8"
            )
            logger.info(
                "cookies_saved",
                path=str(self._cookie_path),
                count=len(cookies),
            )
        except Exception as exc:
            logger.warning("cookies_save_failed", error=str(exc))

    async def _load_cookies(self) -> bool:
        """Restore cookies from disk into the browser context.

        Returns True if cookies were loaded successfully.
        """
        if not self._cookie_path.exists():
            return False
        try:
            raw = self._cookie_path.read_text(encoding="utf-8")
            cookies = json.loads(raw)
            if not cookies:
                return False
            await self.context.add_cookies(cookies)
            logger.info(
                "cookies_loaded",
                path=str(self._cookie_path),
                count=len(cookies),
            )
            return True
        except Exception as exc:
            logger.warning("cookies_load_failed", error=str(exc))
            return False

    # ------------------------------------------------------------------
    # Authentication flows
    # ------------------------------------------------------------------

    async def _perform_auth(self) -> bool:
        """Run the configured authentication flow.

        Supported auth types (via target.connection.auth.type):
          - "none"        : skip auth entirely (default)
          - "form_login"  : fill username/password on a login page
          - "manual"      : open visible browser, wait for human to log in
          - "cookie_file" : inject cookies from a pre-exported JSON file

        Returns True if auth succeeded (or was skipped).
        """
        auth_type = self.auth_config.get("type", "none")
        if auth_type == "none":
            return True

        if auth_type == "form_login":
            return await self._auth_form_login()
        if auth_type == "manual":
            return await self._auth_manual()
        if auth_type == "cookie_file":
            return await self._auth_cookie_file()

        logger.warning("unknown_auth_type", auth_type=auth_type)
        return True

    async def _auth_form_login(self) -> bool:
        """Authenticate by filling a login form."""
        login_url = self.auth_config.get("login_url", self.url)
        creds = self.auth_config.get("credentials", {})

        username = os.getenv(creds.get("username_env", ""), "") or creds.get("username", "")
        password = os.getenv(creds.get("password_env", ""), "") or creds.get("password", "")

        if not username or not password:
            logger.error("auth_credentials_missing")
            return False

        auth_selectors = self.auth_config.get("selectors", {})
        user_sel = auth_selectors.get("username_field", "input[name='username']")
        pass_sel = auth_selectors.get("password_field", "input[name='password']")
        submit_sel = auth_selectors.get("submit_button", "button[type='submit']")
        success_sel = auth_selectors.get("success_indicator")

        try:
            logger.info("auth_form_login_start", login_url=login_url)
            await self.page.goto(login_url, wait_until="domcontentloaded")
            await self.page.wait_for_selector(user_sel, timeout=15000)

            await self.page.fill(user_sel, username)
            await self.page.fill(pass_sel, password)
            await self.page.click(submit_sel)

            if success_sel:
                await self.page.wait_for_selector(success_sel, timeout=30000)
            else:
                await self.page.wait_for_load_state("networkidle", timeout=30000)

            await self._save_cookies()
            logger.info("auth_form_login_success")
            return True

        except Exception as exc:
            logger.error("auth_form_login_failed", error=str(exc))
            return False

    async def _auth_manual(self) -> bool:
        """Pause execution so a human can log in manually.

        The browser must be launched in headed mode (headless=false).
        After the user presses Enter in the terminal, cookies are saved.
        """
        login_url = self.auth_config.get("login_url", self.url)
        try:
            logger.info("auth_manual_start", login_url=login_url)
            await self.page.goto(login_url, wait_until="domcontentloaded")

            print(
                "\n  [PenBot] Browser is open. Please log in manually.\n"
                "           Press ENTER here when you're done...\n"
            )
            await asyncio.get_event_loop().run_in_executor(None, input)

            await self._save_cookies()
            logger.info("auth_manual_success")
            return True
        except Exception as exc:
            logger.error("auth_manual_failed", error=str(exc))
            return False

    async def _auth_cookie_file(self) -> bool:
        """Load cookies from an explicit file path (e.g. exported from a browser extension)."""
        cookie_file = self.auth_config.get("cookie_file")
        if not cookie_file or not Path(cookie_file).exists():
            logger.error("auth_cookie_file_missing", path=cookie_file)
            return False
        try:
            raw = Path(cookie_file).read_text(encoding="utf-8")
            cookies = json.loads(raw)
            await self.context.add_cookies(cookies)
            logger.info("auth_cookie_file_loaded", count=len(cookies))
            return True
        except Exception as exc:
            logger.error("auth_cookie_file_failed", error=str(exc))
            return False

    async def _is_session_valid(self) -> bool:
        """After restoring cookies, check if the session is still alive.

        Uses `auth.selectors.session_check_selector` — a CSS selector that
        should be visible only when the user is logged in.  Falls back to
        checking that the chat input field is present.
        """
        check_sel = self.auth_config.get("selectors", {}).get("session_check_selector")
        fallback_sel = self.selectors.get("input_field")
        sel = check_sel or fallback_sel

        if not sel:
            return True  # No way to check — assume valid

        try:
            await self.page.goto(self.url, wait_until="domcontentloaded")
            await self.page.wait_for_selector(sel, timeout=10000)
            logger.info("session_still_valid")
            return True
        except Exception:
            logger.info("session_expired_or_invalid")
            return False

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self) -> bool:
        """Launch the browser, restore session or authenticate, navigate to target."""
        try:
            logger.info("playwright_initializing", target=self.name, url=self.url)
            self.playwright = await async_playwright().start()

            headless = self.browser_config.get("headless", True)
            launch_args = {}
            slow_mo = self.browser_config.get("slow_mo")
            if slow_mo:
                launch_args["slow_mo"] = slow_mo

            self.browser = await self.playwright.chromium.launch(headless=headless, **launch_args)

            viewport = self.browser_config.get("viewport")
            ctx_kwargs: Dict[str, Any] = {}
            if viewport:
                ctx_kwargs["viewport"] = viewport

            self.context = await self.browser.new_context(**ctx_kwargs)
            self.page = await self.context.new_page()

            # --- Session restoration flow ---
            # 1. Try to restore cookies from a previous run
            cookies_restored = await self._load_cookies()

            if cookies_restored:
                # 2. Navigate and check if the session is still valid
                if await self._is_session_valid():
                    logger.info("session_restored_skipping_auth")
                else:
                    # Session expired — clear stale cookies and re-auth
                    await self.context.clear_cookies()
                    if not await self._perform_auth():
                        logger.error("auth_failed_after_session_expired")
                        return False
            else:
                # No saved cookies — run the auth flow
                if not await self._perform_auth():
                    logger.error("auth_failed_no_saved_session")
                    return False

            # --- Navigate to target URL (if auth redirected elsewhere) ---
            if self.page.url != self.url:
                await self.page.goto(self.url, wait_until="domcontentloaded")

            # --- Post-navigation setup ---
            await self._dismiss_cookie_banner()
            await self._open_chat_widget()
            await self._wait_for_input_ready()

            self.last_message_count = await self._count_messages()

            # Persist cookies after successful init (captures any new tokens)
            await self._save_cookies()

            logger.info("playwright_initialized", target=self.name)
            return True

        except Exception as e:
            logger.error("playwright_init_failed", error=str(e))
            return False

    async def _dismiss_cookie_banner(self) -> None:
        cookie_selector = self.selectors.get("cookie_button")
        if not cookie_selector:
            return
        try:
            cookie_btn = self.page.locator(cookie_selector).first
            if await cookie_btn.is_visible(timeout=5000):
                await cookie_btn.click()
                logger.info("playwright_cookies_accepted")
                await asyncio.sleep(1)
        except Exception:
            pass

    async def _open_chat_widget(self) -> None:
        open_selector = self.selectors.get("open_chat_button")
        if not open_selector:
            return
        try:
            open_btn = await self._get_locator(open_selector)
            await open_btn.click()
            await asyncio.sleep(2)
        except Exception as e:
            logger.warning("playwright_failed_to_open_chat", error=str(e))

    async def _wait_for_input_ready(self) -> None:
        input_selector = self.selectors.get("input_field")
        if not input_selector:
            return
        try:
            await self.page.wait_for_selector(input_selector, timeout=15000)
        except Exception:
            logger.warning("playwright_input_not_found_immediately", selector=input_selector)

    # ------------------------------------------------------------------
    # Element location (iframe-aware)
    # ------------------------------------------------------------------

    async def _get_locator(self, selector: str, timeout: int = 10000):
        """Find an element checking both the main page and iframes."""
        start_time = datetime.now()

        while (datetime.now() - start_time).total_seconds() * 1000 < timeout:
            loc = self.page.locator(selector)
            if await loc.count() > 0 and await loc.first.is_visible():
                return loc.first

            for frame in self.page.frames:
                try:
                    loc = frame.locator(selector)
                    if await loc.count() > 0 and await loc.first.is_visible():
                        logger.info(
                            "playwright_found_in_frame",
                            selector=selector,
                            frame_url=frame.url,
                        )
                        return loc.first
                except Exception:
                    continue

            await asyncio.sleep(0.5)

        logger.warning("playwright_selector_not_found_anywhere", selector=selector)
        return self.page.locator(selector)

    # ------------------------------------------------------------------
    # Messaging
    # ------------------------------------------------------------------

    async def send_message(
        self,
        message: str,
        context: Optional[Dict[str, Any]] = None,
        image_data: Optional[str] = None,
        image_mime_type: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Type a message into the browser chat and wait for a reply."""
        if not self.page:
            await self.initialize()

        if image_data:
            logger.warning("playwright_image_not_supported_yet")

        input_selector = self.selectors.get("input_field")
        submit_selector = self.selectors.get("submit_button")

        try:
            input_el = await self._get_locator(input_selector)
            await input_el.fill(message)

            if submit_selector:
                submit_el = await self._get_locator(submit_selector)
                await submit_el.click()
            else:
                await input_el.press("Enter")

            timeout = self.config.get("timeout", 30000)
            new_reply = await self._wait_for_new_message(timeout)

            # Persist any session cookies that the target may have refreshed
            await self._save_cookies()

            return {
                "content": new_reply,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "metadata": {"source": "playwright"},
            }

        except Exception as e:
            logger.error("playwright_interaction_error", error=str(e))
            return {
                "content": "[Error: Browser interaction failed]",
                "metadata": {"error": str(e)},
            }

    async def receive_message(self) -> Dict[str, Any]:
        """Not used in request-response flow, handled inside send_message."""

    async def health_check(self) -> bool:
        """Check if the page is open and the input is interactive."""
        try:
            if not self.page:
                return False
            input_selector = self.selectors.get("input_field", "body")
            el = await self._get_locator(input_selector)
            return await el.is_visible()
        except Exception:
            return False

    async def close(self) -> None:
        """Save cookies and close the browser."""
        await self._save_cookies()
        if self.context:
            await self.context.close()
        if self.browser:
            await self.browser.close()
        if self.playwright:
            await self.playwright.stop()
        logger.info("playwright_closed")

    # ------------------------------------------------------------------
    # File upload
    # ------------------------------------------------------------------

    async def upload_file(
        self,
        file_content: bytes,
        filename: str,
        mime_type: str = "application/octet-stream",
        context: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Upload a file via browser automation (for RAG testing)."""
        upload_selector = self.selectors.get("file_upload_input")
        if not upload_selector:
            raise NotImplementedError(
                "File upload not configured. Add 'file_upload_input' selector to config."
            )

        try:
            import tempfile
            import os as _os

            with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=f"_{filename}") as tmp:
                tmp.write(file_content)
                tmp_path = tmp.name

            logger.info("file_upload_starting", filename=filename, size=len(file_content))
            file_input = await self._get_locator(upload_selector, timeout=10000)
            await file_input.set_input_files(tmp_path)
            await asyncio.sleep(2)

            upload_submit = self.selectors.get("file_upload_submit")
            if upload_submit:
                submit_btn = await self._get_locator(upload_submit, timeout=5000)
                await submit_btn.click()
                await asyncio.sleep(2)

            try:
                _os.unlink(tmp_path)
            except Exception as exc:
                logger.debug("temp_file_cleanup_failed", path=tmp_path, error=str(exc))

            logger.info("file_uploaded", filename=filename)
            return {
                "success": True,
                "file_id": None,
                "message": "File uploaded successfully via browser automation",
                "metadata": {
                    "filename": filename,
                    "size": len(file_content),
                    "mime_type": mime_type,
                    "method": "playwright",
                },
            }

        except Exception as e:
            logger.error("file_upload_failed", filename=filename, error=str(e))
            return {
                "success": False,
                "message": f"Upload failed: {str(e)}",
                "metadata": {"error": str(e)},
            }

    # ------------------------------------------------------------------
    # Response detection helpers
    # ------------------------------------------------------------------

    async def _count_messages(self) -> int:
        """Count current message bubbles (checks iframes too)."""
        selector = self.selectors.get("response_container")
        if not selector:
            return 0

        try:
            for frame in self.page.frames:
                count = await frame.locator(selector).count()
                if count > 0:
                    return count
            return await self.page.locator(selector).count()
        except Exception:
            return 0

    async def _wait_for_new_message(self, timeout_ms: int) -> str:
        """Wait for a new message bubble and return its text once streaming stabilises."""
        selector = self.selectors.get("response_container")
        if not selector:
            return ""

        start_time = datetime.now()

        # Identify the frame containing the chat
        active_frame = self.page
        for frame in self.page.frames:
            if await frame.locator(self.selectors.get("input_field")).count() > 0:
                active_frame = frame
                break

        target_element = None
        while (datetime.now() - start_time).total_seconds() * 1000 < timeout_ms:
            try:
                current_count = await active_frame.locator(selector).count()
                if current_count > self.last_message_count:
                    target_element = active_frame.locator(selector).last
                    self.last_message_count = current_count
                    break
            except Exception as e:
                logger.debug("message_count_poll_error", error=str(e))
            await asyncio.sleep(0.5)

        if not target_element:
            return "[Timeout: No reply received]"

        # Text stability check — wait until streaming finishes
        last_text = ""
        stable_start = datetime.now()
        stream_timeout = 30000
        stream_start = datetime.now()

        while (datetime.now() - stream_start).total_seconds() * 1000 < stream_timeout:
            try:
                current_text = await target_element.inner_text()
                if current_text == last_text and current_text.strip():
                    if (datetime.now() - stable_start).total_seconds() > 1.5:
                        return current_text
                else:
                    last_text = current_text
                    stable_start = datetime.now()
                await asyncio.sleep(0.5)
            except Exception:
                break

        return last_text
