#!/usr/bin/env bash
# OpenLimits Auto-Setup
# Generated for: http://openlimits.app
set -euo pipefail

# ── Ensure python3 is available ──────────────────────────────────────
if ! command -v python3 &>/dev/null; then
    echo ""
    echo -e "\033[1mOpenLimits Setup\033[0m"
    echo ""
    echo -e "\033[0;33m⚠ python3 is required but not installed.\033[0m"
    echo ""

    INSTALL_CMD=""
    if [[ "$(uname)" == "Darwin" ]]; then
        if command -v brew &>/dev/null; then
            INSTALL_CMD="brew install python3"
        else
            INSTALL_CMD='/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" && brew install python3'
        fi
    elif command -v apt-get &>/dev/null; then
        INSTALL_CMD="sudo apt-get update && sudo apt-get install -y python3"
    elif command -v dnf &>/dev/null; then
        INSTALL_CMD="sudo dnf install -y python3"
    elif command -v pacman &>/dev/null; then
        INSTALL_CMD="sudo pacman -S --noconfirm python"
    elif command -v apk &>/dev/null; then
        INSTALL_CMD="sudo apk add python3"
    fi

    if [[ -n "$INSTALL_CMD" ]] && { [[ -t 0 ]] || [[ -t 2 ]]; }; then
        echo "  Suggested command:"
        echo -e "  \033[2m$INSTALL_CMD\033[0m"
        echo ""
        read -p "Install python3 now? [Y/n] " CHOICE < /dev/tty
        CHOICE="${CHOICE:-y}"
        if [[ "${CHOICE}" =~ ^[Yy]([Ee][Ss])?$ ]]; then
            echo ""
            eval "$INSTALL_CMD"
            echo ""
            if ! command -v python3 &>/dev/null; then
                echo -e "\033[0;31m✗ python3 still not found after install. Please install manually and re-run.\033[0m"
                exit 1
            fi
        else
            echo -e "\033[0;31m✗ Setup requires python3. Please install it and re-run.\033[0m"
            exit 1
        fi
    else
        if [[ -n "$INSTALL_CMD" ]]; then
            echo "  Install it with:"
            echo -e "  \033[1m$INSTALL_CMD\033[0m"
        else
            echo "  Install python3 using your system's package manager."
        fi
        echo ""
        echo "  Then re-run this setup command."
        exit 1
    fi
fi

# ── Hand off to Python ───────────────────────────────────────────────
exec python3 -u - <<'PYTHON_SETUP'
import json, os, re, shutil, subprocess, sys, platform

PROXY_URL = "http://openlimits.app"
TOKEN = ""

# ── Colors ───────────────────────────────────────────────────────────
GREEN  = "\033[0;32m"
YELLOW = "\033[0;33m"
RED    = "\033[0;31m"
BOLD   = "\033[1m"
DIM    = "\033[2m"
NC     = "\033[0m"

def ok(msg):   print(f"{GREEN}✓{NC} {msg}")
def warn(msg): print(f"{YELLOW}⚠ {msg}{NC}")
def fail(msg): print(f"{RED}✗ {msg}{NC}")

def prompt_yn(question, default_yes=True):
    """Ask a yes/no question interactively."""
    try:
        suffix = "[Y/n]" if default_yes else "[y/N]"
        answer = input(f"{question} {suffix} ").strip().lower()
        if not answer:
            return default_yes
        return answer.startswith("y")
    except (EOFError, KeyboardInterrupt):
        return default_yes if default_yes else False

def is_interactive():
    return sys.stdin.isatty() or sys.stderr.isatty()

# ── Token ────────────────────────────────────────────────────────────
print()
print(f"{BOLD}OpenLimits Setup{NC}")
print(f"URL: {PROXY_URL}")

token = TOKEN
if not token:
    print()
    print("Get your API token from your admin.")
    print()
    try:
        token = input(f"{BOLD}Enter your API token:{NC} ").strip()
    except (EOFError, KeyboardInterrupt):
        token = ""
    if not token:
        fail("Token cannot be empty")
        sys.exit(1)

masked = token[:8] + "..." + token[-4:] if len(token) > 12 else token
print(f"Token: {masked}")
print()

# ── Claude Code install check ────────────────────────────────────────
home = os.path.expanduser("~")

