# https://github.com/ultrafunkamsterdam/undetected-chromedriver/issues/1954
# Fix for nodriver in .venv/lib/python3.11/site-packages/nodriver/core/browser.py
from __future__ import annotations

import logging
import os
from pathlib import Path

pyenv_path = os.getenv("PYENV_PATH")
if pyenv_path is None:
    env_path = os.getenv("VIRTUAL_ENV")
    if env_path is None:
        env_path = Path(os.__file__).parent.parent.parent.as_posix()
    pyenv_path = Path(env_path + "/lib/python3.12")
nodriver_pkg_path = Path(pyenv_path) / "site-packages/nodriver"
nodriver_path = nodriver_pkg_path / "cdp/network.py"
if not nodriver_path.exists():
    msg = f"{nodriver_path} not found"
    raise FileNotFoundError(msg)

new_cookie_partition_key = """\
        if isinstance(json, str):
            return cls(top_level_site=json, has_cross_site_ancestor=False)
        elif isinstance(json, dict):
            return cls(
                top_level_site=str(json["topLevelSite"]),
                has_cross_site_ancestor=bool(json["hasCrossSiteAncestor"]),
            )
"""

logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.info(f"Fixing nodriver in {nodriver_path}")
# delete CookiePartitionKey declaration
with nodriver_path.open("r+") as f:
    lines = f.readlines()
    found_def = False
    found_body = False
    i = -1
    while i < len(lines):
        i += 1
        line = lines[i]
        strip_line = line.strip("\n")
        if not found_def and line.startswith("class CookiePartitionKey:"):
            logger.info(f"Found line {i}: {strip_line}")
            found_def = True
            continue
        if found_def:
            if line.startswith("    def from_json"):
                logger.info(f"Found line {i}: {strip_line}")
                found_body = True
                continue
            if found_body:
                if line.startswith(("\t\t", "        ")):
                    logger.info(f"Removing line {i}: {strip_line}")
                    lines.pop(i)
                    i -= 1
                    continue
                else:
                    lines = lines[:i] + [new_cookie_partition_key] + lines[i:]
                    break


with nodriver_path.open("w") as f:
    f.writelines(lines)


browser_path = nodriver_pkg_path / "core/browser.py"
if not browser_path.exists():
    msg = f"{browser_path} not found"
    raise FileNotFoundError(msg)

with browser_path.open("r") as f:
    browser_path_txt = f.read()

browser_path_txt2 = browser_path_txt.replace("for _ in range(5):", "for _ in range(20):")

if browser_path_txt2 != browser_path_txt:
    logger.info("increasing browser open timeout")
    with browser_path.open("w") as f:
        f.write(browser_path_txt2)