#!/usr/bin/env python3
"""
Notion Writer - Push markdown content to Notion pages or databases.
Supports both page content updates and database row creation.
"""

import os
import sys
import re
import argparse
from pathlib import Path
from typing import Optional, List, Dict, Any

from dotenv import load_dotenv
from notion_client import Client

# Load environment variables
load_dotenv(Path(__file__).parent / '.env')

NOTION_TOKEN = os.getenv('NOTION_API_KEY')


def extract_notion_id(url_or_id: str) -> Optional[str]:
    """Extract Notion page/database ID from URL or raw ID."""
    # Already a raw ID (32 chars hex, with or without dashes)
    clean_id = url_or_id.replace('-', '')
    if re.match(r'^[a-f0-9]{32}$', clean_id):
        return clean_id

    # Notion URL patterns
    patterns = [
        r'notion\.so/(?:[^/]+/)?([a-f0-9]{32})',  # notion.so/workspace/page-id
        r'notion\.so/(?:[^/]+/)?[^-]+-([a-f0-9]{32})',  # notion.so/workspace/Page-Title-id
        r'notion\.site/(?:[^/]+/)?([a-f0-9]{32})',  # public notion.site
        r'notion\.site/(?:[^/]+/)?[^-]+-([a-f0-9]{32})',
        r'([a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12})',  # UUID format
    ]

    for pattern in patterns:
        match = re.search(pattern, url_or_id)
        if match:
            return match.group(1).replace('-', '')

    return None


def format_id_with_dashes(raw_id: str) -> str:
    """Format 32-char ID to UUID format with dashes."""
    if len(raw_id) == 32:
        return f"{raw_id[:8]}-{raw_id[8:12]}-{raw_id[12:16]}-{raw_id[16:20]}-{raw_id[20:]}"
    return raw_id


def markdown_to_notion_blocks(markdown_text: str) -> List[Dict[str, Any]]:
    """Convert markdown text to Notion block objects."""
    blocks = []
    lines = markdown_text.split('\n')
    i = 0

    while i < len(lines):
        line = lines[i]

        # Skip empty lines
        if not line.strip():
            i += 1
            continue

        # Headers
        if line.startswith('######'):
            blocks.append(create_heading_block(line[6:].strip(), 3))
        elif line.startswith('#####'):
            blocks.append(create_heading_block(line[5:].strip(), 3))
        elif line.startswith('####'):
            blocks.append(create_heading_block(line[4:].strip(), 3))
        elif line.startswith('###'):
            blocks.append(create_heading_block(line[3:].strip(), 3))
        elif line.startswith('##'):
            blocks.append(create_heading_block(line[2:].strip(), 2))
        elif line.startswith('#'):
            blocks.append(create_heading_block(line[1:].strip(), 1))

        # Code blocks
        elif line.startswith('```'):
            language = line[3:].strip() or 'plain text'
            code_lines = []
            i += 1
            while i < len(lines) and not lines[i].startswith('```'):
                code_lines.append(lines[i])
                i += 1
            blocks.append(create_code_block('\n'.join(code_lines), language))

        # Bullet list
        elif line.strip().startswith('- ') or line.strip().startswith('* '):
            text = line.strip()[2:]
            blocks.append(create_bulleted_list_block(text))

        # Numbered list
        elif re.match(r'^\d+\.\s', line.strip()):
            text = re.sub(r'^\d+\.\s', '', line.strip())
            blocks.append(create_numbered_list_block(text))

        # Checkbox / Todo
        elif line.strip().startswith('- [ ]'):
            text = line.strip()[5:].strip()
            blocks.append(create_todo_block(text, False))
        elif line.strip().startswith('- [x]') or line.strip().startswith('- [X]'):
            text = line.strip()[5:].strip()
            blocks.append(create_todo_block(text, True))

        # Blockquote
        elif line.startswith('>'):
            text = line[1:].strip()
            blocks.append(create_quote_block(text))

        # Horizontal rule
        elif line.strip() in ['---', '***', '___']:
            blocks.append(create_divider_block())

        # Regular paragraph
        else:
            blocks.append(create_paragraph_block(line))

        i += 1

    return blocks


