from __future__ import annotations

import logging
import os
import shutil
import subprocess
import time
import winreg
from dataclasses import dataclass
from pathlib import Path

from .browser_profile import FirefoxHardening, ensure_firefox_profile, persona_env
from .paths import app_data_dir
from .signup_prep import SignupDraft, install_signup_extension

log = logging.getLogger(__name__)

# CREATE_NO_WINDOW suppresses the console; DETACHED_PROCESS fully separates
# the child from Python's job object so Firefox lives past Python's exit.
_CREATE_FLAGS = (
    getattr(subprocess, "CREATE_NO_WINDOW", 0x08000000)
    | getattr(subprocess, "DETACHED_PROCESS", 0x00000008)
    | getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0x00000200)
)


@dataclass
class BrowserConfig:
    firefox_path: str = ""
    profile_dir: str = ""
    clear_on_close: bool = True
    disposable_profile: bool = False
    kill_on_chain_drop: bool = True
    auto_relaunch: bool = False
    lock_managed_profile: bool = True
    force_proxy: bool = True
    disable_webrtc: bool = True
    resist_fingerprinting: bool = True
    disable_telemetry: bool = True
    first_party_isolation: bool = True
    strict_tracking_protection: bool = True
    timezone_utc: bool = True
    persona: str = "blend_windows_chrome"
    cookie_mode: str = "block_third_party"


def _find_firefox_via_registry() -> str:
    """Look up Firefox install path in the Windows registry."""
    hives = [
        (winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Mozilla\Mozilla Firefox"),
        (winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\WOW6432Node\Mozilla\Mozilla Firefox"),
        (winreg.HKEY_CURRENT_USER,  r"SOFTWARE\Mozilla\Mozilla Firefox"),
    ]
    for hive, key_path in hives:
        try:
            with winreg.OpenKey(hive, key_path) as root:
                version = winreg.QueryValue(root, None)
                for sub in (f"{version}\\Main", "bin"):
                    try:
                        with winreg.OpenKey(root, sub) as k:
                            path = winreg.QueryValueEx(k, "PathToExe")[0]
                            if path and Path(path).is_file():
                                return str(path)
                    except OSError:
                        continue
        except OSError:
            continue

    # Also try the "Uninstall" key
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        for root_path in (
            r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall",
            r"SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Uninstall",
        ):
            try:
                with winreg.OpenKey(hive, root_path) as base:
                    i = 0
                    while True:
                        try:
                            sub = winreg.EnumKey(base, i)
                            i += 1
                            if "firefox" not in sub.lower():
                                continue
                            with winreg.OpenKey(base, sub) as k:
                                loc = winreg.QueryValueEx(k, "InstallLocation")[0]
                                candidate = Path(loc) / "firefox.exe"
                                if candidate.is_file():
                                    return str(candidate)
                        except OSError:
                            break
            except OSError:
                continue
    return ""


def default_firefox_path() -> str:
    """Find Firefox.exe — registry first, then well-known paths."""
    # 1) Registry
    reg = _find_firefox_via_registry()
    if reg:
        return reg

    # 2) Standard install locations
    candidates: list[Path] = [
        Path(r"C:\Program Files\Mozilla Firefox\firefox.exe"),
        Path(r"C:\Program Files (x86)\Mozilla Firefox\firefox.exe"),
        Path(os.path.expandvars(r"%LOCALAPPDATA%\Mozilla Firefox\firefox.exe")),
        Path(os.path.expandvars(r"%APPDATA%\Mozilla Firefox\firefox.exe")),
    ]
    # 3) Scoop / winget / portable in PATH
    which = shutil.which("firefox") or shutil.which("firefox.exe")
    if which:
        candidates.insert(0, Path(which))

    for c in candidates:
        if c.is_file():
            return str(c)
    return ""


def default_profile_dir() -> str:
    return str(app_data_dir() / "browser_profiles" / "firefox_hardened")


def _firefox_is_running_on_system() -> bool:
    """Check if any firefox.exe process is alive (tasklist, no admin needed)."""
    try:
        r = subprocess.run(
            ["tasklist", "/FI", "IMAGENAME eq firefox.exe", "/NH", "/FO", "CSV"],
            capture_output=True, text=True, timeout=6,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0x08000000),
        )
        return "firefox.exe" in (r.stdout or "").lower()
    except Exception:
        return False


def _has_descendants(root_pid: int) -> bool:
    """True iff *root_pid* or any descendant process is currently alive.

    Walks the process tree from root_pid using WMIC ProcessId/ParentProcessId
    (no admin required).  Falls back to checking the root PID's existence via
    tasklist if WMIC is missing on the host.
    """
    if not root_pid:
        return False
    # Fast path: is the root PID itself still alive?
    try:
        r = subprocess.run(
            ["tasklist", "/FI", f"PID eq {root_pid}", "/NH", "/FO", "CSV"],
            capture_output=True, text=True, timeout=6,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0x08000000),
        )
        if str(root_pid) in (r.stdout or ""):
            return True
    except Exception:  # noqa: BLE001
        pass
    # Walk descendants via WMIC.  Builds {parent: [child, ...]} then BFS.
    try:
        w = subprocess.run(
            ["wmic", "process", "get", "ProcessId,ParentProcessId", "/FORMAT:CSV"],
            capture_output=True, text=True, timeout=10,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0x08000000),
        )
        children: dict[int, list[int]] = {}
        for line in (w.stdout or "").splitlines():
            parts = [p.strip() for p in line.split(",")]
            # CSV header: Node,ParentProcessId,ProcessId
            if len(parts) < 3 or not parts[1].isdigit() or not parts[2].isdigit():
                continue
            ppid, pid = int(parts[1]), int(parts[2])
            children.setdefault(ppid, []).append(pid)
        stack = [root_pid]
        seen = {root_pid}
        while stack:
            cur = stack.pop()
            kids = children.get(cur, [])
            for k in kids:
                if k in seen:
                    continue
                seen.add(k)
                # Any live descendant = our session is still alive.
                return True
        return False
    except Exception:  # noqa: BLE001
        # WMIC missing (newer Windows) — be conservative and return False
        # rather than the old "any firefox.exe = mine" heuristic.
        return False


