main.py

import os
import json
import re
import time
import random
from playwright.sync_api import sync_playwright
from markdownify import MarkdownConverter

# If this list is NOT empty, the script will ONLY extract conversations
# whose titles contain at least one of these words (case-insensitive).
TITLE_WHITELIST = []

CLONE_AND_CLEAN_JS = """
(el) => {
    const clone = el.cloneNode(true);

    // 1. Purge all the garbage UI we don't want polluting the markdown/text
    const selectorsToRemove = [
        '.cdk-visually-hidden',
        '.screen-reader-user-query-label',
        '.screen-reader-model-response-label',
        'message-actions',
        'elicitations',
        'follow-up',
        '.response-container-header',
        'mat-progress-bar',
        '.mini-app-placeholder',
        '.floating-action-buttons',
        '.model-response-label-announcer',
        '.hide-from-message-actions',
        'source-footnote'
    ];

    clone.querySelectorAll(selectorsToRemove.join(', ')).forEach(node => node.remove());

    // 2. Off-screen render trick. Browsers strip line-breaks from innerText
    // if the node isn't actively part of the document layout tree.
    const hiddenDiv = document.createElement('div');
    hiddenDiv.style.position = 'absolute';
    hiddenDiv.style.left = '-9999px';
    hiddenDiv.appendChild(clone);
    document.body.appendChild(hiddenDiv);

    const result = {
        innerText: clone.innerText.trim(),
        innerHTML: clone.innerHTML,
        outerHTML: clone.outerHTML
    };

    // Cleanup the live DOM
    document.body.removeChild(hiddenDiv);

    return result;
}
"""

class GeminiMarkdownConverter(MarkdownConverter):
    """
    Custom markdownify converter to handle Gemini's specific DOM structures,
    like their custom web components and data attributes.
    """
    def convert_code(self, el, text, parent_tags):
        # Gemini uses <code data-test-id="code-content"> for code blocks
        if el.get('data-test-id') == 'code-content':
            return f"\n```\n{text}\n```\n"
        return super().convert_code(el, text, parent_tags)

def md_gemini(html):
    return GeminiMarkdownConverter(heading_style="ATX").convert(html)