def parse_rich_text(text: str) -> List[Dict[str, Any]]:
    """Parse markdown inline formatting to Notion rich text."""
    rich_text = []

    # Simple implementation - just plain text for now
    # TODO: Add support for bold, italic, code, links
    if text:
        rich_text.append({
            "type": "text",
            "text": {"content": text[:2000]}  # Notion limit
        })

    return rich_text


def create_paragraph_block(text: str) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "paragraph",
        "paragraph": {
            "rich_text": parse_rich_text(text)
        }
    }


def create_heading_block(text: str, level: int) -> Dict[str, Any]:
    heading_type = f"heading_{level}"
    return {
        "object": "block",
        "type": heading_type,
        heading_type: {
            "rich_text": parse_rich_text(text)
        }
    }


def create_bulleted_list_block(text: str) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "bulleted_list_item",
        "bulleted_list_item": {
            "rich_text": parse_rich_text(text)
        }
    }


def create_numbered_list_block(text: str) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "numbered_list_item",
        "numbered_list_item": {
            "rich_text": parse_rich_text(text)
        }
    }


def create_todo_block(text: str, checked: bool) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "to_do",
        "to_do": {
            "rich_text": parse_rich_text(text),
            "checked": checked
        }
    }


def create_quote_block(text: str) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "quote",
        "quote": {
            "rich_text": parse_rich_text(text)
        }
    }


def create_code_block(code: str, language: str) -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "code",
        "code": {
            "rich_text": parse_rich_text(code),
            "language": language.lower()
        }
    }


def create_divider_block() -> Dict[str, Any]:
    return {
        "object": "block",
        "type": "divider",
        "divider": {}
    }


def list_accessible_content(notion: Client, filter_type: str = 'all') -> None:
    """List all accessible pages and databases."""
    print("=" * 70)
    print("Accessible Notion Content (Claude-D.intelligence)")
    print("=" * 70)

    databases = []
    pages = []

    try:
        # Search all content and filter locally
        response = notion.search()
        all_results = response.get('results', [])

        # Handle pagination
        while response.get('has_more'):
            response = notion.search(start_cursor=response.get('next_cursor'))
            all_results.extend(response.get('results', []))

        # Separate databases and pages
        for item in all_results:
            if item.get('object') == 'database':
                databases.append(item)
            elif item.get('object') == 'page':
                pages.append(item)

        # Show databases
        if filter_type in ['all', 'database']:
            print("\n📊 DATABASES")
            print("-" * 70)

            if databases:
                for i, db in enumerate(databases, 1):
                    title = ""
                    if db.get('title'):
                        title = db['title'][0].get('plain_text', 'Untitled') if db['title'] else 'Untitled'
                    db_id = db['id']
                    url = db.get('url', '')
                    props = list(db.get('properties', {}).keys())[:5]
                    print(f"\n{i}. {title}")
                    print(f"   ID: {db_id}")
                    print(f"   URL: {url}")
                    print(f"   Properties: {', '.join(props)}")
            else:
                print("   No databases accessible")

        # Show pages
        if filter_type in ['all', 'page']:
            print("\n\n📄 PAGES")
            print("-" * 70)

            if pages:
                for i, page in enumerate(pages, 1):
                    title = "Untitled"
                    if 'properties' in page:
                        for prop in page['properties'].values():
                            if prop.get('type') == 'title':
                                title_arr = prop.get('title', [])
                                if title_arr:
                                    title = title_arr[0].get('plain_text', 'Untitled')
                                    break
                    page_id = page['id']
                    url = page.get('url', '')
                    parent_type = page.get('parent', {}).get('type', 'unknown')
                    print(f"\n{i}. {title}")
                    print(f"   ID: {page_id}")
                    print(f"   URL: {url}")
                    print(f"   Parent: {parent_type}")
            else:
                print("   No pages accessible")

        print("\n" + "=" * 70)
        db_count = len(databases) if filter_type != 'page' else 0
        page_count = len(pages) if filter_type != 'database' else 0
        print(f"Total: {db_count} databases, {page_count} pages")
        print("=" * 70)

    except Exception as e:
        print(f"Error listing content: {e}")


