#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 JacobJandon — https://github.com/JacobJandon/OnionClaw
"""
OnionClaw — First-Run Setup Wizard
===================================
Creates .env from .env.example (if not present), prompts for LLM credentials,
and writes a minimal Tor torrc that enables the ControlPort needed for
renew_identity() (circuit rotation).

Usage:
  python3 setup.py
"""
import os
import sys
import shutil
import subprocess
import json
import urllib.request

_DIR   = os.path.dirname(os.path.abspath(__file__))
_ENV   = os.path.join(_DIR, ".env")
_EXAMPLE = os.path.join(_DIR, ".env.example")
# Also check parent dir for the Sicry repo copy
_ENV_PARENT   = os.path.join(_DIR, "..", ".env")
_EXAMPLE_PARENT = os.path.join(_DIR, "..", ".env.example")

TORRC_SNIPPET = """\
# ── OnionClaw required settings ──────────────────────────────────
ControlPort 9051
CookieAuthentication 1
# CookieAuthFileGroupReadable 1  ← uncomment + add user to debian-tor group
#                                   for permanent cookie auth across restarts
# DataDirectory /tmp/tor_data   # uncomment to use a custom data dir
# ─────────────────────────────────────────────────────────────────
"""

# AUTH-1: torrc snippet that permanently fixes cookie permissions
TORRC_COOKIE_FIX = "CookieAuthFileGroupReadable 1\n"

# Systemd drop-in path for Tor (used as fallback when torrc can't be written)
SYSTEMD_DROPIN_DIR  = "/etc/systemd/system/tor.service.d"
SYSTEMD_DROPIN_PATH = os.path.join(SYSTEMD_DROPIN_DIR, "onionclaw-cookie.conf")
SYSTEMD_DROPIN_CONTENT = """\
# ── OnionClaw AUTH-1 fix ─────────────────────────────────────────────────────
# Makes the Tor cookie file group-readable so users in the debian-tor group
# can authenticate to the ControlPort without needing root.
# Generated by:  python3 setup.py
# ─────────────────────────────────────────────────────────────────────────────
[Service]
ExecStartPost=+/bin/sh -c 'chmod g+r /run/tor/control.authcookie 2>/dev/null; chmod g+r /var/lib/tor/control_auth_cookie 2>/dev/null; true'
"""


TORRC_CANDIDATE_PATHS = [
    "/etc/tor/torrc",
    "/usr/local/etc/tor/torrc",
    os.path.expanduser("~/.torrc"),
]

# Custom torrc written to /tmp/tor_data/torrc when system torrc is unwritable
CUSTOM_TORRC_DIR  = "/tmp/tor_data"
CUSTOM_TORRC_PATH = os.path.join(CUSTOM_TORRC_DIR, "torrc")

BOLD  = "\033[1m"
GREEN = "\033[32m"
YELLOW= "\033[33m"
RED   = "\033[31m"
CYAN  = "\033[36m"
RESET = "\033[0m"


def _p(msg, color=""):
    print(f"{color}{msg}{RESET}" if color else msg)


def _ask(prompt, default=""):
    try:
        val = input(f"{CYAN}{prompt}{RESET} [{default}]: ").strip()
        return val if val else default
    except (KeyboardInterrupt, EOFError):
        print()
        sys.exit(0)


def _yes(prompt):
    try:
        ans = input(f"{CYAN}{prompt}{RESET} [Y/n]: ").strip().lower()
        return ans in ("", "y", "yes")
    except (KeyboardInterrupt, EOFError):
        print()
        sys.exit(0)


# ─────────────────────────────────────────────────────────────────
# 1.  .env setup
# ─────────────────────────────────────────────────────────────────

