#!/usr/bin/env python3
"""
provision-agent.py — Playwright automation for AI agent identity provisioning.

Enrolls AI agent identities in Cloudron, logs into Gitea/Discourse/Redmine via
SSO, generates API keys, and stores all credentials in Bitwarden.

Usage:
    python3 provision-agent.py                       # provision all agents in manifest
    python3 provision-agent.py --agent vp-techops    # provision one agent
    python3 provision-agent.py --phase1-only         # Cloudron enrollment only
    python3 provision-agent.py --dry-run             # validate manifest without browser

Manifest: agents.yaml (see agents.yaml.example)

See: ~/Q3/agent-identity-bootstrap.md for the full architecture.
"""

import argparse
import json
import logging
import os
import sys
import time
from pathlib import Path

import yaml
from playwright.sync_api import BrowserContext, Page, sync_playwright

from bw_helper import BitwardenHelper

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("provision")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

CLOUDRON_BASE = os.environ.get("CLOUDRON_BASE", "https://tsys-cloudron.knel.net")
GITEA_URL = os.environ.get("GITEA_URL", "https://git.knownelement.com")
DISCOURSE_URL = os.environ.get("DISCOURSE_URL", "https://community.turnsys.com")
REDMINE_URL = os.environ.get("REDMINE_URL", "https://projects.knownelement.com")

STATE_DIR = Path("/app/state")

# ---------------------------------------------------------------------------
# Cloudron enrollment
# ---------------------------------------------------------------------------


def enroll_cloudron(
    page: Page,
    agent: dict,
    bw: BitwardenHelper,
) -> dict:
    """
    Phase 1: Accept Cloudron invite, set password, enable 2FA.

    Returns a dict with the agent's Cloudron credentials.
    """
    name = agent["name"]
    invite_url = agent["cloudron_invite"]
    display_name = agent.get("display_name", name)

    log.info(f"[{name}] Phase 1: Cloudron enrollment — {invite_url}")

    # Check if already provisioned
    item_name = f"{name} Cloudron"
    if bw.item_exists(item_name):
        log.info(f"[{name}] Cloudron credential already exists in Bitwarden — skipping")
        cloudron_email = agent.get("cloudron_email", f"{name}@turnsys.com")
        return {
            "username": cloudron_email,
            "password": bw.get_item_password(item_name),
        }

    # Generate a strong password
    password = bw.generate_password(length=32)
    log.info(f"[{name}] Generated password ({len(password)} chars)")

    # Navigate to invite link
    page.goto(invite_url, wait_until="networkidle")

    # Fill in the invite acceptance form
    # Cloudron invite page typically has: username (may be pre-filled), password, confirm password
    page.wait_for_selector('input[type="password"]', timeout=15000)

    password_inputs = page.query_selector_all('input[type="password"]')
    if len(password_inputs) >= 2:
        password_inputs[0].fill(password)
        password_inputs[1].fill(password)
    elif len(password_inputs) == 1:
        password_inputs[0].fill(password)
    else:
        raise RuntimeError(f"[{name}] No password field found on Cloudron invite page")

    # Set display name if field exists
    name_field = page.query_selector('input[name="displayName"], input[name="name"]')
    if name_field:
        name_field.fill(display_name)

    # Submit
    submit = page.query_selector('button[type="submit"], button:has-text("Setup"), button:has-text("Create"), button:has-text("Accept")')
    if submit:
        submit.click()

    page.wait_for_load_state("networkidle")
    log.info(f"[{name}] Invite accepted")

    # Enable 2FA
    totp_secret = enable_cloudron_2fa(page, agent, bw)

    # Store credential in Bitwarden
    cloudron_email = agent.get("cloudron_email", f"{name}@turnsys.com")
    bw.create_item(
        name=item_name,
        username=cloudron_email,
        password=password,
        uris=[CLOUDRON_BASE],
        collection_name=name,
        totp_secret=totp_secret,
    )
    log.info(f"[{name}] Cloudron credential stored in Bitwarden (collection: {name})")

    return {"username": cloudron_email, "password": password, "totp_secret": totp_secret}