def get_page_info(notion: Client, page_id: str) -> Optional[Dict]:
    """Get page information."""
    try:
        formatted_id = format_id_with_dashes(page_id)
        return notion.pages.retrieve(page_id=formatted_id)
    except Exception as e:
        print(f"Error retrieving page: {e}")
        return None


def get_database_info(notion: Client, database_id: str) -> Optional[Dict]:
    """Get database information."""
    try:
        formatted_id = format_id_with_dashes(database_id)
        return notion.databases.retrieve(database_id=formatted_id)
    except Exception as e:
        print(f"Error retrieving database: {e}")
        return None


def append_to_page(notion: Client, page_id: str, blocks: List[Dict]) -> bool:
    """Append blocks to an existing Notion page."""
    try:
        formatted_id = format_id_with_dashes(page_id)

        # Notion API limits to 100 blocks per request
        for i in range(0, len(blocks), 100):
            batch = blocks[i:i+100]
            notion.blocks.children.append(
                block_id=formatted_id,
                children=batch
            )

        return True
    except Exception as e:
        print(f"Error appending to page: {e}")
        return False


def clear_page_content(notion: Client, page_id: str) -> bool:
    """Clear all content from a page."""
    try:
        formatted_id = format_id_with_dashes(page_id)

        # Get all child blocks
        children = notion.blocks.children.list(block_id=formatted_id)

        # Delete each block
        for block in children.get('results', []):
            notion.blocks.delete(block_id=block['id'])

        return True
    except Exception as e:
        print(f"Error clearing page: {e}")
        return False


def create_database_row(notion: Client, database_id: str, properties: Dict, content_blocks: List[Dict] = None) -> Optional[str]:
    """Create a new row in a Notion database."""
    try:
        formatted_id = format_id_with_dashes(database_id)

        page_data = {
            "parent": {"database_id": formatted_id},
            "properties": properties
        }

        if content_blocks:
            page_data["children"] = content_blocks[:100]  # Limit on creation

        result = notion.pages.create(**page_data)

        # If more than 100 blocks, append the rest
        if content_blocks and len(content_blocks) > 100:
            append_to_page(notion, result['id'], content_blocks[100:])

        return result['id']
    except Exception as e:
        print(f"Error creating database row: {e}")
        return None


def write_to_page(notion: Client, page_id: str, markdown_content: str, mode: str = 'append') -> bool:
    """Write markdown content to a Notion page."""
    blocks = markdown_to_notion_blocks(markdown_content)

    if not blocks:
        print("No content to write")
        return False

    if mode == 'replace':
        if not clear_page_content(notion, page_id):
            return False

    return append_to_page(notion, page_id, blocks)