def setup_env():
    _p("\n═══════════════════════════════════════════════════", BOLD)
    _p("  OnionClaw Setup — Step 1: .env configuration", BOLD)
    _p("═══════════════════════════════════════════════════", BOLD)

    if os.path.exists(_ENV):
        _p(f"\n✓ .env already exists: {_ENV}", GREEN)
        # Always enforce correct permissions — chmod 600 even when skipping reconfigure.
        # Existing installs that ran setup before the permission fix still need this.
        try:
            os.chmod(_ENV, 0o600)
            _p(f"✓ .env permissions set to 600 (owner read/write only)", GREEN)
        except OSError as e:
            _p(f"⚠  Could not set .env permissions: {e}", YELLOW)
        if not _yes("  Re-configure anyway?"):
            return

    # Find the best .env.example
    example = _EXAMPLE if os.path.exists(_EXAMPLE) else (
              _EXAMPLE_PARENT if os.path.exists(_EXAMPLE_PARENT) else None)

    if example:
        shutil.copy2(example, _ENV)
        _p(f"\n✓ Copied {os.path.relpath(example, _DIR)} → {os.path.relpath(_ENV, _DIR)}", GREEN)
    else:
        # Write a minimal .env if no example found
        with open(_ENV, "w") as f:
            f.write("# OnionClaw .env — generated by setup.py\n")
            f.write("LLM_PROVIDER=openai\n")
            f.write("TOR_SOCKS_HOST=127.0.0.1\nTOR_SOCKS_PORT=9050\n")
            f.write("TOR_CONTROL_HOST=127.0.0.1\nTOR_CONTROL_PORT=9051\n")
            f.write("SICRY_MAX_CHARS=8000\nSICRY_CACHE_TTL=600\n")
        _p(f"\n✓ Created minimal .env at {_ENV}", GREEN)

    # Interactive prompts ──────────────────────────────────────────
    _p("\nChoose your LLM provider:")
    _p("  1) openai      (OPENAI_API_KEY required)")
    _p("  2) anthropic   (ANTHROPIC_API_KEY required)")
    _p("  3) gemini      (GEMINI_API_KEY required)")
    _p("  4) ollama      (local — no key needed)")
    _p("  5) llamacpp    (local — no key needed)")
    choice = _ask("Enter number", "1")
    provider_map = {"1": "openai", "2": "anthropic", "3": "gemini",
                    "4": "ollama",  "5": "llamacpp"}
    provider = provider_map.get(choice, "openai")

    api_key = ""
    key_var = ""
    if provider == "openai":
        key_var = "OPENAI_API_KEY"
        api_key = _ask("OpenAI API key (sk-...)", "")
    elif provider == "anthropic":
        key_var = "ANTHROPIC_API_KEY"
        api_key = _ask("Anthropic API key (sk-ant-...)", "")
    elif provider == "gemini":
        key_var = "GEMINI_API_KEY"
        api_key = _ask("Gemini API key (AIza...)", "")
    elif provider in ("ollama", "llamacpp"):
        _p(f"  → {provider} uses a local server — no API key needed.", GREEN)

    # Patch the .env file
    _patch_env("LLM_PROVIDER", provider)
    if key_var and api_key:
        _patch_env(key_var, api_key)
    _patch_env("SICRY_CACHE_TTL", "600")  # ensure cache TTL is set

    _p(f"\n✓ .env configured with provider={provider}", GREEN)

    # Lock permissions — .env contains API keys; must not be world-readable
    try:
        os.chmod(_ENV, 0o600)
        _p(f"✓ .env permissions set to 600 (owner read/write only)", GREEN)
    except OSError as e:
        _p(f"⚠  Could not set .env permissions: {e}", YELLOW)

    # Copy to parent Sicry repo as well
    parent_env = os.path.join(_DIR, "..", ".env")
    if not os.path.exists(parent_env):
        try:
            shutil.copy2(_ENV, parent_env)
            try:
                os.chmod(parent_env, 0o600)
            except OSError:
                pass
            _p(f"✓ Also copied to parent Sicry repo: {os.path.abspath(parent_env)}", GREEN)
        except Exception:
            pass


def _patch_env(key: str, value: str):
    """Update or append a key=value line in .env."""
    if not os.path.exists(_ENV):
        return
    with open(_ENV, "r") as f:
        lines = f.readlines()
    found = False
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith(f"{key}=") or stripped.startswith(f"# {key}="):
            lines[i] = f"{key}={value}\n"
            found = True
            break
    if not found:
        lines.append(f"{key}={value}\n")
    with open(_ENV, "w") as f:
        f.writelines(lines)