def enable_cloudron_2fa(page: Page, agent: dict, bw: BitwardenHelper) -> str:
    """
    Navigate to Cloudron 2FA settings and enable TOTP.

    Returns the TOTP secret.
    """
    name = agent["name"]
    log.info(f"[{name}] Enabling 2FA on Cloudron account")

    # Navigate to account settings
    page.goto(f"{CLOUDRON_BASE}/settings.html#account", wait_until="networkidle")
    page.wait_for_timeout(2000)

    # Click "Enable 2FA" button
    enable_btn = page.query_selector('button:has-text("Enable"), button:has-text("2FA"), a:has-text("Enable")')
    if not enable_btn:
        log.warning(f"[{name}] Could not find 2FA enable button — may already be enabled")
        return ""

    enable_btn.click()
    page.wait_for_timeout(2000)

    # Extract TOTP secret from the QR code or the manual entry text
    # Cloudron shows a QR code and a text secret
    secret_text = page.query_selector('.modal-body code, .two-factor-secret, input[readonly]')
    if secret_text:
        totp_secret = secret_text.text_content().strip().replace(" ", "")
    else:
        # Try to extract from QR image source (base64)
        qr_img = page.query_selector('img[src*="data:image"]')
        if qr_img:
            qr_src = qr_img.get_attribute("src")
            totp_secret = decode_qr_from_base64(qr_src)
        else:
            log.error(f"[{name}] Could not extract TOTP secret from 2FA page")
            return ""

    log.info(f"[{name}] Extracted TOTP secret: {totp_secret[:4]}...")

    # Generate current TOTP code and confirm
    import pyotp
    totp_code = pyotp.TOTP(totp_secret).now()

    code_input = page.query_selector('input[name="totpToken"], input[name="token"], input[placeholder*="code"]')
    if code_input:
        code_input.fill(totp_code)
        confirm_btn = page.query_selector('button:has-text("Confirm"), button:has-text("Enable"), button[type="submit"]')
        if confirm_btn:
            confirm_btn.click()
            page.wait_for_timeout(2000)
            log.info(f"[{name}] 2FA confirmed")
    else:
        log.warning(f"[{name}] Could not find TOTP confirmation input")

    return totp_secret


def decode_qr_from_base64(data_uri: str) -> str:
    """Decode a TOTP secret from a base64 QR code data URI."""
    import base64
    import io

    from PIL import Image
    from pyzbar.pyzbar import decode

    # Extract base64 data from data URI
    header, b64data = data_uri.split(",", 1)
    img_bytes = base64.b64decode(b64data)
    img = Image.open(io.BytesIO(img_bytes))

    decoded = decode(img)
    if decoded:
        # TOTP URIs look like: otpauth://totp/Label?secret=XXXX&...
        uri = decoded[0].data.decode()
        if "secret=" in uri:
            return uri.split("secret=")[1].split("&")[0]

    raise RuntimeError("Could not decode TOTP secret from QR code")


# ---------------------------------------------------------------------------
# System access (Phase 2)
# ---------------------------------------------------------------------------


def sso_login(page: Page, system_url: str, agent: dict, bw: BitwardenHelper) -> bool:
    """
    Log into a Cloudron-managed app via SSO.

    Returns True if login succeeded.
    """
    name = agent["name"]
    log.info(f"[{name}] SSO login: {system_url}")

    # Navigate to the app — should redirect to Cloudron SSO
    page.goto(system_url, wait_until="networkidle")

    # If already logged in (SSO session), we're done
    if not page.query_selector('input[type="password"]'):
        log.info(f"[{name}] SSO session active — already logged in")
        return True

    # Fill Cloudron SSO login form
    cloudron_item = f"{name} Cloudron"
    cloudron_email = agent.get("cloudron_email", f"{name}@turnsys.com")
    password = bw.get_item_password(cloudron_item)

    user_input = page.query_selector('input[name="username"], input[type="email"], input[name="email"]')
    pass_input = page.query_selector('input[type="password"]')

    if user_input:
        user_input.fill(cloudron_email)
    if pass_input:
        pass_input.fill(password)

    # Handle 2FA if prompted
    submit = page.query_selector('button[type="submit"], button:has-text("Sign in"), button:has-text("Log in")')
    if submit:
        submit.click()
        page.wait_for_load_state("networkidle")

    # Check for TOTP prompt
    totp_input = page.query_selector('input[name="totpToken"], input[name="token"], input[placeholder*="code"], input[autocomplete*="one-time-code"]')
    if totp_input:
        totp_code = bw.get_totp(cloudron_item)
        totp_input.fill(totp_code)
        submit2 = page.query_selector('button[type="submit"]')
        if submit2:
            submit2.click()
            page.wait_for_load_state("networkidle")

    log.info(f"[{name}] SSO login complete for {system_url}")
    return True