def main():
    os.makedirs("conversations", exist_ok=True)

    with sync_playwright() as p:
        print("Connecting to live Chrome session over CDP (port 9222)...")
        try:
            browser = p.chromium.connect_over_cdp("http://localhost:9222")
        except Exception as e:
            print("Failed to connect. Is Chrome running with --remote-debugging-port=9222?")
            print(f"Error: {e}")
            return

        # Find the Gemini tab
        gemini_page = None
        for context in browser.contexts:
            for page in context.pages:
                if "gemini.google.com" in page.url:
                    gemini_page = page
                    break
            if gemini_page:
                break

        if not gemini_page:
            print("Could not find a Gemini tab. Please open gemini.google.com in your browser.")
            return

        print(f"Attached to active session.")

        # 1. Sidebar Infinite Scroll & Extraction
        print("Locating sidebar to load all historical chats...")
        sidebar_scroller = gemini_page.locator('bard-sidenav infinite-scroller')
        links_locator = gemini_page.locator('[data-test-id="chats-expandable-section"] gem-nav-list-item a')

        try:
            sidebar_scroller.wait_for(state="visible", timeout=5000)
            print("Scrolling sidebar to the bottom...")
            while True:
                old_height = sidebar_scroller.evaluate("el => el.scrollHeight")

                if links_locator.count() > 0:
                    try:
                        links_locator.last.scroll_into_view_if_needed()
                    except Exception:
                        pass # Ignore if Angular unmounts it mid-scroll

                # Dumb but reliable wait for Angular to trigger the load and inject elements
                time.sleep(2.0)

                new_height = sidebar_scroller.evaluate("el => el.scrollHeight")
                if new_height <= old_height:
                    break
        except Exception as e:
            print(f"Could not find or scroll the sidebar. Ensure the left menu is open. {e}")

        print("Extracting chat links...")
        chat_links = []

        for i in range(links_locator.count()):
            link = links_locator.nth(i)
            href = link.get_attribute("href")
            title = link.inner_text().strip()

            # Fallback if text is visually hidden
            if not title:
                title = link.get_attribute("aria-label") or f"Untitled_Chat_{i}"
                print(f"  -> Chat link {i+1} has no visible text. Using fallback title: '{title}'.")

            chat_links.append({
                "locator": link,
                "href": href,
                "title": title
            })

        print(f"Found {len(chat_links)} total conversations in the sidebar.")

        # 2. Master Extraction Loop
        for index, chat in enumerate(chat_links):
            raw_title = chat["title"]

            # Whitelist filter check
            if TITLE_WHITELIST:
                if not any(term.lower() in raw_title.lower() for term in TITLE_WHITELIST):
                    print(f"[{index+1}/{len(chat_links)}] Skipping '{raw_title}' (not in whitelist).")
                    continue

            safe_title = re.sub(r'[^a-zA-Z0-9_\-\s]', '', raw_title).strip().replace(' ', '_')
            convo_dir = os.path.join("conversations", safe_title)

            # Idempotency check
            if os.path.exists(convo_dir) and os.path.exists(os.path.join(convo_dir, "chat.json")):
                print(f"[{index+1}/{len(chat_links)}] Skipping '{raw_title}' (already saved).")
                continue

            print(f"\n[{index+1}/{len(chat_links)}] Loading conversation: '{raw_title}'...")

            # Navigation
            target_href = chat["href"]
            is_new_page = target_href and not gemini_page.url.endswith(target_href)

            # Snapshot the page title *before* we click
            old_page_title = gemini_page.title()

            chat["locator"].click()

            if is_new_page:
                try:
                    gemini_page.wait_for_url(f"**{target_href}", timeout=5000)
                except Exception:
                    pass

                # Wait for the document title to change (Angular routing complete)
                try:
                    gemini_page.wait_for_function(f"document.title !== {json.dumps(old_page_title)}", timeout=5000)
                except Exception:
                    print("  -> Warning: Document title didn't change within 5s timeout. Falling back to sleep.")
                    pass

            # Human-like delay to let DOM components finalize rendering after the title swap
            time.sleep(random.uniform(1.0, 2.0))

            # CANVAS MODE CHECK
            # If the conversation loaded into Canvas mode, it collapses the sidebar. Let's fix that before continuing.
            close_panel_btn = gemini_page.locator('gem-icon-button[aria-label="Close panel"], gem-icon-button[arialabel="Close panel"]')
            if close_panel_btn.count() > 0 and close_panel_btn.first.is_visible():
                print("  -> Canvas mode detected. Closing panel and restoring sidebar...")
                try:
                    close_panel_btn.first.click()
                    time.sleep(2.0)

                    sparkle_btn = gemini_page.locator('side-nav-sparkle-button')
                    if sparkle_btn.count() > 0 and sparkle_btn.first.is_visible():
                        print("  -> Hovering sparkle button to reveal sidebar toggle...")
                        sparkle_btn.first.hover()
                        time.sleep(2.0)

                        menu_toggle = gemini_page.locator('gem-icon-button[aria-label*="menu" i], gem-icon-button[aria-label*="sidebar" i]')

                        if menu_toggle.count() > 0 and menu_toggle.first.is_visible():
                            print("  -> Clicking revealed menu toggle...")
                            menu_toggle.first.click()
                        else:
                            print("  -> Clicking sparkle button directly...")
                            sparkle_btn.first.click()

                        gemini_page.locator('bard-sidenav').wait_for(state="visible", timeout=5000)
                        time.sleep(0.5)
                    else:
                        print("  -> Sparkle button not found. Sidebar may not be restored correctly.")
                except Exception as e:
                    print(f"  -> Warning: Issue handling Canvas mode UI: {e}")

            chat_container = gemini_page.locator('infinite-scroller[data-test-id="chat-history-container"]')

            json_messages = []
            md_lines = []
            html_blocks = []
            total_extracted = 0

            print(f"  -> Scrolling up to extract history...")

            # 3. Chat Infinite Scroll & Extraction
            while True:
                untagged_locator = chat_container.locator('user-query:not([data-exporter-extracted="true"]), model-response:not([data-exporter-extracted="true"])')

                # Snapshot the DOM nodes so we don't break the locator when we mutate their attributes
                handles = untagged_locator.element_handles()

                if handles:
                    chunk_json = []
                    chunk_md = []
                    chunk_html = []

                    for handle in handles:
                        tag_name = handle.evaluate("el => el.tagName.toLowerCase()")
                        role = "user" if tag_name == "user-query" else "model"

                        extracted_data = handle.evaluate(CLONE_AND_CLEAN_JS)

                        text_content = extracted_data['innerText']
                        raw_html = extracted_data['innerHTML']
                        raw_outer_html = extracted_data['outerHTML']

                        markdown_content = md_gemini(raw_html).strip()

                        chunk_json.append({"role": role, "text": text_content})
                        header = "# User" if role == "user" else "# Gemini"
                        chunk_md.append(f"{header}\n\n{markdown_content}")
                        chunk_html.append(f'<div class="message-wrapper">\n{raw_outer_html}\n</div>')

                        # Mark as extracted so it's ignored next loop
                        handle.evaluate("el => el.setAttribute('data-exporter-extracted', 'true')")

                    # Prepend the chunks since we are moving backwards in time (upwards)
                    json_messages = chunk_json + json_messages
                    md_lines = chunk_md + md_lines
                    html_blocks = chunk_html + html_blocks

                    total_extracted += len(handles)
                    print(f"  -> Extracted {len(handles)} messages (Total so far: {total_extracted}). Scrolling higher...")

                # Scroll up using Option B (Scroll to the absolute topmost message in the DOM)
                old_height = chat_container.evaluate("el => el.scrollHeight")

                all_messages = chat_container.locator('user-query, model-response')
                if all_messages.count() > 0:
                    all_messages.first.scroll_into_view_if_needed()

                # Dumb but reliable wait for Angular to trigger the load and inject elements
                time.sleep(2.0)

                new_height = chat_container.evaluate("el => el.scrollHeight")
                has_more = new_height > old_height

                if not has_more:
                    # Final safety check in case the last tiny scroll revealed elements without expanding scrollHeight
                    if chat_container.locator('user-query:not([data-exporter-extracted="true"]), model-response:not([data-exporter-extracted="true"])').count() > 0:
                        continue
                    else:
                        break

            print(f"  -> Reached the top! Saving {total_extracted} messages.")
            # 4. Write Output Files
            if total_extracted == 0:
                print(f"  -> No messages found. Skipping save.")
                continue

            os.makedirs(convo_dir, exist_ok=True)

            json_file = os.path.join(convo_dir, "chat.json")
            md_file = os.path.join(convo_dir, "chat.md")
            html_file = os.path.join(convo_dir, "chat.html")

            with open(json_file, 'w', encoding='utf-8') as f:
                json.dump({
                    "title": raw_title,
                    "message_count": total_extracted,
                    "messages": json_messages
                }, f, indent=2, ensure_ascii=False)

            with open(md_file, 'w', encoding='utf-8') as f:
                f.write("\n\n---\n\n".join(md_lines))

            html_template = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>{raw_title}</title>
    <style>
        body {{
            background: #f0f4f9;
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
            padding: 2rem;
            max-width: 900px;
            margin: 0 auto;
            line-height: 1.6;
            color: #1f1f1f;
        }}
        .message-wrapper {{
            margin-bottom: 1.5rem;
            clear: both;
            overflow: hidden;
            width: 100%;
        }}
        user-query {{
            display: block;
            background: #e3e3e3;
            padding: 1.25rem 1.5rem 0.25rem;
            border-radius: 24px;
            border-bottom-right-radius: 4px;
            max-width: 80%;
            float: right;
            font-size: 1rem;
        }}
        model-response {{
            display: block;
            background: #ffffff;
            padding: 1.25rem 1.5rem 0.25rem;
            border-radius: 24px;
            border-top-left-radius: 4px;
            max-width: 85%;
            float: left;
            box-shadow: 0 2px 6px rgba(0,0,0,0.02);
            font-size: 1rem;
        }}
        .cdk-visually-hidden,
        .screen-reader-user-query-label,
        .screen-reader-model-response-label,
        .luminous-toggle-container,
        user-query-file-carousel,
        .attachment-container,
        gem-icon-button,
        .luminous-actions-container,
        message-actions,
        elicitations,
        follow-up,
        .response-container-header,
        mat-progress-bar,
        .mini-app-placeholder,
        .floating-action-buttons,
        .model-response-label-announcer,
        .code-block-decoration,
        .hide-from-message-actions,
        source-footnote {{
            display: none !important;
        }}
        pre {{
            background: #1e1e1e;
            color: #d4d4d4;
            padding: 1rem;
            border-radius: 8px;
            overflow-x: auto;
            margin: 1rem 0;
        }}
        code {{
            font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, "Liberation Mono", monospace;
            font-size: 0.9em;
        }}
        p code {{
            background: #f1f3f4;
            color: #1f1f1f;
            padding: 0.2rem 0.4rem;
            border-radius: 4px;
        }}
        a {{ color: #1a73e8; text-decoration: none; }}
        a:hover {{ text-decoration: underline; }}
        p {{ margin-bottom: 1rem; margin-top: 0; }}
        h1, h2, h3, h4 {{ margin-top: 1.5rem; margin-bottom: 0.75rem; }}
        ul, ol {{ margin-bottom: 1rem; padding-left: 2rem; }}
        li {{ margin-bottom: 0.5rem; }}
        li p, ol p, ul p {{ margin-bottom: 0rem; }}
    </style>
</head>
<body>
{''.join(html_blocks)}
</body>
</html>"""
            with open(html_file, 'w', encoding='utf-8') as f:
                f.write(html_template)

if __name__ == "__main__":
    main()