def main():
    parser = argparse.ArgumentParser(
        description='Push markdown content to Notion pages or databases',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  # Append markdown file to a page
  python notion_writer.py --page https://notion.so/My-Page-abc123 --file content.md

  # Replace page content
  python notion_writer.py --page PAGE_ID --file content.md --replace

  # Pipe content from stdin
  cat content.md | python notion_writer.py --page PAGE_ID --stdin

  # Create database row with title
  python notion_writer.py --database DB_URL --title "New Entry" --file content.md

  # Test connection
  python notion_writer.py --test
'''
    )

    parser.add_argument('--page', '-p', help='Notion page URL or ID')
    parser.add_argument('--database', '-d', help='Notion database URL or ID')
    parser.add_argument('--file', '-f', help='Markdown file to push')
    parser.add_argument('--stdin', action='store_true', help='Read content from stdin')
    parser.add_argument('--replace', '-r', action='store_true', help='Replace page content instead of append')
    parser.add_argument('--title', '-t', help='Title for database row (required for database)')
    parser.add_argument('--test', action='store_true', help='Test Notion API connection')
    parser.add_argument('--list', '-l', nargs='?', const='all', choices=['all', 'pages', 'databases'],
                        help='List accessible pages and/or databases (default: all)')
    parser.add_argument('--info', action='store_true', help='Show page/database info')

    args = parser.parse_args()

    if not NOTION_TOKEN:
        print("Error: NOTION_API_KEY not set in .env file")
        print("Get your integration token from: https://www.notion.so/my-integrations")
        sys.exit(1)

    notion = Client(auth=NOTION_TOKEN)

    # Test connection
    if args.test:
        try:
            me = notion.users.me()
            print("=" * 50)
            print("Notion API Connection Test")
            print("=" * 50)
            print(f"\n✅ Connected successfully!")
            print(f"   Bot: {me.get('name', 'Unknown')}")
            print(f"   Type: {me.get('type', 'Unknown')}")
            print("\n" + "=" * 50)
            return
        except Exception as e:
            print(f"❌ Connection failed: {e}")
            sys.exit(1)

    # List accessible content
    if args.list:
        filter_map = {'all': 'all', 'pages': 'page', 'databases': 'database'}
        list_accessible_content(notion, filter_map.get(args.list, 'all'))
        return

    # Show info
    if args.info:
        if args.page:
            page_id = extract_notion_id(args.page)
            if page_id:
                info = get_page_info(notion, page_id)
                if info:
                    print("=" * 50)
                    print("Page Information")
                    print("=" * 50)
                    title = ""
                    if 'properties' in info:
                        for prop in info['properties'].values():
                            if prop.get('type') == 'title':
                                title_arr = prop.get('title', [])
                                if title_arr:
                                    title = title_arr[0].get('plain_text', '')
                    print(f"Title: {title}")
                    print(f"ID: {info['id']}")
                    print(f"Created: {info.get('created_time', 'N/A')}")
                    print(f"URL: {info.get('url', 'N/A')}")
                    return

        if args.database:
            db_id = extract_notion_id(args.database)
            if db_id:
                info = get_database_info(notion, db_id)
                if info:
                    print("=" * 50)
                    print("Database Information")
                    print("=" * 50)
                    title = ""
                    if info.get('title'):
                        title = info['title'][0].get('plain_text', '') if info['title'] else ''
                    print(f"Title: {title}")
                    print(f"ID: {info['id']}")
                    print(f"Properties: {', '.join(info.get('properties', {}).keys())}")
                    return

        print("Please specify --page or --database with --info")
        return

    # Get content
    content = None
    if args.stdin:
        content = sys.stdin.read()
    elif args.file:
        file_path = Path(args.file)
        if not file_path.exists():
            print(f"Error: File not found: {args.file}")
            sys.exit(1)
        content = file_path.read_text(encoding='utf-8')

    # Write to page
    if args.page:
        if not content:
            print("Error: No content provided. Use --file or --stdin")
            sys.exit(1)

        page_id = extract_notion_id(args.page)
        if not page_id:
            print(f"Error: Invalid Notion page URL/ID: {args.page}")
            sys.exit(1)

        mode = 'replace' if args.replace else 'append'
        print(f"{'Replacing' if mode == 'replace' else 'Appending'} content to page...")

        if write_to_page(notion, page_id, content, mode):
            print(f"✅ Successfully wrote content to page")
            formatted_id = format_id_with_dashes(page_id)
            print(f"   https://notion.so/{formatted_id.replace('-', '')}")
        else:
            print("❌ Failed to write content")
            sys.exit(1)
        return

    # Create database row
    if args.database:
        if not content and not args.title:
            print("Error: Provide --title and/or --file for database row")
            sys.exit(1)

        db_id = extract_notion_id(args.database)
        if not db_id:
            print(f"Error: Invalid Notion database URL/ID: {args.database}")
            sys.exit(1)

        # Get database schema to find title property
        db_info = get_database_info(notion, db_id)
        if not db_info:
            sys.exit(1)

        # Find the title property
        title_prop = None
        for prop_name, prop_config in db_info.get('properties', {}).items():
            if prop_config.get('type') == 'title':
                title_prop = prop_name
                break

        if not title_prop:
            print("Error: Could not find title property in database")
            sys.exit(1)

        properties = {
            title_prop: {
                "title": [{"text": {"content": args.title or "Untitled"}}]
            }
        }

        content_blocks = markdown_to_notion_blocks(content) if content else None

        print(f"Creating database row...")
        row_id = create_database_row(notion, db_id, properties, content_blocks)

        if row_id:
            print(f"✅ Successfully created database row")
            print(f"   https://notion.so/{row_id.replace('-', '')}")
        else:
            print("❌ Failed to create database row")
            sys.exit(1)
        return

    # No action specified
    parser.print_help()


if __name__ == "__main__":
    main()