# ─────────────────────────────────────────────────────────────────
# 2.  Tor ControlPort / torrc setup
# ─────────────────────────────────────────────────────────────────

def setup_tor():
    _p("\n═══════════════════════════════════════════════════", BOLD)
    _p("  OnionClaw Setup — Step 2: Tor ControlPort", BOLD)
    _p("═══════════════════════════════════════════════════", BOLD)
    _p("""
renew_identity() needs a Tor ControlPort to rotate circuits.
Without it, Tor keeps the same identity for the entire session.

Required torrc lines:
  ControlPort 9051
  CookieAuthentication 1
""")

    # Check if tor is installed
    tor_bin = shutil.which("tor")
    if not tor_bin:
        _p("⚠  Tor binary not found in PATH.", YELLOW)
        _p("   Install:  sudo apt install tor   (Debian/Ubuntu)", YELLOW)
        _p("             brew install tor        (macOS)", YELLOW)
        _p("             pacman -S tor           (Arch)", YELLOW)
    else:
        _p(f"✓ Tor found: {tor_bin}", GREEN)

    # Try to patch system torrc
    patched_system = False
    for torrc_path in TORRC_CANDIDATE_PATHS:
        if os.path.exists(torrc_path):
            if _already_has_controlport(torrc_path):
                _p(f"\n✓ ControlPort already configured in {torrc_path}", GREEN)
                patched_system = True
                break
            if _yes(f"\nPatch {torrc_path} to add ControlPort settings? (requires write access)"):
                try:
                    with open(torrc_path, "r") as f:
                        existing = f.read()
                    with open(torrc_path, "a") as f:
                        f.write(f"\n{TORRC_SNIPPET}")
                    _p(f"✓ Patched {torrc_path}", GREEN)
                    patched_system = True
                    break
                except PermissionError:
                    _p(f"  Permission denied — try:  sudo python3 setup.py", YELLOW)
                except Exception as e:
                    _p(f"  Failed: {e}", RED)

    if not patched_system:
        # Fall back: write custom torrc + instruct user to start tor with it
        _write_custom_torrc()

    # Print instructions regardless
    _p("""
─────────────────────────────────────────────────────────
  AFTER configuring torrc, restart Tor:

    sudo systemctl restart tor        # systemd (Debian/Ubuntu)
    sudo service tor restart          # SysV init
    brew services restart tor         # macOS Homebrew
    tor -f /tmp/tor_data/torrc &      # custom torrc (no root needed)

─────────────────────────────────────────────────────────
""", CYAN)

    # Verify ControlPort is reachable now
    _verify_controlport()


def _already_has_controlport(path: str) -> bool:
    try:
        with open(path, "r") as f:
            content = f.read()
        return any(
            line.strip().lower().startswith("controlport")
            for line in content.splitlines()
            if not line.strip().startswith("#")
        )
    except Exception:
        return False


def _write_custom_torrc():
    _p(f"\nWriting custom torrc → {CUSTOM_TORRC_PATH}", CYAN)
    try:
        os.makedirs(CUSTOM_TORRC_DIR, exist_ok=True)
        if os.path.exists(CUSTOM_TORRC_PATH) and _already_has_controlport(CUSTOM_TORRC_PATH):
            _p(f"✓ ControlPort already in {CUSTOM_TORRC_PATH}", GREEN)
        else:
            full_torrc = (
                "# OnionClaw custom torrc\n"
                f"DataDirectory {CUSTOM_TORRC_DIR}\n"
                + TORRC_SNIPPET
            )
            with open(CUSTOM_TORRC_PATH, "w") as f:
                f.write(full_torrc)
            _p(f"✓ Written: {CUSTOM_TORRC_PATH}", GREEN)
            _p(f"\n  Also set in .env:")
            _p(f"    TOR_DATA_DIR={CUSTOM_TORRC_DIR}")
        # Set TOR_DATA_DIR in .env so sicry.py finds the cookie
        _patch_env("TOR_DATA_DIR", CUSTOM_TORRC_DIR)
        _p(f"✓ TOR_DATA_DIR={CUSTOM_TORRC_DIR} written to .env", GREEN)
    except Exception as e:
        _p(f"  Could not write custom torrc: {e}", RED)
        _p("  Manual setup required — see README.md#tor-control-port", YELLOW)


