#!/usr/bin/env python3
"""
Simple Google Docs Downloader

Downloads a Google Docs document in both .txt and .md formats.

Usage:
    python download_gdoc.py <google_docs_url>
"""

import re
import sys
import requests


def extract_doc_id(url):
    """Extract the document ID from a Google Docs URL."""
    patterns = [
        r'/document/d/([a-zA-Z0-9-_]+)',
        r'id=([a-zA-Z0-9-_]+)',
    ]

    for pattern in patterns:
        match = re.search(pattern, url)
        if match:
            return match.group(1)

    raise ValueError("Could not extract document ID from URL")


def extract_code_sections(content, output_dir="."):
    """
    Extract code blocks from a document and save them to separate files.

    The document should have sections marked with # and code blocks in ```py ... ```
    """

    # Pattern to find section headers (# filename)
    # and their corresponding code blocks (```py ... ```)

    # Split by section headers
    sections = re.split(r'\n# (.+?\.py)\n', content)

    # Process sections (skip first element if empty)
    for i in range(1, len(sections), 2):
        if i + 1 < len(sections):
            filename = sections[i].strip()
            section_content = sections[i + 1]

            # Extract code from ```py ... ``` blocks
            code_blocks = re.findall(r'```py\n(.*?)```', section_content, re.DOTALL)

            if code_blocks:
                # Combine all code blocks for this file
                code = '\n'.join(code_blocks)

                # Write to file
                output_path = f"{output_dir}/{filename}"
                with open(output_path, 'w', encoding='utf-8') as out_file:
                    out_file.write(code)

                print(f"Created: {output_path}")


def download_google_doc(doc_id, format='txt'):
    """
    Download Google Doc in specified format.

    Args:
        doc_id: The Google Docs document ID
        format: 'txt', 'md', 'html', 'pdf', 'docx', etc.

    Returns:
        Document content as string (for text formats) or bytes (for binary formats)
    """
    export_url = f"https://docs.google.com/document/d/{doc_id}/export?format={format}"

    print(f"Downloading as {format.upper()}...")

    response = requests.get(export_url)
    response.raise_for_status()

    if format in ['txt', 'md', 'html']:
        return response.text
    else:
        return response.content


def main():
    if len(sys.argv) < 3:
        print("Usage: python download_gdoc.py <google_docs_url> <output_directory>")
        print("\nExample:")
        print('  python download_gdoc.py "https://docs.google.com/document/d/YOUR_DOC_ID/edit" code')
        sys.exit(1)

    url = sys.argv[1]

    try:
        # Extract document ID
        doc_id = extract_doc_id(url)
        print(f"Document ID: {doc_id}\n")

        # Download as TXT
        txt_content = download_google_doc(doc_id, 'md')
        # with open('document.txt', 'w', encoding='utf-8') as f:
        #    f.write(txt_content)

        extract_code_sections(txt_content, sys.argv[2])

        # Download as MD

        # print(f"✓ Saved: document.txt ({len(txt_content)} characters)")
        # print("✓ Download complete!")
        print("Done!")

    except requests.exceptions.HTTPError as e:
        print(f"\n❌ Error downloading document: {e}")
        print("Make sure the document is publicly accessible (shared with 'Anyone with the link')")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