def provision_gitea(page: Page, agent: dict, bw: BitwardenHelper) -> str:
    """Generate a Gitea API token via SSO login. Returns the token."""
    name = agent["name"]
    systems = agent.get("systems", {})
    gitea_cfg = systems.get("gitea", {})

    if not gitea_cfg:
        log.info(f"[{name}] No Gitea config — skipping")
        return ""

    item_name = f"{name} Gitea"
    if bw.item_exists(item_name):
        log.info(f"[{name}] Gitea token already exists — skipping")
        return bw.get_item_password(item_name)

    url = gitea_cfg.get("url", GITEA_URL)
    token_name = gitea_cfg.get("token_name", f"{name}-api")

    sso_login(page, f"{url}/user/login", agent, bw)

    # Navigate to API token settings
    page.goto(f"{url}/user/settings/applications", wait_until="networkidle")

    # Generate new token
    name_input = page.query_selector('input[name="name"]')
    if name_input:
        name_input.fill(token_name)

    # Select scopes if checkboxes exist
    for scope in gitea_cfg.get("scopes", ["api", "repo", "read:org"]):
        scope_cb = page.query_selector(f'input[value="{scope}"]')
        if scope_cb and not scope_cb.is_checked():
            scope_cb.check()

    gen_btn = page.query_selector('button:has-text("Generate Token")')
    if gen_btn:
        gen_btn.click()
        page.wait_for_timeout(2000)

    # Extract the generated token
    token_el = page.query_selector('.ui.info.message code, .ui.message code, input[readonly]')
    if not token_el:
        # Try the new Gitea UI
        token_el = page.query_selector('.token-value, .access-token')

    token = token_el.text_content().strip() if token_el else ""
    if not token:
        log.error(f"[{name}] Could not extract Gitea API token")
        return ""

    log.info(f"[{name}] Gitea token generated: {token[:8]}...")

    # Store in Bitwarden
    bw.create_item(
        name=item_name,
        username=name,
        password=token,
        uris=[url],
        collection_name=name,
        custom_fields={"token_name": token_name},
    )
    log.info(f"[{name}] Gitea token stored in Bitwarden")

    return token


def provision_discourse(page: Page, agent: dict, bw: BitwardenHelper) -> str:
    """
    Generate a Discourse API key via SSO login.

    Note: Discourse API keys typically require admin to create.
    If the agent can't self-generate, this logs a warning.
    Returns the API key (empty string if not possible).
    """
    name = agent["name"]
    systems = agent.get("systems", {})
    discourse_cfg = systems.get("discourse", {})

    if not discourse_cfg:
        log.info(f"[{name}] No Discourse config — skipping")
        return ""

    item_name = f"{name} Discourse"
    if bw.item_exists(item_name):
        log.info(f"[{name}] Discourse key already exists — skipping")
        return bw.get_item_password(item_name)

    url = discourse_cfg.get("url", DISCOURSE_URL)

    sso_login(page, f"{url}/", agent, bw)

    # Try to generate an API key from user preferences
    # Note: In Discourse, only admin can create API keys via UI
    # Non-admin users may not have this option
    page.goto(f"{url}/u/{name}/preferences/account", wait_until="networkidle")

    api_key_section = page.query_selector('.api-keys, [data-section="api-keys"]')

    if not api_key_section:
        log.warning(
            f"[{name}] Discourse API key self-generation not available. "
            "An admin must create the key. The agent will need a manually-created key."
        )
        return ""

    # If the section exists, try to create a key
    revoke_btn = page.query_selector('.api-keys button:has-text("Revoke")')
    if not revoke_btn:
        # No existing keys — create one
        gen_btn = page.query_selector('button:has-text("New API Key"), button:has-text("Create")')
        if gen_btn:
            gen_btn.click()
            page.wait_for_timeout(2000)

            # Read the key
            key_el = page.query_selector('.api-key-value, code')
            api_key = key_el.text_content().strip() if key_el else ""

            if api_key:
                bw.create_item(
                    name=item_name,
                    username=name,
                    password=api_key,
                    uris=[url],
                    collection_name=name,
                )
                log.info(f"[{name}] Discourse API key stored in Bitwarden")
                return api_key

    log.warning(f"[{name}] Could not generate Discourse API key")
    return ""


