"""CPA Chromium session lifecycle, reuse, cleanup, and cookies."""
from __future__ import annotations

import os
import threading
import time
from pathlib import Path
from typing import Any, Callable, Optional

LogFn = Callable[[str], None]
_mint_tls = threading.local()
_mint_registry_lock = threading.Lock()
_mint_registry = set()

def _noop_log(_: str) -> None:
    return None

class BrowserConfirmError(RuntimeError):
    pass

def _sleep(sec: float) -> None:
    time.sleep(sec)

def create_standalone_page(proxy: Optional[str] = None, headless: bool = False, log: Optional[LogFn] = None):
    logger = log or _noop_log
    try:
        from DrissionPage import Chromium, ChromiumOptions
    except ImportError as exc:
        raise BrowserConfirmError("DrissionPage not installed") from exc

    options = None
    package_root = Path(__file__).resolve().parents[1]
    try:
        from browser_runtime import create_browser_options

        options = create_browser_options(
            extension_path=package_root / "turnstilePatch"
        )
        logger("using shared browser_runtime.create_browser_options")
    except Exception as exc:  # noqa: BLE001
        logger("shared browser options unavailable: %s" % exc)
        options = None

    if options is None:
        options = ChromiumOptions()
        options.auto_port()
        options.set_timeouts(base=2)
        for flag in (
            "--disable-gpu",
            "--no-sandbox",
            "--disable-dev-shm-usage",
            "--mute-audio",
            "--no-first-run",
            "--disable-background-networking",
            "--window-size=1280,900",
        ):
            options.set_argument(flag)
        extension = str(package_root / "turnstilePatch")
        if os.path.isdir(extension):
            try:
                options.add_extension(extension)
                logger("added extension %s" % extension)
            except Exception as exc:  # noqa: BLE001
                logger("extension add failed: %s" % exc)

    if headless:
        try:
            options.headless(True)
        except Exception:
            options.set_argument("--headless=new")
        logger("headless=True (may hit Cloudflare / break real clicks)")
    else:
        try:
            options.headless(False)
        except Exception:
            pass

    for candidate in (
        "/usr/bin/chromium",
        "/usr/bin/chromium-browser",
        "/usr/bin/google-chrome",
        "/usr/bin/google-chrome-stable",
    ):
        if os.path.isfile(candidate):
            try:
                options.set_browser_path(candidate)
            except Exception:
                pass
            break

    from .proxyutil import prepare_chromium_proxy, proxy_log_label, resolve_proxy

    resolved = resolve_proxy(proxy)
    proxy_bridge = None
    browser = None
    chrome_proxy, proxy_bridge = prepare_chromium_proxy(resolved, log=logger)
    try:
        if chrome_proxy:
            options.set_argument("--proxy-server=%s" % chrome_proxy)
            logger("browser proxy=%s (chromium %s)" % (proxy_log_label(resolved), chrome_proxy))
        else:
            logger("browser proxy=(none)")
        browser = Chromium(options)
        if proxy_bridge is not None:
            setattr(browser, "_cpa_proxy_bridge", proxy_bridge)
        page = browser.latest_tab
        _register_mint_browser(browser)
        logger("standalone chromium started")
        return browser, page
    except Exception:
        if browser is not None:
            close_standalone(browser)
        if proxy_bridge is not None:
            try:
                proxy_bridge.stop()
            except Exception:
                pass
        raise

def close_standalone(browser: Any) -> None:
    if browser is None:
        return
    _unregister_mint_browser(browser)
    bridge = getattr(browser, "_cpa_proxy_bridge", None)
    try:
        browser.quit()
    except Exception:
        pass
    if bridge is not None:
        try:
            bridge.stop()
        except Exception:
            pass

def _register_mint_browser(browser: Any) -> None:
    if browser is None:
        return
    with _mint_registry_lock:
        _mint_registry.add(browser)

def _unregister_mint_browser(browser: Any) -> None:
    if browser is None:
        return
    with _mint_registry_lock:
        _mint_registry.discard(browser)

def _mint_tls_get():
    state = getattr(_mint_tls, "state", None)
    if state is None:
        state = {"browser": None, "page": None, "served": 0, "proxy": None, "headless": None}
        _mint_tls.state = state
    return state

def clear_page_session(page: Any, browser: Optional[Any] = None, log: Optional[LogFn] = None) -> None:
    logger = log or _noop_log
    try:
        if page is not None:
            try:
                page.get("about:blank")
            except Exception:
                pass
            for javascript in (
                "try{localStorage.clear()}catch(e){}",
                "try{sessionStorage.clear()}catch(e){}",
            ):
                try:
                    page.run_js(javascript)
                except Exception:
                    pass
        for target in (page, browser):
            if target is None:
                continue
            try:
                target.set.cookies.clear()  # type: ignore[attr-defined]
                logger("mint session cookies cleared")
                break
            except Exception:
                try:
                    cookies = target.cookies()
                    if isinstance(cookies, list):
                        for cookie in cookies:
                            try:
                                target.set.cookies.remove(cookie)  # type: ignore[attr-defined]
                            except Exception:
                                pass
                except Exception:
                    pass
    except Exception as exc:
        logger("clear_page_session: %s" % exc)