def _fix_cookie_auth():
    """AUTH-1: apply permanent cookie-auth fix — group membership + torrc + systemd drop-in.

    The Tor cookie file (/run/tor/control.authcookie) is recreated on every
    restart with permissions that only the root / debian-tor group can read.
    Three complementary fixes are applied here so the fix survives reboots:

      1.  Add user to debian-tor group   (usermod -aG debian-tor $USER)
      2.  Add CookieAuthFileGroupReadable 1 to the active torrc
      3.  Install a systemd drop-in that chmod g+r the cookie after Tor starts
    """
    import grp
    import getpass
    current_user = getpass.getuser()
    if current_user == "root":
        _p("  Running as root — cookie auth should already work.", GREEN)
        return

    # ── Check / apply group membership ───────────────────────────
    try:
        tor_gid = grp.getgrnam("debian-tor").gr_gid
        user_groups = [g.gr_name for g in grp.getgrall() if current_user in g.gr_mem]
        already_in_group = "debian-tor" in user_groups
    except KeyError:
        # No debian-tor group — different distro; skip
        _p("  ⚠  No debian-tor group found (non-Debian system).", YELLOW)
        _p("     Use password auth instead (FIX B above).", YELLOW)
        return

    if already_in_group:
        _p(f"✓ {current_user} is already in debian-tor group", GREEN)
    else:
        _p(f"\n  {current_user} is NOT in the debian-tor group.", YELLOW)
        if _yes(f"  Apply FIX A: run 'sudo usermod -aG debian-tor {current_user}' now?"):
            try:
                ret = subprocess.run(
                    ["sudo", "usermod", "-aG", "debian-tor", current_user],
                    capture_output=True, text=True,
                )
                if ret.returncode == 0:
                    _p(f"✓ Added {current_user} to debian-tor", GREEN)
                    _p("  ⚠  You must log out and back in (or run: newgrp debian-tor) for the group to take effect.", YELLOW)
                else:
                    _p(f"✗ usermod failed: {ret.stderr.strip()}", RED)
                    _p("  Try manually:  sudo usermod -aG debian-tor $USER", YELLOW)
            except FileNotFoundError:
                _p("✗ sudo not available — run manually:  sudo usermod -aG debian-tor $USER", RED)

    # ── Add CookieAuthFileGroupReadable 1 to the active torrc ─────
    for torrc_path in TORRC_CANDIDATE_PATHS:
        if not os.path.exists(torrc_path):
            continue
        try:
            with open(torrc_path, "r") as f:
                content = f.read()
            if "CookieAuthFileGroupReadable" in content:
                _p(f"✓ CookieAuthFileGroupReadable already in {torrc_path}", GREEN)
            else:
                with open(torrc_path, "a") as f:
                    f.write(TORRC_COOKIE_FIX)
                _p(f"✓ Added CookieAuthFileGroupReadable 1 to {torrc_path}", GREEN)
                _p("  Restart Tor for it to take effect.", YELLOW)
        except PermissionError:
            _p(f"  ⚠  Cannot write {torrc_path} (permission denied)", YELLOW)
            _p(f"     Run:  echo 'CookieAuthFileGroupReadable 1' | sudo tee -a {torrc_path}", YELLOW)
        except Exception as e:
            _p(f"  ⚠  Error patching torrc: {e}", YELLOW)
        break  # only patch the first found torrc

    # ── Install systemd drop-in as belt-and-suspenders ───────────
    if _yes("\n  Also install systemd drop-in (belt-and-suspenders: chmod g+r cookie after every Tor start)?"):
        try:
            os.makedirs(SYSTEMD_DROPIN_DIR, exist_ok=True)
            with open(SYSTEMD_DROPIN_PATH, "w") as f:
                f.write(SYSTEMD_DROPIN_CONTENT)
            subprocess.run(["sudo", "systemctl", "daemon-reload"],
                           capture_output=True, check=False)
            _p(f"✓ Systemd drop-in written to {SYSTEMD_DROPIN_PATH}", GREEN)
            _p("  After restarting Tor, the cookie will always be group-readable.", GREEN)
        except PermissionError:
            _p(f"  ⚠  Cannot write {SYSTEMD_DROPIN_PATH} — try with sudo:", YELLOW)
            _p(f"     sudo mkdir -p {SYSTEMD_DROPIN_DIR}", YELLOW)
            _p(f"     sudo tee {SYSTEMD_DROPIN_PATH} << 'EOF'", YELLOW)
            _p(SYSTEMD_DROPIN_CONTENT + "EOF", YELLOW)
            _p("     sudo systemctl daemon-reload && sudo systemctl restart tor", YELLOW)
        except Exception as e:
            _p(f"  ⚠  Drop-in install failed: {e}", YELLOW)