def provision_redmine(page: Page, agent: dict, bw: BitwardenHelper) -> str:
    """Get the Redmine API access key via SSO login. Returns the key."""
    name = agent["name"]
    systems = agent.get("systems", {})
    redmine_cfg = systems.get("redmine", {})

    if not redmine_cfg:
        log.info(f"[{name}] No Redmine config — skipping")
        return ""

    item_name = f"{name} Redmine"
    if bw.item_exists(item_name):
        log.info(f"[{name}] Redmine key already exists — skipping")
        return bw.get_item_password(item_name)

    url = redmine_cfg.get("url", REDMINE_URL)

    sso_login(page, f"{url}/login", agent, bw)

    # Navigate to account page where API key lives
    page.goto(f"{url}/my/account", wait_until="networkidle")

    # The API key is in the right sidebar under "API access key"
    # Click "Show" to reveal it
    show_link = page.query_selector('a:has-text("Show"), #api_access_key + a, a[href*="access_key"]')
    if show_link:
        show_link.click()
        page.wait_for_timeout(1000)

    key_el = page.query_selector('#api_access_key, .api-key code, .api-access-key')
    api_key = key_el.text_content().strip() if key_el else ""

    if not api_key:
        # If there's no existing key, try to reset/generate
        reset_link = page.query_selector('a:has-text("Reset"), a:has-text("Generate")')
        if reset_link:
            reset_link.click()
            page.wait_for_timeout(2000)
            page.click('button:has-text("OK"), button:has-text("Confirm")')
            page.wait_for_timeout(1000)
            key_el = page.query_selector('#api_access_key, .api-key code')
            api_key = key_el.text_content().strip() if key_el else ""

    if not api_key:
        log.error(f"[{name}] Could not get Redmine API key")
        return ""

    log.info(f"[{name}] Redmine API key obtained: {api_key[:8]}...")

    # Store in Bitwarden
    bw.create_item(
        name=item_name,
        username=name,
        password=api_key,
        uris=[url],
        collection_name=name,
    )
    log.info(f"[{name}] Redmine API key stored in Bitwarden")

    return api_key


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------


def verify_gitea(token: str, agent: dict) -> bool:
    """Verify the Gitea API token works."""
    import urllib.request

    url = agent.get("systems", {}).get("gitea", {}).get("url", GITEA_URL)
    req = urllib.request.Request(f"{url}/api/v1/user", headers={"Authorization": f"token {token}"})
    try:
        resp = urllib.request.urlopen(req, timeout=10)
        data = json.loads(resp.read())
        log.info(f"  Gitea verify: user={data.get('login', '?')}")
        return resp.status == 200
    except Exception as e:
        log.error(f"  Gitea verify FAILED: {e}")
        return False


def verify_redmine(key: str, agent: dict) -> bool:
    """Verify the Redmine API key works."""
    import urllib.request

    url = agent.get("systems", {}).get("redmine", {}).get("url", REDMINE_URL)
    req = urllib.request.Request(f"{url}/users/current.json", headers={"X-Redmine-API-Key": key})
    try:
        resp = urllib.request.urlopen(req, timeout=10)
        data = json.loads(resp.read())
        log.info(f"  Redmine verify: user={data.get('user', {}).get('login', '?')}")
        return resp.status == 200
    except Exception as e:
        log.error(f"  Redmine verify FAILED: {e}")
        return False


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def load_manifest(path: str) -> list[dict]:
    """Load and validate the agent manifest."""
    with open(path) as f:
        data = yaml.safe_load(f)

    agents = data.get("agents", [])
    if not agents:
        log.error("No agents found in manifest")
        sys.exit(1)

    for agent in agents:
        if "cloudron_invite" not in agent or "REPLACE" in agent["cloudron_invite"]:
            log.warning(f"Agent {agent.get('name', '?')} has no valid Cloudron invite")

    return agents


