import asyncio
import random
import os
from dataclasses import dataclass, asdict
from typing import List, Optional

import pandas as pd
from playwright.async_api import async_playwright

try:
    from playwright_stealth import stealth_async
    STEALTH_AVAILABLE = True
except ImportError:
    STEALTH_AVAILABLE = False


@dataclass
class Video:
    rank: int = 0
    title: str = ""
    views: str = ""
    duration: str = ""
    url: str = ""
    video_id: str = ""
    content_type: str = "video"


@dataclass
class Short:
    rank: int = 0
    title: str = ""
    views: str = ""
    url: str = ""
    video_id: str = ""
    content_type: str = "short"


class YouTubeScraper:
    USER_AGENTS = [
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
    ]

    def __init__(self):
        self.browser = None
        self.page = None
        self.channel_name = ""

    async def run(
        self,
        link: str,
        filename: Optional[str] = None,
        folder: str = "data",
        scrape_videos: Optional[bool] = None,
        scrape_shorts: Optional[bool] = None,
        max_videos: Optional[int] = None,
        max_shorts: Optional[int] = None
    ) -> pd.DataFrame:
        """
        Main scraping function
        
        Args:
            link: Channel URL (e.g., https://www.youtube.com/@3blue1brown)
            filename: Output filename (without extension). If None, uses channel name
            folder: Output folder (default: "data")
            scrape_videos: Whether to scrape videos (None = auto-detect based on other params)
            scrape_shorts: Whether to scrape shorts (None = auto-detect based on other params)
            max_videos: Max videos to scrape (None = all available)
            max_shorts: Max shorts to scrape (None = all available)
        
        Returns:
            Combined DataFrame of all scraped content
        
        Examples:
            # Scrape both videos and shorts (default)
            await scraper.run("https://www.youtube.com/@channel")
            
            # Scrape only shorts
            await scraper.run("https://www.youtube.com/@channel", scrape_shorts=True)
            
            # Scrape only videos
            await scraper.run("https://www.youtube.com/@channel", scrape_videos=True)
            
            # Scrape both explicitly
            await scraper.run("https://www.youtube.com/@channel", scrape_videos=True, scrape_shorts=True)
        """
        
        # ============================================================
        # SMART DEFAULTS: Determine what to scrape
        # ============================================================
        # If neither specified → scrape both
        # If only one specified as True → scrape only that one
        # If explicitly set to False → respect that
        
        if scrape_videos is None and scrape_shorts is None:
            # Default: scrape both
            scrape_videos = True
            scrape_shorts = True
        else:
            # If one is specified, default the other to False
            if scrape_videos is None:
                scrape_videos = False
            if scrape_shorts is None:
                scrape_shorts = False
        
        # Clean base URL
        base_url = link.rstrip('/')
        for suffix in ['/videos', '/shorts', '/streams', '/playlists', '/community', '/channels', '/about', '/featured']:
            if base_url.endswith(suffix):
                base_url = base_url[:-len(suffix)]
        
        # Extract channel name
        self.channel_name = base_url.split('/')[-1].replace('@', '')
        
        # Set filename
        if filename is None:
            filename = self.channel_name
        
        all_content = []
        
        async with async_playwright() as p:
            await self._setup_browser(p)
            
            print("=" * 60)
            print(f"🎬 YouTube Channel Scraper")
            print(f"📺 Channel: {self.channel_name}")
            print(f"📁 Output: {folder}/{filename}.csv")
            print(f"🎯 Scraping: {'Videos' if scrape_videos else ''}{' + ' if scrape_videos and scrape_shorts else ''}{'Shorts' if scrape_shorts else ''}")
            print("=" * 60)
            
            # Scrape Videos
            if scrape_videos:
                print(f"\n{'='*60}")
                print("📹 SCRAPING VIDEOS")
                print("=" * 60)
                videos_url = f"{base_url}/videos"
                videos = await self._scrape_videos(videos_url, max_videos)
                all_content.extend(videos)
                
                if scrape_shorts:
                    await self.page.wait_for_timeout(random.randint(2000, 4000))
            
            # Scrape Shorts
            if scrape_shorts:
                print(f"\n{'='*60}")
                print("📱 SCRAPING SHORTS")
                print("=" * 60)
                shorts_url = f"{base_url}/shorts"
                shorts = await self._scrape_shorts(shorts_url, max_shorts)
                all_content.extend(shorts)
            
            await self.browser.close()
        
        # Convert to DataFrame
        if not all_content:
            print("\n⚠️ No content found.")
            return pd.DataFrame()
        
        df = pd.DataFrame([asdict(item) for item in all_content])
        
        # Print summary
        print(f"\n{'='*60}")
        print("📊 SCRAPING COMPLETE")
        print("=" * 60)
        
        videos_count = len(df[df['content_type'] == 'video']) if 'content_type' in df.columns else 0
        shorts_count = len(df[df['content_type'] == 'short']) if 'content_type' in df.columns else 0
        
        if scrape_videos:
            print(f"   📹 Videos: {videos_count}")
        if scrape_shorts:
            print(f"   📱 Shorts: {shorts_count}")
        print(f"   📦 Total:  {len(df)}")
        
        # Print preview
        print(f"\n{'='*60}")
        print("📋 PREVIEW")
        print("=" * 60)
        print(df.to_string())
        
        # Save to CSV
        os.makedirs(folder, exist_ok=True)
        output_path = f"{folder}/{filename}.csv"
        df.to_csv(output_path, index=False)
        print(f"\n💾 Saved: {output_path} ({len(df)} rows)")
        
        return df

    async def _setup_browser(self, playwright):
        """Setup browser with stealth"""
        self.browser = await playwright.chromium.launch(
            headless=True,
            args=[
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-blink-features=AutomationControlled',
            ]
        )
        
        context = await self.browser.new_context(
            viewport={"width": 1920, "height": 1080},
            user_agent=random.choice(self.USER_AGENTS),
            locale="en-US",
        )
        
        await context.set_extra_http_headers({
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            "Accept-Language": "en-US,en;q=0.9",
            "Accept-Encoding": "gzip, deflate, br",
        })
        
        if STEALTH_AVAILABLE:
            await stealth_async(context)
        
        self.page = await context.new_page()
        
        await self.page.add_init_script("""
            Object.defineProperty(navigator, 'webdriver', {get: () => undefined});
        """)

    async def _navigate_to_page(self, url: str) -> bool:
        """Navigate to a page"""
        print(f"🔗 URL: {url}")
        print(f"📄 Loading page...")
        
        try:
            await self.page.goto(url, wait_until="networkidle", timeout=60000)
            await self.page.wait_for_timeout(3000)
            
            await self._handle_consent()
            
            return True
            
        except Exception as e:
            print(f"⚠️ Failed to load page: {e}")
            return False

    async def _handle_consent(self):
        """Handle YouTube cookie consent dialog"""
        try:
            consent_selectors = [
                'button[aria-label*="Accept"]',
                'button:has-text("Accept all")',
                'button:has-text("I agree")',
                '[aria-label="Accept the use of cookies and other data for the purposes described"]',
                'tp-yt-paper-button:has-text("Accept all")',
            ]
            
            for selector in consent_selectors:
                try:
                    btn = self.page.locator(selector).first
                    if await btn.count() > 0 and await btn.is_visible(timeout=2000):
                        await btn.click(timeout=3000)
                        await self.page.wait_for_timeout(1500)
                        print("   ✅ Handled consent dialog")
                        break
                except:
                    continue
        except:
            pass

    async def _click_popular_button(self):
        """Click the Popular button to sort (if available)"""
        print(f"🔘 Looking for 'Popular' button...")
        
        try:
            await self.page.wait_for_timeout(2000)
            
            chip_bar_exists = await self.page.evaluate('''() => {
                const chipBar = document.querySelector('yt-chip-cloud-renderer, #chips-wrapper, iron-selector#chips');
                return chipBar !== null && chipBar.offsetParent !== null;
            }''')
            
            if not chip_bar_exists:
                print("   ℹ️ No filter buttons available")
                return
            
            popular_clicked = False
            
            try:
                clicked = await self.page.evaluate('''() => {
                    const chips = document.querySelectorAll('yt-chip-cloud-chip-renderer, yt-formatted-string, iron-selector#chips yt-formatted-string');
                    for (const chip of chips) {
                        const text = chip.textContent.trim().toLowerCase();
                        if (text === 'popular') {
                            chip.click();
                            return true;
                        }
                    }
                    return false;
                }''')
                popular_clicked = clicked
            except:
                pass
            
            if not popular_clicked:
                try:
                    popular_btn = self.page.locator('yt-chip-cloud-chip-renderer:has-text("Popular")').first
                    if await popular_btn.count() > 0 and await popular_btn.is_visible(timeout=2000):
                        await popular_btn.click(timeout=5000)
                        popular_clicked = True
                except:
                    pass
            
            if popular_clicked:
                await self.page.wait_for_timeout(3000)
                print("   ✅ Clicked Popular button")
            else:
                print("   ⚠️ Popular button not found, using default sort")
                
        except Exception as e:
            print(f"   ⚠️ Error: {e}")

    async def _scroll_to_load_all(self, content_type: str, max_items: Optional[int] = None):
        """Scroll page to load all content"""
        print(f"📜 Scrolling to load {content_type}...")
        
        prev_count = 0
        no_change_count = 0
        max_no_change = 10
        
        while no_change_count < max_no_change:
            if content_type == "videos":
                current_count = await self.page.evaluate('''() => {
                    const items = document.querySelectorAll('ytd-rich-item-renderer');
                    let count = 0;
                    items.forEach(item => {
                        const link = item.querySelector('a[href*="/watch?v="]');
                        if (link) count++;
                    });
                    return count;
                }''')
            else:
                current_count = await self.page.evaluate('''() => {
                    const items = document.querySelectorAll('ytd-rich-item-renderer');
                    let count = 0;
                    items.forEach(item => {
                        const link = item.querySelector('a[href*="/shorts/"]');
                        if (link) count++;
                    });
                    return count;
                }''')
            
            print(f"   Loaded: {current_count} {content_type}", end='\r')
            
            if max_items and current_count >= max_items:
                print(f"\n   ✅ Reached target: {current_count} {content_type}")
                break
            
            if current_count == prev_count:
                no_change_count += 1
            else:
                no_change_count = 0
            
            prev_count = current_count
            
            await self.page.evaluate('window.scrollBy(0, 1500)')
            await self.page.wait_for_timeout(random.randint(1000, 2000))
            
            if no_change_count > 5:
                await self.page.keyboard.press('End')
                await self.page.wait_for_timeout(1500)
        
        print(f"\n   📦 Total loaded: {prev_count} {content_type}")

    async def _scrape_videos(self, url: str, max_items: Optional[int] = None) -> List[Video]:
        """Scrape videos from channel"""
        
        success = await self._navigate_to_page(url)
        if not success:
            return []
        
        await self._click_popular_button()
        await self._scroll_to_load_all("videos", max_items)
        
        print(f"🔍 Extracting video data...")
        
        try:
            data = await self.page.evaluate('''() => {
                const videos = [];
                const seen = new Set();
                
                const items = document.querySelectorAll('ytd-rich-item-renderer');
                
                items.forEach((item) => {
                    try {
                        const titleLink = item.querySelector('a#video-title-link, a#video-title, a[href*="/watch?v="]');
                        if (!titleLink) return;
                        
                        const href = titleLink.href || '';
                        if (href.includes('/shorts/')) return;
                        
                        const videoIdMatch = href.match(/watch\\?v=([a-zA-Z0-9_-]{11})/);
                        if (!videoIdMatch) return;
                        
                        const videoId = videoIdMatch[1];
                        if (seen.has(videoId)) return;
                        seen.add(videoId);
                        
                        const title = titleLink.textContent.trim() || 
                                     titleLink.getAttribute('title') || '';
                        if (!title) return;
                        
                        let views = '';
                        const metadataLine = item.querySelector('#metadata-line');
                        if (metadataLine) {
                            const viewSpan = metadataLine.querySelector('span.inline-metadata-item');
                            if (viewSpan) {
                                views = viewSpan.textContent.trim();
                            }
                        }
                        
                        let duration = '';
                        const durationEl = item.querySelector(
                            'ytd-thumbnail-overlay-time-status-renderer #text, ' +
                            'span.ytd-thumbnail-overlay-time-status-renderer, ' +
                            'badge-shape .badge-shape-wiz__text'
                        );
                        if (durationEl) {
                            duration = durationEl.textContent.trim();
                        }
                        
                        videos.push({
                            title,
                            views,
                            duration,
                            videoId,
                            url: 'https://www.youtube.com/watch?v=' + videoId
                        });
                        
                    } catch(e) {}
                });
                
                return videos;
            }''')
            
            videos = []
            items_to_process = data[:max_items] if max_items else data
            
            for idx, v in enumerate(items_to_process, 1):
                videos.append(Video(
                    rank=idx,
                    title=v.get('title', ''),
                    views=v.get('views', ''),
                    duration=v.get('duration', ''),
                    url=v.get('url', ''),
                    video_id=v.get('videoId', ''),
                    content_type='video'
                ))
            
            print(f"✅ Extracted {len(videos)} videos")
            return videos
            
        except Exception as e:
            print(f"⚠️ Extraction error: {e}")
            return []

    async def _scrape_shorts(self, url: str, max_items: Optional[int] = None) -> List[Short]:
        """Scrape shorts from channel"""
        
        success = await self._navigate_to_page(url)
        if not success:
            return []
        
        current_url = self.page.url
        if '/shorts' not in current_url:
            print("   ⚠️ Channel has no shorts")
            return []
        
        await self._click_popular_button()
        await self._scroll_to_load_all("shorts", max_items)
        
        print(f"🔍 Extracting shorts data...")
        
        try:
            data = await self.page.evaluate('''() => {
                const shorts = [];
                const seen = new Set();
                
                const items = document.querySelectorAll('ytd-rich-item-renderer');
                
                items.forEach((item) => {
                    try {
                        const thumbnailLink = item.querySelector('a#thumbnail[href*="/shorts/"], a[href*="/shorts/"]');
                        if (!thumbnailLink) return;
                        
                        const href = thumbnailLink.href || '';
                        
                        const shortIdMatch = href.match(/\\/shorts\\/([a-zA-Z0-9_-]{11})/);
                        if (!shortIdMatch) return;
                        
                        const videoId = shortIdMatch[1];
                        if (seen.has(videoId)) return;
                        seen.add(videoId);
                        
                        // Get title
                        let title = '';
                        const titleEl = item.querySelector('#video-title, yt-formatted-string#video-title, a#video-title-link');
                        if (titleEl) {
                            title = titleEl.textContent.trim() || titleEl.getAttribute('title') || '';
                        }
                        
                        if (!title) {
                            const h3Link = item.querySelector('h3 a, #details h3');
                            if (h3Link) {
                                title = h3Link.textContent.trim() || h3Link.getAttribute('title') || '';
                            }
                        }
                        
                        if (!title && thumbnailLink.getAttribute('aria-label')) {
                            const ariaLabel = thumbnailLink.getAttribute('aria-label');
                            title = ariaLabel.split(' - ')[0] || ariaLabel;
                        }
                        
                        // Get views
                        let views = '';
                        
                        const metadataLine = item.querySelector('#metadata-line, ytd-video-meta-block #metadata-line');
                        if (metadataLine) {
                            const spans = metadataLine.querySelectorAll('span');
                            for (const span of spans) {
                                const text = span.textContent.trim();
                                if (text.toLowerCase().includes('view')) {
                                    views = text;
                                    break;
                                }
                            }
                        }
                        
                        if (!views) {
                            const allSpans = item.querySelectorAll('span');
                            for (const span of allSpans) {
                                const text = span.textContent.trim();
                                if (text.toLowerCase().includes('view') && /\\d/.test(text)) {
                                    views = text;
                                    break;
                                }
                            }
                        }
                        
                        if (!views && thumbnailLink.getAttribute('aria-label')) {
                            const ariaLabel = thumbnailLink.getAttribute('aria-label');
                            const viewMatch = ariaLabel.match(/(\\d[\\d,\\.]*[KMB]?\\s*views?)/i);
                            if (viewMatch) {
                                views = viewMatch[1];
                            }
                        }
                        
                        if (!views) {
                            const overlay = item.querySelector('ytd-thumbnail-overlay-bottom-panel-renderer, .ytd-thumbnail-overlay-bottom-panel-renderer');
                            if (overlay) {
                                const text = overlay.textContent.trim();
                                if (text.toLowerCase().includes('view')) {
                                    views = text;
                                }
                            }
                        }
                        
                        shorts.push({
                            title,
                            views,
                            videoId,
                            url: 'https://www.youtube.com/shorts/' + videoId
                        });
                        
                    } catch(e) {}
                });
                
                return shorts;
            }''')
            
            shorts = []
            items_to_process = data[:max_items] if max_items else data
            
            for idx, s in enumerate(items_to_process, 1):
                shorts.append(Short(
                    rank=idx,
                    title=s.get('title', ''),
                    views=s.get('views', ''),
                    url=s.get('url', ''),
                    video_id=s.get('videoId', ''),
                    content_type='short'
                ))
            
            print(f"✅ Extracted {len(shorts)} shorts")
            return shorts
            
        except Exception as e:
            print(f"⚠️ Extraction error: {e}")
            return []