def normalize_cookies(cookies: Any):
    output = []
    if not cookies:
        return output
    if isinstance(cookies, dict):
        for name, value in cookies.items():
            if name and value is not None:
                output.append({"name": str(name), "value": str(value), "domain": ".x.ai", "path": "/"})
        cookies = output
        output = []
    if not isinstance(cookies, (list, tuple)):
        return output
    for cookie in cookies:
        if not isinstance(cookie, dict):
            continue
        name = cookie.get("name") or cookie.get("Name")
        value = cookie.get("value") or cookie.get("Value")
        if not name or value is None:
            continue
        domain = str(cookie.get("domain") or cookie.get("Domain") or ".x.ai")
        path = str(cookie.get("path") or cookie.get("Path") or "/")
        item = {"name": str(name), "value": str(value), "domain": domain, "path": path}
        for source, target in (
            ("expiry", "expiry"),
            ("expires", "expiry"),
            ("secure", "secure"),
            ("httpOnly", "httpOnly"),
            ("sameSite", "sameSite"),
        ):
            if source in cookie and cookie[source] is not None:
                item[target] = cookie[source]
        output.append(item)
    sso_names = {"sso", "sso-rw", "cf_clearance", "sso_jwt", "__cf_bm"}
    extras = []
    seen = {(item["name"], item["domain"], item["path"]) for item in output}
    for item in list(output):
        name = item["name"]
        if name not in sso_names and not name.startswith("sso"):
            continue
        for domain in (".x.ai", "accounts.x.ai", ".accounts.x.ai", "auth.x.ai", ".auth.x.ai"):
            key = (name, domain, item["path"])
            if key in seen:
                continue
            clone = dict(item)
            clone["domain"] = domain
            extras.append(clone)
            seen.add(key)
    output.extend(extras)
    return output

def inject_cookies(page: Any, cookies: Any, log: Optional[LogFn] = None) -> int:
    logger = log or _noop_log
    items = normalize_cookies(cookies)
    if not items or page is None:
        return 0
    for url in ("https://accounts.x.ai/", "https://auth.x.ai/", "https://grok.com/"):
        try:
            page.get(url)
            _sleep(0.4)
        except Exception:
            continue
    count = 0
    for target_name, target in (("page", page), ("browser", getattr(page, "browser", None))):
        if target is None:
            continue
        try:
            target.set.cookies(items)  # type: ignore[attr-defined]
            count = len(items)
            logger("injected cookies bulk via %s=%s" % (target_name, count))
            break
        except Exception as exc:
            logger("bulk set via %s failed: %s" % (target_name, exc))
    if count == 0:
        for item in items:
            ok = False
            for target in (page, getattr(page, "browser", None)):
                if target is None:
                    continue
                try:
                    target.set.cookies(item)  # type: ignore[attr-defined]
                    ok = True
                    break
                except Exception:
                    continue
            if ok:
                count += 1
        logger("injected cookies one-by-one=%s/%s" % (count, len(items)))
    try:
        javascript_items = [
            cookie
            for cookie in items
            if not bool(cookie.get("httpOnly")) and str(cookie.get("name") or "").startswith("sso")
        ]
        if javascript_items:
            page.run_js(
                """
                const items = arguments[0] || [];
                for (const c of items) {
                  let cookie = `${c.name}=${c.value}; path=${c.path || '/'}; domain=${c.domain || '.x.ai'}`;
                  if (c.secure !== false) cookie += '; Secure';
                  document.cookie = cookie;
                }
                return document.cookie;
                """,
                javascript_items,
            )
            logger("injected non-httpOnly sso cookies via document.cookie")
    except Exception as exc:
        logger("document.cookie injection failed: %s" % exc)
    return count

def acquire_mint_browser(proxy: Optional[str] = None, headless: bool = False, reuse: bool = True, recycle_every: int = 15, log: Optional[LogFn] = None):
    logger = log or _noop_log
    state = _mint_tls_get()
    if not reuse:
        browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
        return browser, page, True
    proxy_key = str(proxy or "")
    if state["browser"] is None:
        browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
        state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
        return browser, page, False
    if state["proxy"] != proxy_key or bool(state["headless"]) != bool(headless):
        try:
            close_standalone(state["browser"])
        except Exception:
            pass
        browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
        state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
        return browser, page, False
    if recycle_every and state["served"] and state["served"] % max(int(recycle_every), 1) == 0:
        try:
            close_standalone(state["browser"])
        except Exception:
            pass
        browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
        state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
        return browser, page, False
    clear_page_session(state["page"], browser=state["browser"], log=logger)
    return state["browser"], state["page"], False

def release_mint_browser(owned: bool, success: bool, log: Optional[LogFn] = None) -> None:
    logger = log or _noop_log
    if owned:
        return
    state = _mint_tls_get()
    if success:
        state["served"] = int(state.get("served", 0) or 0) + 1
        logger("mint browser served=%s" % state["served"])

def shutdown_mint_browsers() -> None:
    state = _mint_tls_get()
    with _mint_registry_lock:
        browsers = list(_mint_registry)
    for browser in browsers:
        try:
            close_standalone(browser)
        except Exception:
            pass
    state.update({"browser": None, "page": None, "served": 0, "proxy": None, "headless": None})