def provision_agent(
    context: BrowserContext,
    agent: dict,
    bw: BitwardenHelper,
    phase1_only: bool = False,
) -> dict:
    """Provision a single agent identity."""
    name = agent["name"]
    results = {"name": name, "cloudron": False, "gitea": False, "discourse": False, "redmine": False}

    page = context.new_page()

    try:
        # Phase 1: Cloudron enrollment
        enroll_cloudron(page, agent, bw)
        results["cloudron"] = True
        log.info(f"[{name}] Phase 1 complete: Cloudron identity enrolled")

        if phase1_only:
            log.info(f"[{name}] Phase 1 only — skipping system access")
            return results

        systems = agent.get("systems", {})
        if not systems:
            log.info(f"[{name}] No systems configured — Phase 1 only")
            return results

        # Phase 2: System access (fresh page for each system to avoid SSO conflicts)
        for system_name in ["gitea", "discourse", "redmine"]:
            system_page = context.new_page()
            try:
                if system_name == "gitea":
                    token = provision_gitea(system_page, agent, bw)
                    results["gitea"] = bool(token) and verify_gitea(token, agent)
                elif system_name == "discourse":
                    key = provision_discourse(system_page, agent, bw)
                    results["discourse"] = bool(key)
                elif system_name == "redmine":
                    key = provision_redmine(system_page, agent, bw)
                    results["redmine"] = bool(key) and verify_redmine(key, agent)
            except Exception as e:
                log.error(f"[{name}] {system_name} provisioning failed: {e}")
            finally:
                system_page.close()

        log.info(f"[{name}] All phases complete: {results}")

    except Exception as e:
        log.error(f"[{name}] Provisioning failed: {e}")
        results["error"] = str(e)
    finally:
        page.close()

    # Save state (in finally so partial results survive failures)
    state_file = STATE_DIR / f"{name}.json"
    try:
        with open(state_file, "w") as f:
            json.dump(results, f, indent=2)
    except OSError:
        log.warning(f"[{name}] Could not write state file: {state_file}")

    return results


def main():
    parser = argparse.ArgumentParser(description="Provision AI agent identities")
    parser.add_argument("--manifest", default="agents.yaml", help="Path to manifest file")
    parser.add_argument("--agent", help="Provision only this agent")
    parser.add_argument("--phase1-only", action="store_true", help="Cloudron enrollment only")
    parser.add_argument("--dry-run", action="store_true", help="Validate manifest without browser")
    parser.add_argument("--headed", action="store_true", help="Show browser (debugging)")
    args = parser.parse_args()

    # Ensure state directory exists (deferred from module level so --dry-run works)
    STATE_DIR.mkdir(parents=True, exist_ok=True)

    # Load manifest
    agents = load_manifest(args.manifest)
    if args.agent:
        agents = [a for a in agents if a["name"] == args.agent]
        if not agents:
            log.error(f"Agent '{args.agent}' not found in manifest")
            sys.exit(1)

    log.info(f"Manifest: {len(agents)} agent(s) to provision")
    for a in agents:
        log.info(f"  - {a['name']} ({a.get('display_name', '?')}) [{a.get('priority', '?')}]")

    if args.dry_run:
        log.info("Dry run — manifest validated successfully")
        return

    # Initialize Bitwarden
    bw = BitwardenHelper(
        client_id=os.environ["BW_CLIENTID"],
        client_secret=os.environ["BW_CLIENTSECRET"],
        password=os.environ["BW_PASSWORD"],
    )
    log.info("Connecting to Bitwarden...")
    bw.login()
    log.info("Bitwarden session established")

    # Launch Playwright
    headless = not args.headed
    if os.environ.get("HEADFUL", "false").lower() == "true":
        headless = False

    all_results = []

    with sync_playwright() as pw:
        browser = pw.chromium.launch(headless=headless)

        for agent in agents:
            # Fresh context per agent (no cookie/session bleed)
            context = browser.new_context(
                accept_downloads=False,
                java_script_enabled=True,
            )
            log.info(f"=== Provisioning: {agent['name']} ===")
            try:
                result = provision_agent(context, agent, bw, args.phase1_only)
                all_results.append(result)
            except Exception as e:
                log.error(f"FAILED: {agent['name']}: {e}")
                all_results.append({"name": agent["name"], "error": str(e)})
            finally:
                context.close()

        browser.close()

    # Summary
    log.info("\n=== PROVISIONING SUMMARY ===")
    for r in all_results:
        status = "OK" if "error" not in r else "FAILED"
        systems = []
        for s in ["cloudron", "gitea", "discourse", "redmine"]:
            if r.get(s):
                systems.append(s)
        log.info(f"  {r['name']}: {status} — {', '.join(systems) if systems else '(none)'}")

    # Exit non-zero if any failed
    if any("error" in r for r in all_results):
        sys.exit(1)


if __name__ == "__main__":
    main()