"""
client.py — Browser Automation & Payload Execution Layer
=========================================================
Uses Playwright (sync API) to drive headless Chromium against https://os.solidpoint.ai.
Handles login, chat-message injection, auto-continue monitoring, and response extraction.

⚠️  NO print() calls anywhere — all output goes to logging (stderr / file).
"""

import csv
import json
import logging
import os
import sys
import time
import tempfile
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# Logging — stderr + file, NEVER stdout
# ---------------------------------------------------------------------------
LOG_FILE = Path(__file__).parent.parent / "logs" / "mcp_manager.log"

_fmt = logging.Formatter(
    "[%(asctime)s] %(levelname)-7s %(name)s — %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

_sh = logging.StreamHandler(sys.stderr)
_sh.setFormatter(_fmt)

_fh = logging.FileHandler(LOG_FILE, encoding="utf-8")
_fh.setFormatter(_fmt)

log = logging.getLogger("redteam.client")
log.setLevel(logging.DEBUG)
log.addHandler(_sh)
log.addHandler(_fh)

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).parent.parent
ENV_FILE = BASE_DIR / "config" / ".env"
ACCOUNTS_CSV = BASE_DIR / "data" / "accounts.csv"
TESTCASES_JSON = BASE_DIR / "data" / "redteam_testcases.json"
EVIDENCE_DIR = BASE_DIR / "evidence_screenshots"
EVIDENCE_DIR.mkdir(exist_ok=True)


# ---------------------------------------------------------------------------
# Configuration loader
# ---------------------------------------------------------------------------
def load_env() -> dict:
    """Parse .enc file into a dict (simple KEY=VALUE format)."""
    env = {}
    if not ENV_FILE.exists():
        log.warning(".env file not found at %s", ENV_FILE)
        return env
    for line in ENV_FILE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip()
    return env