def find_claude():
    if shutil.which("claude"):
        return shutil.which("claude")
    for p in [
        os.path.join(home, ".claude", "local", "claude"),
        os.path.join(home, ".local", "bin", "claude"),
        "/usr/local/bin/claude",
        "/opt/homebrew/bin/claude",
    ]:
        if os.path.isfile(p) and os.access(p, os.X_OK):
            return p
    return None

claude_bin = find_claude()
if claude_bin:
    ok(f"Claude Code found at {DIM}{claude_bin}{NC}")
else:
    warn("Claude Code is not installed")
    if is_interactive() and prompt_yn("  Install Claude Code now?", default_yes=True):
        print()
        print(f"  {DIM}curl -fsSL https://claude.ai/install.sh | bash{NC}")
        print()
        result = subprocess.run(
            ["bash", "-c", "curl -fsSL https://claude.ai/install.sh | bash"],
            capture_output=False,
        )
        if result.returncode == 0:
            print()
            ok("Claude Code installed")
            claude_bin = find_claude()
        else:
            print()
            fail("Installation failed")
            print("  Try manually: curl -fsSL https://claude.ai/install.sh | bash")
    elif not is_interactive():
        print(f"  Install it with: {BOLD}curl -fsSL https://claude.ai/install.sh | bash{NC}")

# ── Merge ~/.claude/settings.json ────────────────────────────────────
settings_dir = os.path.join(home, ".claude")
settings_file = os.path.join(settings_dir, "settings.json")
os.makedirs(settings_dir, exist_ok=True)

env_updates = {
    "ANTHROPIC_BASE_URL": PROXY_URL,
    "ANTHROPIC_AUTH_TOKEN": token,
    "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "true",
    "API_TIMEOUT_MS": "3000000",
    "ANTHROPIC_DEFAULT_HAIKU_MODEL": "claude-haiku-4-5-20251001",
    "ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-sonnet-4-6",
    "ANTHROPIC_DEFAULT_OPUS_MODEL": "claude-opus-4-6",
}

if os.path.isfile(settings_file):
    try:
        with open(settings_file, "r") as f:
            settings = json.load(f)
        if not isinstance(settings, dict):
            settings = {}
    except (json.JSONDecodeError, ValueError):
        settings = {}

    if "env" not in settings or not isinstance(settings.get("env"), dict):
        settings["env"] = {}

    for k, v in env_updates.items():
        settings["env"][k] = v
    settings["env"].pop("ANTHROPIC_API_KEY", None)

    with open(settings_file, "w") as f:
        json.dump(settings, f, indent=2)
        f.write("\n")
    ok("Merged into existing settings.json (hooks and other settings preserved)")
else:
    settings = {"env": env_updates}
    with open(settings_file, "w") as f:
        json.dump(settings, f, indent=2)
        f.write("\n")
    ok("Created settings.json")

# ── Merge ~/.claude.json (onboarding) ────────────────────────────────
claude_json = os.path.join(home, ".claude.json")
if os.path.isfile(claude_json):
    try:
        with open(claude_json, "r") as f:
            cj = json.load(f)
        if not isinstance(cj, dict):
            cj = {}
    except (json.JSONDecodeError, ValueError):
        cj = {}
else:
    cj = {}

cj["hasCompletedOnboarding"] = True
slugs = cj.get("completedTosSlugs", [])
if not isinstance(slugs, list):
    slugs = []
if "termsOfService" not in slugs:
    slugs.append("termsOfService")
cj["completedTosSlugs"] = slugs

with open(claude_json, "w") as f:
    json.dump(cj, f, indent=2)
    f.write("\n")
ok("Onboarding and login skipped")

# ── Merge ~/.codex/config.toml ───────────────────────────────────────
codex_dir = os.path.join(home, ".codex")
codex_config = os.path.join(codex_dir, "config.toml")
os.makedirs(codex_dir, exist_ok=True)

provider_block = f'''model_provider = "openlimits"

[model_providers.openlimits]
name = "OpenLimits"
base_url = "{PROXY_URL}/v1"
env_key = "OPENAI_API_KEY"
wire_api = "responses"
'''

