#!/usr/bin/env python3
"""
Download All NotebookLM Assets Script

This script downloads ALL assets from a NotebookLM notebook, including
multiple audio and video files that were created.
"""
import os
import sys
import json
import asyncio
from datetime import datetime

# Simple database access without full Flask app
try:
    import psycopg2
    from psycopg2.extras import RealDictCursor
    DB_AVAILABLE = True
except ImportError:
    DB_AVAILABLE = False
    print("Warning: Database not available, will only download files")

def download_all_notebooklm_assets():
    """Download all assets from NotebookLM notebook."""
    notebook_id = "0f38d000-7497-442c-9b67-8773b02a558b"
    nlm_account_id = 10
    report_id = 35
    output_dir = f"/home/ashraffarid2010/jadwaai.com/app/static/uploads/notebooks/{report_id}"

    print("=" * 80)
    print("DOWNLOAD ALL NOTEBOOKLM ASSETS")
    print("=" * 80)
    print(f"\n📋 Configuration:")
    print(f"  Notebook ID: {notebook_id}")
    print(f"  Report ID: {report_id}")
    print(f"  NLM Account ID: {nlm_account_id}")
    print(f"  Output Directory: {output_dir}")

    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    print(f"\n🔍 Step 1: Checking current files...")
    current_files = set(os.listdir(output_dir)) if os.path.exists(output_dir) else set()
    print(f"  Current files: {len(current_files)}")
    for f in sorted(current_files):
        filepath = os.path.join(output_dir, f)
        if os.path.isfile(filepath):
            size = os.path.getsize(filepath) / (1024 * 1024)
            print(f"    - {f} ({size:.2f} MB)")

    print(f"\n🔍 Step 2: Listing all artifacts in NotebookLM...")
    print("  This will show us what's actually available in NotebookLM")

    try:
        # Import notebooklm library
        from notebooklm import NotebookLMClient
        import sys

        # Set up environment for admin account
        admin_session_path = f"/home/ashraffarid2010/jadwaai.com/data/nlm_admin_accounts/{nlm_account_id}/storage_state.json"

        if not os.path.exists(admin_session_path):
            print(f"  ❌ Admin session not found: {admin_session_path}")
            return

        # Set NOTEBOOKLM_HOME
        session_dir = os.path.dirname(admin_session_path)
        os.environ["NOTEBOOKLM_HOME"] = session_dir

        print(f"\n🔍 Step 3: Connecting to NotebookLM...")
        print(f"  Session path: {admin_session_path}")

        async def list_and_download():
            """List and download all artifacts."""
            try:
                timeout = 300.0  # 5 minutes for operations
                client = await NotebookLMClient.from_storage(timeout=timeout)

                print(f"\n✅ Connected to NotebookLM")

                # Get all artifacts from the notebook
                print(f"\n🔍 Step 4: Listing all artifacts in notebook...")
                artifacts_data = await client.notebooks.list_artifacts(notebook_id)

                print(f"  Raw artifacts response type: {type(artifacts_data)}")
                print(f"  Artifacts data: {artifacts_data}")

                # Try to parse the artifacts
                all_artifacts = []
                if isinstance(artifacts_data, dict):
                    # Check for common keys
                    for key in ['artifacts', 'items', 'data', 'results']:
                        if key in artifacts_data:
                            items = artifacts_data[key]
                            if isinstance(items, list):
                                all_artifacts.extend(items)
                                break
                    # If no list found, try dict values
                    if not all_artifacts:
                        for value in artifacts_data.values():
                            if isinstance(value, list):
                                all_artifacts.extend(value)
                elif isinstance(artifacts_data, list):
                    all_artifacts = artifacts_data

                print(f"\n📊 Found {len(all_artifacts)} artifacts")

                # Artifact storage
                downloaded_files = {}
                artifact_counts = {}

                for idx, artifact in enumerate(all_artifacts):
                    print(f"\n--- Artifact {idx + 1} ---")
                    print(f"  Type: {type(artifact)}")
                    print(f"  Content: {artifact}")

                    # Try to extract artifact type and info
                    artifact_type = "unknown"
                    artifact_id = None

                    if isinstance(artifact, dict):
                        artifact_type = artifact.get('type', artifact.get('artifact_type', 'unknown'))
                        artifact_id = artifact.get('id', artifact.get('artifact_id'))
                    elif hasattr(artifact, 'type'):
                        artifact_type = artifact.type
                    elif hasattr(artifact, 'artifact_type'):
                        artifact_type = artifact.artifact_type

                    print(f"  Artifact Type: {artifact_type}")
                    print(f"  Artifact ID: {artifact_id}")

                    # Count artifacts by type
                    if artifact_type not in artifact_counts:
                        artifact_counts[artifact_type] = 0
                    artifact_counts[artifact_type] += 1

                    # Download logic based on type
                    count = artifact_counts[artifact_type]

                    if artifact_type in ['audio', 'audio_overview']:
                        try:
                            output_path = os.path.join(output_dir, f"audio_{count}.mp3")
                            print(f"  🎧 Downloading audio to: {output_path}")
                            await client.artifacts.download_audio(notebook_id, output_path)
                            downloaded_files[f'audio_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading audio: {e}")

                    elif artifact_type in ['video', 'video_overview']:
                        try:
                            output_path = os.path.join(output_dir, f"video_{count}.mp4")
                            print(f"  🎬 Downloading video to: {output_path}")
                            await client.artifacts.download_video(notebook_id, output_path)
                            downloaded_files[f'video_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading video: {e}")

                    elif artifact_type in ['mind_map', 'mindmap']:
                        try:
                            output_path = os.path.join(output_dir, f"mind_map_{count}.json")
                            print(f"  🧠 Downloading mind map to: {output_path}")
                            await client.artifacts.download_mind_map(notebook_id, output_path, output_format="json")
                            downloaded_files[f'mind_map_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading mind map: {e}")

                    elif artifact_type == 'report':
                        try:
                            output_path = os.path.join(output_dir, f"report_{count}.md")
                            print(f"  📄 Downloading report to: {output_path}")
                            await client.artifacts.download_report(notebook_id, output_path)
                            downloaded_files[f'report_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading report: {e}")

                    elif artifact_type in ['flashcards', 'flashcard']:
                        try:
                            output_path = os.path.join(output_dir, f"flashcards_{count}.json")
                            print(f"  📝 Downloading flashcards to: {output_path}")
                            await client.artifacts.download_flashcards(notebook_id, output_path, output_format="json")
                            downloaded_files[f'flashcards_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading flashcards: {e}")

                    elif artifact_type in ['quiz', 'quizzes']:
                        try:
                            output_path = os.path.join(output_dir, f"quiz_{count}.json")
                            print(f"  ❓ Downloading quiz to: {output_path}")
                            await client.artifacts.download_quiz(notebook_id, output_path, output_format="json")
                            downloaded_files[f'quiz_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading quiz: {e}")

                    elif artifact_type in ['infographic']:
                        try:
                            output_path = os.path.join(output_dir, f"infographic_{count}.png")
                            print(f"  📊 Downloading infographic to: {output_path}")
                            await client.artifacts.download_infographic(notebook_id, output_path)
                            downloaded_files[f'infographic_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading infographic: {e}")

                    elif artifact_type in ['slide_deck', 'slidedeck', 'slide_deck']:
                        try:
                            output_path = os.path.join(output_dir, f"slide_deck_{count}.pdf")
                            print(f"  📑 Downloading slide deck to: {output_path}")
                            await client.artifacts.download_slide_deck(notebook_id, output_path)
                            downloaded_files[f'slide_deck_{count}'] = output_path
                            print(f"  ✅ Downloaded: {output_path}")
                        except Exception as e:
                            print(f"  ❌ Error downloading slide deck: {e}")

                return downloaded_files, artifact_counts

            except Exception as e:
                print(f"❌ Error: {e}")
                import traceback
                traceback.print_exc()
                return {}, {}

        # Run the async function
        downloaded_files, artifact_counts = asyncio.run(list_and_download())

        print(f"\n📊 SUMMARY")
        print(f"=" * 80)
        print(f"Downloaded Files: {len(downloaded_files)}")
        for name, path in downloaded_files.items():
            size = os.path.getsize(path) / (1024 * 1024)
            print(f"  ✅ {name}: {path} ({size:.2f} MB)")

        print(f"\nArtifact Counts: {json.dumps(artifact_counts, indent=2)}")

        # Create database update suggestions
        if DB_AVAILABLE:
            print(f"\n📊 DATABASE UPDATES NEEDED:")
            print(f"=" * 80)

            conn = psycopg2.connect(
                host="127.0.0.1",
                port=5432,
                database="jadwa",
                user="postgres",
            )

            cursor = conn.cursor(cursor_factory=RealDictCursor)

            # Get current database records
            cursor.execute("""
                SELECT asset_type, COUNT(*) as count
                FROM notebook_assets
                WHERE report_id = 35
                GROUP BY asset_type
            """)
            current_db_records = {row['asset_type']: row['count'] for row in cursor.fetchall()}

            print(f"Current database records: {current_db_records}")
            print(f"Actual files available: {artifact_counts}")

            cursor.close()
            conn.close()

    except Exception as e:
        print(f"❌ Error during download process: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    download_all_notebooklm_assets()