def _verify_controlport():
    _p("\nChecking ControlPort connectivity …")
    try:
        import socket
        s = socket.create_connection(("127.0.0.1", 9051), timeout=2)
        s.close()
        _p("✓ ControlPort 9051 is open", GREEN)
    except Exception:
        _p("✗ ControlPort 9051 not reachable yet", YELLOW)
        _p("  → Start/restart Tor with the updated torrc config above.", YELLOW)
        return
    # AUTH-1/AUTH-2: actually test authentication, not just the port
    _p("  Testing Tor authentication (stem) …")
    try:
        import stem
        from stem.control import Controller
        env_password = os.environ.get("TOR_CONTROL_PASSWORD", "")
        env_file = os.path.join(_DIR, ".env")
        if not env_password and os.path.exists(env_file):
            for line in open(env_file):
                if line.strip().startswith("TOR_CONTROL_PASSWORD="):
                    env_password = line.strip().split("=", 1)[1].strip().strip('"')
                    break
        with Controller.from_port(port=9051) as ctrl:
            if env_password:
                ctrl.authenticate(password=env_password)
                _p("✓ Auth OK   (password from TOR_CONTROL_PASSWORD)", GREEN)
            else:
                ctrl.authenticate()  # tries cookie + no-auth
                _p("✓ Auth OK   (cookie / no-auth)", GREEN)
    except ImportError:
        _p("  ⚠  stem not installed — skipping auth test (install with: pip install stem)", YELLOW)
    except Exception as auth_err:
        _p(f"✗ Auth failed: {auth_err}", RED)
        _p("""
  Authentication test failed. Three reliable fixes:

  FIX A — Add your user to the debian-tor group (cookie auth survives restarts):
    sudo usermod -aG debian-tor $USER
    newgrp debian-tor    # (or log out and back in)

  FIX B — Use password auth (most reliable — never breaks on restarts):
    1. Generate a hashed password:
         tor --hash-password 'YOUR_SECRET_PASS'
    2. Add to /etc/tor/torrc (or your custom torrc):
         HashedControlPassword <output from step 1>
    3. Add to .env in this directory:
         TOR_CONTROL_PASSWORD=YOUR_SECRET_PASS
    4. Restart Tor.

  FIX C — CookieAuthFileGroupReadable 1 (works with FIX A — permanent):
    Add to /etc/tor/torrc:
         CookieAuthFileGroupReadable 1
    Then restart Tor. The cookie stays group-readable across all restarts.
""", YELLOW)
        # AUTH-1: Offer to apply the permanent fix automatically
        if _yes("  Apply the permanent AUTH-1 fix now (FIX A + FIX C + systemd drop-in)?"):
            _fix_cookie_auth()



# ─────────────────────────────────────────────────────────────────
# 3.  Dependency check
# ─────────────────────────────────────────────────────────────────