if os.path.isfile(codex_config):
    with open(codex_config, "r", encoding="utf-8-sig") as f:
        content = f.read()
    # Strip existing openlimits provider line and section
    content = re.sub(r'^model_provider\s*=\s*"openlimits"\s*\n?', "", content, flags=re.MULTILINE)
    content = re.sub(r'\[model_providers\.openlimits\][^\[]*', "", content, flags=re.DOTALL)
    content = content.rstrip()
    if content:
        content += "\n\n"
    content += provider_block
    with open(codex_config, "w") as f:
        f.write(content)
    ok("Merged Codex config (existing settings preserved)")
else:
    with open(codex_config, "w") as f:
        f.write(provider_block)
    ok("Created Codex config")

# ── Set OPENAI_API_KEY in shell profile ──────────────────────────────
shell_rc = None
for rc in [".zshrc", ".bashrc", ".bash_profile"]:
    p = os.path.join(home, rc)
    if os.path.isfile(p):
        shell_rc = p
        break

if shell_rc:
    # Resolve symlinks
    shell_rc_real = os.path.realpath(shell_rc)
    with open(shell_rc_real, "r") as f:
        lines = f.readlines()
    # Remove old export lines for these vars
    drop_patterns = ["export OPENAI_API_KEY=", "export OPENAI_BASE_URL=", "export CODEX_OPENAI_BASE_URL="]
    lines = [l for l in lines if not any(pat in l for pat in drop_patterns)]
    lines.append(f"export OPENAI_API_KEY='{token}'\n")
    with open(shell_rc_real, "w") as f:
        f.writelines(lines)
    os.environ["OPENAI_API_KEY"] = token
    ok(f"Set OPENAI_API_KEY in {os.path.basename(shell_rc)}")
else:
    warn(f"Add to your shell profile: export OPENAI_API_KEY='{token}'")

# On macOS, also set via launchctl for GUI apps
if platform.system() == "Darwin":
    try:
        subprocess.run(["launchctl", "setenv", "OPENAI_API_KEY", token],
                       capture_output=True, timeout=5)
        ok("Set OPENAI_API_KEY for GUI apps (launchctl)")
    except Exception:
        pass

# ── Warn about ANTHROPIC_API_KEY conflicts ───────────────────────────
if os.environ.get("ANTHROPIC_API_KEY"):
    warn("ANTHROPIC_API_KEY is set — this will bypass OpenLimits")
    print(f"  Run: {BOLD}unset ANTHROPIC_API_KEY{NC}")
    print("  And remove it from your shell profile (~/.bashrc, ~/.zshrc)")
else:
    ok("No conflicting env vars")

# ── Test connection ──────────────────────────────────────────────────
proxy_ok = False
try:
    from urllib.request import urlopen, Request
    from urllib.error import URLError
    r = urlopen(Request(f"{PROXY_URL}/status"), timeout=5)
    if r.status == 200:
        ok("OpenLimits reachable")
        proxy_ok = True
    else:
        warn(f"OpenLimits returned HTTP {r.status}")
except Exception:
    warn("Could not connect to OpenLimits (check VPN)")

# ── Validate token ───────────────────────────────────────────────────
if proxy_ok:
    try:
        req = Request(f"{PROXY_URL}/health", headers={"x-api-key": token})
        r = urlopen(req, timeout=10)
        if r.status == 200:
            ok("Token valid")
    except Exception as e:
        code = getattr(getattr(e, "response", None), "status", None) or getattr(e, "code", None)
        if code == 401:
            fail("Token rejected — check with your admin")
        elif code == 403:
            fail("Token has a permission or spending limit issue — contact your admin")
        elif code:
            warn(f"Got HTTP {code} — token may still work")
        else:
            warn("Could not validate token")

# ── Done ─────────────────────────────────────────────────────────────
print()
if shell_rc:
    print(f"{GREEN}{BOLD}Done!{NC} Now run:")
    print()
    print(f"  {BOLD}source {shell_rc}{NC}")
    print()
    print(f"Then run {BOLD}claude{NC} or {BOLD}codex{NC} in any project directory.")
else:
    print(f"{GREEN}{BOLD}Done!{NC} Open a new terminal, then run {BOLD}claude{NC} or {BOLD}codex{NC} in any project directory.")
print()

# ── Offer to launch Claude Code ─────────────────────────────────────
if claude_bin and is_interactive():
    if prompt_yn("Launch Claude Code now?", default_yes=False):
        print()
        os.execv(claude_bin, [claude_bin])
PYTHON_SETUP