def load_accounts() -> list[dict]:
    """Load test accounts from accounts.csv, filtering only SUCCESS rows."""
    accounts = []
    if not ACCOUNTS_CSV.exists():
        log.warning("accounts.csv not found at %s", ACCOUNTS_CSV)
        return accounts
    with open(ACCOUNTS_CSV, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row.get("Status", "").strip().upper() == "SUCCESS":
                accounts.append({
                    "email": row["Email"].strip().strip('"'),
                    "password": row["Password"].strip().strip('"'),
                })
    log.info("Loaded %d valid accounts from CSV", len(accounts))
    return accounts


def load_testcases() -> list[dict]:
    """Load red-team test cases from JSON."""
    if not TESTCASES_JSON.exists():
        log.error("redteam_testcases.json not found at %s", TESTCASES_JSON)
        return []
    with open(TESTCASES_JSON, encoding="utf-8") as f:
        cases = json.load(f)
    log.info("Loaded %d test cases", len(cases))
    return cases


# ---------------------------------------------------------------------------
# Browser helpers
# ---------------------------------------------------------------------------
# Selectors observed from the os.solidpoint.ai UI (Preact SPA)
SEL_EMAIL = '#login-email, input[type="email"]'
SEL_PASS = '#login-pass, input[type="password"]'
SEL_LOGIN_BTN = ".login-btn"
SEL_SIGNUP_TOGGLE = ".login-toggle button"
SEL_LOGIN_ERROR = ".login-error"
SEL_TEXTAREA = ".input-row textarea"
SEL_SEND_BTN = ".send-btn"
SEL_MSG_BODY = ".m-body"
SEL_TYPING = ".typing"
SEL_WELCOME = ".welcome"
SEL_SIDEBAR = ".sidebar"
SEL_CONTINUE_BTN = 'button:has-text("Continue"), button:has-text("continue")'
SEL_UPGRADE_CLOSE = ".upgrade-cancel, .upgrade-overlay"


def _wait_for_app(page, timeout: int = 15_000):
    """Wait until the SPA renders (login or chat view)."""
    page.wait_for_load_state("networkidle", timeout=timeout)
    # Either login-box or sidebar/welcome must appear
    page.wait_for_selector(
        f"{SEL_EMAIL}, {SEL_SIDEBAR}, {SEL_WELCOME}",
        timeout=timeout,
    )


def login_with_credentials(page, email: str, password: str, is_signup: bool = False) -> bool:
    """
    Log in (or sign up) via the UI.
    Returns True on success, False on failure.
    """
    log.info("Attempting %s for %s", "signup" if is_signup else "signin", email)

    page.goto("https://os.solidpoint.ai", wait_until="networkidle", timeout=30_000)
    _wait_for_app(page)

    # Switch to signup mode if needed
    if is_signup:
        toggle = page.locator(SEL_SIGNUP_TOGGLE)
        if toggle.is_visible():
            txt = (toggle.text_content() or "").lower()
            if "sign up" in txt:
                toggle.click()
                page.wait_for_timeout(500)

    # Fill credentials
    page.locator(SEL_EMAIL).first.fill(email)
    page.locator(SEL_PASS).first.fill(password)
    page.locator(SEL_LOGIN_BTN).first.click()

    # Wait for transition
    page.wait_for_timeout(4_000)

    # Check for errors
    err_el = page.locator(SEL_LOGIN_ERROR)
    if err_el.is_visible():
        err_text = err_el.text_content() or ""
        log.warning("Login error for %s: %s", email, err_text.strip())
        return False

    # Success indicators: sidebar, welcome screen, or chat area visible
    success = (
        page.locator(SEL_SIDEBAR).is_visible()
        or page.locator(SEL_WELCOME).is_visible()
        or page.locator(SEL_TEXTAREA).is_visible()
    )
    if success:
        log.info("Login succeeded for %s", email)
    else:
        log.warning("Login outcome unclear for %s (no sidebar/welcome)", email)

    return success


def login_with_token(page, token: str) -> bool:
    """
    Inject PRO_CHAT_TOKEN into localStorage and reload to authenticate.
    """
    log.info("Injecting token via localStorage…")
    page.goto("https://os.solidpoint.ai", wait_until="networkidle", timeout=30_000)
    page.evaluate(f"localStorage.setItem('chat_token', '{token}')")
    page.reload(wait_until="networkidle", timeout=30_000)
    _wait_for_app(page)

    success = (
        page.locator(SEL_SIDEBAR).is_visible()
        or page.locator(SEL_WELCOME).is_visible()
        or page.locator(SEL_TEXTAREA).is_visible()
    )
    if success:
        log.info("Token-based auth succeeded")
    else:
        log.warning("Token-based auth outcome unclear")
    return success


def _dismiss_upgrade_modal(page):
    """Dismiss the upgrade modal if it appears."""
    cancel = page.locator(".upgrade-cancel")
    if cancel.is_visible():
        cancel.click()
        page.wait_for_timeout(300)


def _create_new_chat(page):
    """Click the new-chat button in sidebar."""
    new_btn = page.locator(".sb-new").first
    if new_btn.is_visible():
        new_btn.click()
        page.wait_for_timeout(1_000)
        log.debug("Created new chat session")


def _upload_file_for_testcase(page, file_spec: dict) -> Optional[str]:
    """
    For indirect-injection test cases that include a file payload,
    create a temp file and attach it via the file input.
    Returns the temp file path or None.
    """
    if not file_spec:
        return None

    fname = file_spec.get("name", "payload.txt")
    content = file_spec.get("content", "")

    tmp_dir = Path(tempfile.mkdtemp(prefix="redteam_"))
    tmp_path = tmp_dir / fname
    tmp_path.write_text(content, encoding="utf-8")
    log.debug("Created temp payload file: %s", tmp_path)

    # Try to locate file input (hidden) — the attach button triggers it
    file_input = page.locator('input[type="file"]')
    if file_input.count() > 0:
        file_input.first.set_input_files(str(tmp_path))
        page.wait_for_timeout(1_000)
        log.debug("Attached file via input[type=file]")
    else:
        # Fallback: inject file content as a paste-style attachment in the message
        log.debug("No file input found; prepending file content to message")

    return str(tmp_path)


def _wait_for_response_complete(page, timeout_sec: int = 120) -> str:
    """
    Monitor the chat for response completion.
    Handles auto-continue clicks and waits for typing indicator to disappear.
    Returns the extracted final assistant response text.
    """
    log.debug("Waiting for response (timeout=%ds)…", timeout_sec)
    start = time.time()
    last_content = ""
    stable_count = 0

    while time.time() - start < timeout_sec:
        # Dismiss upgrade modals
        _dismiss_upgrade_modal(page)

        # Auto-continue: click "Continue" if it appears
        continue_btn = page.locator(SEL_CONTINUE_BTN)
        if continue_btn.first.is_visible():
            log.debug("Clicking 'Continue' button…")
            continue_btn.first.click()
            page.wait_for_timeout(1_000)
            stable_count = 0
            continue

        # Check if typing indicator is present
        typing_visible = page.locator(SEL_TYPING).is_visible()

        # Check if send button is in "busy" (stop) state
        send_busy = page.locator(".send-btn.busy").is_visible()

        if typing_visible or send_busy:
            stable_count = 0
            page.wait_for_timeout(500)
            continue

        # Extract latest assistant message
        msgs = page.locator(f".m-ast {SEL_MSG_BODY}")
        count = msgs.count()
        if count > 0:
            current = msgs.nth(count - 1).inner_text()
            if current == last_content and current.strip():
                stable_count += 1
            else:
                last_content = current
                stable_count = 0

            # Stable for ~2 seconds → response complete
            if stable_count >= 4:
                log.debug("Response stable — extracting")
                return current.strip()

        page.wait_for_timeout(500)

    log.warning("Response extraction timed out after %ds", timeout_sec)
    # Return whatever we have
    msgs = page.locator(f".m-ast {SEL_MSG_BODY}")
    if msgs.count() > 0:
        return msgs.nth(msgs.count() - 1).inner_text().strip()
    return "[TIMEOUT — no response extracted]"


def execute_payload(page, testcase: dict, timeout_sec: int = 120) -> dict:
    """
    Send a single red-team payload into the active chat and extract the response.

    Args:
        page: Playwright page with an authenticated session.
        testcase: Dict with keys: id, cat, title, message, file.
        timeout_sec: Max seconds to wait for response.

    Returns:
        Dict with: id, cat, title, payload, response, elapsed_sec, has_file.
    """
    tc_id = testcase["id"]
    message = testcase["message"]
    file_spec = testcase.get("file")

    log.info("▶ Executing %s [%s]: %s", tc_id, testcase["cat"], testcase["title"])

    # Create a fresh chat for isolation
    _create_new_chat(page)
    _dismiss_upgrade_modal(page)

    # Handle file attachment for indirect injection tests
    tmp_file = None
    if file_spec:
        tmp_file = _upload_file_for_testcase(page, file_spec)
        # If no file input was available, prepend file content to message
        file_input = page.locator('input[type="file"]')
        if file_input.count() == 0:
            message = (
                f"[File: {file_spec['name']}]\n"
                f"```\n{file_spec['content']}\n```\n\n"
                f"{message}"
            )

    # Type the message
    textarea = page.locator(SEL_TEXTAREA).first
    textarea.click()
    textarea.fill(message)
    page.wait_for_timeout(300)

    # Send
    t0 = time.time()
    send_btn = page.locator(SEL_SEND_BTN).first
    send_btn.click()
    page.wait_for_timeout(1_500)

    # Wait for and extract response
    response_text = _wait_for_response_complete(page, timeout_sec)
    elapsed = round(time.time() - t0, 2)

    page.wait_for_timeout(1000)
    screenshot_name = f"{tc_id}_{testcase['cat']}_evidence.png"
    screenshot_path = EVIDENCE_DIR / screenshot_name
    page.screenshot(path=str(screenshot_path), full_page=True)
    log.info("Captured evidence screenshot: %s", screenshot_path)

    log.info("◀ %s completed in %.1fs (%d chars)", tc_id, elapsed, len(response_text))

    return {
        "id": tc_id,
        "cat": testcase["cat"],
        "title": testcase["title"],
        "payload": testcase["message"],
        "response": response_text,
        "elapsed_sec": elapsed,
        "has_file": file_spec is not None,
        "screenshot_path": str(screenshot_path),
    }


# ---------------------------------------------------------------------------
# High-level orchestration (called by server.py)
# ---------------------------------------------------------------------------
def run_full_suite(
    accounts: list[dict],
    testcases: list[dict],
    env: dict,
    headless: bool = True,
    timeout_per_test: int = 120,
) -> list[dict]:
    """
    Execute all test cases across available accounts.
    Uses Pro account (token injection) if available, falls back to CSV accounts.

    Returns list of result dicts.
    """
    from playwright.sync_api import sync_playwright

    results = []
    target_url = env.get("TARGET_URL", "https://os.solidpoint.ai")
    pro_token = env.get("PRO_CHAT_TOKEN", "")
    pro_email = env.get("PRO_EMAIL", "")
    pro_password = env.get("PRO_PASSWORD", "")

    log.info("=" * 60)
    log.info("STARTING RED-TEAM SECURITY SUITE")
    log.info("Target: %s | Tests: %d | Accounts: %d", target_url, len(testcases), len(accounts))
    log.info("=" * 60)

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

        # Determine which accounts to cycle through
        auth_methods = []

        # Priority 1: Pro token
        if pro_token:
            auth_methods.append({"type": "token", "token": pro_token, "label": "PRO (token)"})

        # Priority 2: Pro credentials
        if pro_email and pro_password:
            auth_methods.append({"type": "creds", "email": pro_email, "password": pro_password, "label": f"PRO ({pro_email})"})

        # Priority 3: CSV accounts
        for acc in accounts:
            auth_methods.append({"type": "creds", "email": acc["email"], "password": acc["password"], "label": acc["email"]})

        if not auth_methods:
            log.error("No authentication methods available!")
            return results

        # Round-robin through accounts for test distribution
        account_idx = 0

        for i, tc in enumerate(testcases):
            auth = auth_methods[account_idx % len(auth_methods)]
            log.info(
                "--- Test %d/%d [%s] via %s ---",
                i + 1, len(testcases), tc["id"], auth["label"],
            )

            context = browser.new_context(
                viewport={"width": 1280, "height": 800},
                user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
            )
            page = context.new_page()

            try:
                # Authenticate
                logged_in = False
                if auth["type"] == "token":
                    logged_in = login_with_token(page, auth["token"])
                else:
                    logged_in = login_with_credentials(page, auth["email"], auth["password"])

                if not logged_in:
                    log.warning("Auth failed for %s — skipping test %s", auth["label"], tc["id"])
                    results.append({
                        "id": tc["id"],
                        "cat": tc["cat"],
                        "title": tc["title"],
                        "payload": tc["message"],
                        "response": "[AUTH_FAILED]",
                        "elapsed_sec": 0,
                        "has_file": tc.get("file") is not None,
                        "account": auth["label"],
                        "error": "Authentication failed",
                    })
                    # Try next account for subsequent tests
                    account_idx += 1
                    continue

                # Execute payload
                result = execute_payload(page, tc, timeout_per_test)
                result["account"] = auth["label"]
                results.append(result)

            except Exception as exc:
                log.exception("Exception during test %s: %s", tc["id"], exc)
                results.append({
                    "id": tc["id"],
                    "cat": tc["cat"],
                    "title": tc["title"],
                    "payload": tc["message"],
                    "response": f"[ERROR: {exc}]",
                    "elapsed_sec": 0,
                    "has_file": tc.get("file") is not None,
                    "account": auth["label"],
                    "error": str(exc),
                })
            finally:
                context.close()

            # Rotate account every N tests to distribute load
            if (i + 1) % 5 == 0:
                account_idx += 1

        browser.close()

    log.info("=" * 60)
    log.info("SUITE COMPLETE — %d results collected", len(results))
    log.info("=" * 60)

    return results