def check_deps():
    _p("\n═══════════════════════════════════════════════════", BOLD)
    _p("  OnionClaw Setup — Step 3: Python dependencies", BOLD)
    _p("═══════════════════════════════════════════════════", BOLD)

    required = [
        ("requests",    "requests[socks]"),
        ("bs4",         "beautifulsoup4"),
        ("dotenv",      "python-dotenv"),
        ("stem",        "stem"),
    ]

    optional = [
        # MCP-1: mcp is needed only for `python3 sicry.py serve`
        ("mcp", "mcp",
         "optional — needed only for MCP server mode (python3 sicry.py serve)"),
    ]

    missing = []
    for pkg, pip_name in required:
        try:
            __import__(pkg)
            _p(f"  ✓ {pip_name}", GREEN)
        except ImportError:
            _p(f"  ✗ {pip_name}  — MISSING", RED)
            missing.append(pip_name)

    _p("\n  Optional packages:")
    for pkg, pip_name, note in optional:
        try:
            __import__(pkg)
            _p(f"  ✓ {pip_name}  ({note})", GREEN)
        except ImportError:
            _p(f"  ◦ {pip_name}  ({note})", YELLOW)
            _p(f"    Install: pip install {pip_name} --user")
            _p(f"       or:  pipx install {pip_name}   (isolated environment)")
            _p(f"       or:  pip install {pip_name} --break-system-packages   (Debian override)")

    if missing:
        _p(f"\n  Install missing packages:", YELLOW)
        cmd = f"pip install {' '.join(missing)}"
        _p(f"    {cmd}", YELLOW)
        if _yes("  Install them now?"):
            subprocess.run([sys.executable, "-m", "pip", "install", *missing])
    else:
        _p("\n  All required dependencies satisfied.", GREEN)


# ─────────────────────────────────────────────────────────────────
# 4.  Final summary
# ─────────────────────────────────────────────────────────────────

def summary():
    _p("\n═══════════════════════════════════════════════════", BOLD)
    _p("  Setup complete!", BOLD + GREEN)
    _p("═══════════════════════════════════════════════════", BOLD)
    _p(f"""
Quick-start commands:

  1. Start Tor (if not running):
       sudo systemctl start tor
       — or —
       tor -f {CUSTOM_TORRC_PATH} &

  2. Verify everything works:
       python3 sicry.py check_tor

  3. Run a full pipeline search:
       python3 pipeline.py --query "your topic"

  4. Run without an LLM key (raw results):
       python3 pipeline.py --query "your topic" --no-llm

  5. MCP server mode (for Claude Desktop / Cursor / Zed):
       python3 sicry.py serve
  6. Check engine health (with result caching):
       python3 check_engines.py --cached 10

  7. Sync sicry.py to a newer SICRY™ release:
       python3 sync_sicry.py               # latest main
       python3 sync_sicry.py --tag v1.2.0  # specific tag
       python3 sync_sicry.py --dry-run     # preview only

  8. Clear the fetch result cache:
       python3 pipeline.py --clear-cache
       — or — python3 sicry.py clear-cache
  Docs: https://github.com/JacobJandon/OnionClaw
""", CYAN)


def _check_update_notice() -> None:
    """Print a one-line notice if a newer OnionClaw release exists on GitHub."""
    _RELEASES_API = (
        "https://api.github.com/repos/JacobJandon/OnionClaw/releases/latest"
    )
    try:
        import sicry as _sicry_mod
        current = getattr(_sicry_mod, "__version__", None)
        if not current:
            return
        req = urllib.request.Request(
            _RELEASES_API,
            headers={"User-Agent": f"OnionClaw-setup/{current}"},
        )
        with urllib.request.urlopen(req, timeout=4) as resp:
            data = json.loads(resp.read())
        tag = data.get("tag_name", "").lstrip("v")
        if not tag:
            return
        def _ver(v):
            try:
                return tuple(int(x) for x in v.split("."))
            except Exception:
                return (0,)
        if _ver(tag) > _ver(current):
            _p(f"\n⚡  Update available: v{current} → v{tag}", YELLOW)
            url = data.get("html_url", "https://github.com/JacobJandon/OnionClaw/releases")
            _p(f"   Release page : {url}", CYAN)
            _p(f"   Upgrade with : git -C {_DIR} pull", CYAN)
            _p(f"                  python3 {os.path.join(_DIR, 'sync_sicry.py')}", CYAN)
    except Exception:
        pass  # network unavailable, rate-limited, etc. — always silent


# ─────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    _p(f"""
{'═' * 51}
  OnionClaw First-Run Setup Wizard
  Press Ctrl+C at any prompt to abort.
{'═' * 51}""", BOLD)

    setup_env()
    setup_tor()
    check_deps()
    summary()
    _check_update_notice()