class BrowserSession:
    def __init__(self) -> None:
        self._proc: subprocess.Popen[str] | None = None
        self._profile_path: Path | None = None
        self._launched_at: float = 0.0

    def is_running(self) -> bool:
        if self._proc is None:
            return False
        # Firefox's initial launcher process exits quickly (code 0) while
        # browser child processes carry on.  Poll the parent process first.
        if self._proc.poll() is None:
            return True
        # Parent has exited — check whether any process in our spawned tree
        # (the original PID's children) is still alive.  Using a PID-scoped
        # WMIC query avoids the old failure mode of treating any unrelated
        # firefox.exe on the machine as ours.
        launched_pid = self._proc.pid
        return _has_descendants(launched_pid)

    def pid(self) -> int | None:
        if self._proc is None:
            return None
        return self._proc.pid

    def launch(
        self,
        cfg: BrowserConfig,
        proxy_host: str,
        proxy_port: int,
        start_url: str = "",
        signup_draft: SignupDraft | None = None,
    ) -> tuple[bool, str]:
        raw_exe = (cfg.firefox_path or default_firefox_path()).strip().strip('"').strip("'")
        if not raw_exe:
            return (
                False,
                "Firefox not found.  Install Firefox then set the path in the Browser tab "
                "(or browse to it with the Browse button).",
            )
        exe = Path(raw_exe)
        if not exe.is_file():
            return False, (
                f"Firefox executable not found: {raw_exe}  "
                "Use the Browse button on the Browser tab to locate firefox.exe."
            )

        raw_profile = (cfg.profile_dir or default_profile_dir()).strip().strip('"').strip("'")
        profile_dir = Path(raw_profile)

        hard = FirefoxHardening(
            force_proxy=cfg.force_proxy,
            disable_webrtc=cfg.disable_webrtc,
            resist_fingerprinting=cfg.resist_fingerprinting,
            disable_telemetry=cfg.disable_telemetry,
            first_party_isolation=cfg.first_party_isolation,
            strict_tracking_protection=cfg.strict_tracking_protection,
            clear_on_shutdown=cfg.clear_on_close,
            timezone_utc=cfg.timezone_utc,
            persona_key=cfg.persona,
            cookie_mode=cfg.cookie_mode,
        )

        try:
            if cfg.lock_managed_profile:
                ensure_firefox_profile(profile_dir, proxy_host, int(proxy_port), hard)
                if signup_draft is not None:
                    install_signup_extension(profile_dir, signup_draft)
        except Exception as e:
            return False, f"Profile setup failed: {e!s}"

        if self.is_running():
            self.stop()

        if cfg.lock_managed_profile:
            args = [str(exe), "-no-remote", "-profile", str(profile_dir)]
        else:
            args = [str(exe)]

        url = (start_url or "").strip()
        if url:
            args.append(url)

        env = os.environ.copy()
        env.update(persona_env(hard))

        try:
            self._proc = subprocess.Popen(
                args,
                # DETACHED_PROCESS ensures Firefox outlives Python and is not
                # killed when Python's job object exits.  CREATE_NO_WINDOW
                # suppresses any console window.  CREATE_NEW_PROCESS_GROUP
                # isolates Ctrl-C handling.
                creationflags=_CREATE_FLAGS,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                text=True,
                env=env,
            )
        except Exception as e:
            self._proc = None
            return False, f"Launch failed: {e!s}"

        self._launched_at = time.monotonic()

        # Firefox's multi-process launcher exits (code 0) within ~0.5 s while
        # the actual browser window loads in child processes.  We therefore do
        # NOT treat a quick exit of the parent as an error — we just verify
        # that *some* firefox.exe appears in the process list within 3 s.
        deadline = time.monotonic() + 3.0
        found = False
        while time.monotonic() < deadline:
            time.sleep(0.3)
            rc = self._proc.poll()
            # If it already exited with an error code, report it.
            if rc is not None and rc not in (0, None):
                self._proc = None
                return False, f"Firefox exited with error code {rc}."
            if _firefox_is_running_on_system():
                found = True
                break

        if not found:
            # One last check — maybe the profile write was the only action needed
            # and Firefox opened fast; accept success if parent is still alive.
            if self._proc.poll() is None:
                found = True

        if not found:
            self._proc = None
            return False, (
                "Firefox did not appear in the process list after 3 s.  "
                "Check that the path is correct and Firefox is not already blocked by another -profile lock."
            )

        self._profile_path = profile_dir if cfg.lock_managed_profile else None
        pid = self._proc.pid
        log.info("Launched hardened Firefox pid=%s profile=%s", pid, profile_dir)
        mode = "hardened profile" if cfg.lock_managed_profile else "unlocked mode"
        return True, f"Firefox started ({mode}, pid {pid})."

    def stop(self, dispose: bool = False) -> tuple[bool, str]:
        # Capture our PID before clearing the handle
        launched_pid: int | None = self._proc.pid if self._proc else None

        # Terminate the Popen handle if still alive
        if self._proc and self._proc.poll() is None:
            try:
                self._proc.terminate()
                self._proc.wait(timeout=8)
            except Exception:
                try:
                    self._proc.kill()
                except Exception:
                    pass
            finally:
                self._proc = None

        # Kill only the process tree we spawned (handles the detached-child
        # case where the launcher process already exited naturally).
        # /T kills the entire child tree; /PID scopes to our PID only.
        if launched_pid is not None:
            try:
                subprocess.run(
                    ["taskkill", "/F", "/T", "/PID", str(launched_pid)],
                    capture_output=True,
                    timeout=10,
                    creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0x08000000),
                )
            except Exception:
                pass

        self._proc = None
        if dispose and self._profile_path and self._profile_path.exists():
            try:
                shutil.rmtree(self._profile_path, ignore_errors=True)
            except Exception:
                pass
        self._launched_at = 0.0
        return True, "Browser stopped."