#!/usr/bin/env python3
"""
PTV Multi-Channel EPG Scraper - Enhanced Version
- Real-time progress output
- Comprehensive time parsing
- Robust error handling
- All original functionality preserved
"""

import sys
import signal
import time
import threading
import os
import re
from datetime import datetime, timedelta
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from bs4 import BeautifulSoup
import xml.etree.ElementTree as ET
from xml.dom.minidom import parseString
import pytz
from pathlib import Path
from collections import defaultdict
import undetected_chromedriver as uc
import random

# ======================== CONFIGURATION ========================
OUTPUT_FILENAME = "ptv_epg.xml"
DEFAULT_PROGRAM_DURATION = {
    'sports': timedelta(hours=4),
    'default': timedelta(hours=1)
}

ABBREVIATIONS = {
    # Sports abbreviations
    'ICC', 'EPL', 'IPL', 'PCB', 'FIFA', 'UEFA', 'NBA', 'NFL', 'NHL', 'MLB', 
    'PGA', 'US', 'UK', 'USA', 'T20', 'ODI', 'PSL', 'FIA', 'F1', 'ATP', 
    'WTA', 'WTC', 'UFC', 'CFL', 'PZ', 'KK', 'LQ','QG','MS','IU','IL',
    'GG', 'DV','TRI','WI','NZ','PK','PAK','IN','SL','BD','BG','EN','SA','US','USA','UK',
    'AU','AUS','SF','H/L','T20I'
    
    # Home abbreviations
    'PTV', 'BBC', 'CNN', 'PBC', 'EID', 'EP', 'NO', 'PROG', 
    'QURAN', 'NAAT', 'HADITH', 'PAKISTAN', 'AM', 'PM', 'EID','AZAN','EP','CA','PAK',
    
    # World abbreviations
    'R', 'RPT', 'HEADLINES', 'AL QURAN', 'COVID', 'COVID-19', 'UAE', 
    'UN', 'WHO', 'WTO', 'IMF', 'EU', 'NATO', 'UNSC', 'LIVE',
    
    # News abbreviations
    'PTI', 'AFP', 'AP', 'UNESCO',
    
    # Global abbreviations
    'S.A.W', 'DC', 'TV', 'F',
    
    # National abbreviations
    'NN', 'AIOU', 'SINDHI', 'SINDH', 
    'AL-QURAN', 
    
    # Bolan abbreviations
    'MN', 'BRAHVI', 'BALOCHI', 'PASHTO', 
    'QURAAN', 'IQRA', 'AZAN', 'MAGHRIB', 'FAJAR', 'SEHAR', 
    'DAROOD', 'IBRAHIM'
}

CHANNELS = {
    'sports': {
        'id': 'PTV SPORTS.pk',
        'name': 'PTV Sports',
        'url': 'https://sports.ptv.com.pk/tvguidemaster',
        'duration': 'sports'
    },
    'home': {
        'id': 'PTV HOME.pk',
        'name': 'PTV Home',
        'url': 'https://ptv.com.pk/ptvhome/tvguide',
        'duration': 'default'
    },
    'world': {
        'id': 'PTV WORLD.pk',
        'name': 'PTV World',
        'url': 'https://ptv.com.pk/ptvworld/tvguide',
        'duration': 'default'
    },
    'news': {
        'id': 'PTV NEWS.pk',
        'name': 'PTV News',
        'url': 'https://ptv.com.pk/ptvnews/tvguide',
        'duration': 'default'
    },
    'global': {
        'id': 'PTV GLOBAL.pk',
        'name': 'PTV Global',
        'url': 'https://ptv.com.pk/ptvglobal/tvguide',
        'duration': 'default'
    },
    'national': {
        'id': 'PTV NATIONAL.pk',
        'name': 'PTV National',
        'url': 'https://ptv.com.pk/ptvnational/tvguide',
        'duration': 'default'
    },
    'bolan': {
        'id': 'PTV BOLAN.pk',
        'name': 'PTV Bolan',
        'url': 'https://ptv.com.pk/ptvbolan/tvguide',
        'duration': 'default'
    }
}

# ======================== UTILITY FUNCTIONS ========================
def handler(signum, frame):
    print("\n⏹️ Script terminated by user")
    sys.exit(0)

signal.signal(signal.SIGINT, handler)

def countdown_timer():
    """Display countdown timer before automatic exit"""
    for i in range(5, 0, -1):
        print(f"\rWindow will close in {i} seconds... (Press Enter to exit now)", end="")
        time.sleep(1)
    print("\nClosing program automatically...")
    os._exit(0)

def print_channel_header(channel_name):
    """Print consistent channel header"""
    print(f"\n{'='*50}")
    print(f"📡 Scraping {CHANNELS[channel_name]['name']}")
    print(f"🌐 URL: {CHANNELS[channel_name]['url']}")

def print_day_results(day_name, programs, week_dates):
    """Print day's results with date and end times"""
    date_str = week_dates[day_name].strftime('%a %d %b')
    print(f"\n📅 {day_name} ({date_str}):")
    
    for i, program in enumerate(programs, 1):
        start_str = program['start'].strftime('%H:%M')
        end_str = program['end'].strftime('%H:%M') if program['end'] else '--:--'
        print(f"  {i}. {start_str}-{end_str} | {program['title']}")
    print(f"✔ Found {len(programs)} programs")

def format_title(title):
    """Improved title formatting with fixes for parentheses and specific time parsing"""
    # Protect all abbreviations (they should remain ALL CAPS)
    protected = {}
    for abbr in ABBREVIATIONS:
        if abbr in title.upper():
            # Find the exact case match in the original title
            matches = re.finditer(re.escape(abbr), title, re.IGNORECASE)
            for match in matches:
                protected[match.group()] = match.group()
    
    # Convert to Title Case while handling parentheses
    words = re.split(r'(\()', title)  # Split on parentheses to handle them separately
    formatted_parts = []
    
    for part in words:
        if part == '(':
            formatted_parts.append(part)
        elif part:
            # Process each word separately
            sub_words = part.split()
            formatted_sub_words = []
            for word in sub_words:
                if word in protected:
                    formatted_sub_words.append(word)
                else:
                    # Convert to Title Case while preserving hyphenated words
                    if '-' in word:
                        parts = word.split('-')
                        formatted_word = '-'.join([p.capitalize() for p in parts])
                    else:
                        formatted_word = word.capitalize()
                    formatted_sub_words.append(formatted_word)
            formatted_parts.append(' '.join(formatted_sub_words))
    
    # Reconstruct the title
    formatted_title = ''.join(formatted_parts)
    
    # Handle special cases that should remain ALL CAPS
    special_allcaps = ['LIVE', 'REPEAT', 'TELECAST', 'SPECIAL', 'PROG']
    words = formatted_title.split()
    for i, word in enumerate(words):
        if word.upper() in special_allcaps:
            words[i] = word.upper()
    
    return ' '.join(words)

def parse_time(time_str, channel=None):
    """Time parser with enhanced PTV Global multi-timezone handling"""
    if not time_str or not isinstance(time_str, str):
        print(f"⚠️ Empty/invalid time string for channel {channel}")
        return None

    original_time = time_str.strip()
    time_str = original_time.upper()

    # ===== Special Handling for PTV Global =====
    if channel == 'global':
        # Handle all variants: PST, ST, PSA, and cases with multiple UK times
        time_match = re.search(
            r'(?P<time>\d{3,4})\s*(?:PST|ST|PSA|UK)',
            time_str.replace(',', ' ')  # Handle comma-separated formats
        )
        if time_match:
            raw_time = time_match.group('time')
            if len(raw_time) == 3:  # Handle 3-digit times like '600PST'
                raw_time = '0' + raw_time
            try:
                return datetime.strptime(raw_time.zfill(4), "%H%M").time()
            except ValueError:
                pass

    # ===== Handling for National/Bolan TO formats =====
    if channel in ['national', 'bolan'] and 'TO' in time_str:
        time_str = time_str.split('TO')[0].strip()

    # ===== Common Time Processing =====
    time_str = (time_str.replace('.', ':')
                .replace(' ', '')
                .replace('NOON', 'PM')
                .replace('N00N', 'PM')
                .replace('MN', 'AM')
                .replace('TO', ''))

    formats = [
        "%I%M%p", "%I:%M%p", "%H%M", "%H:%M", 
        "%I%M %p", "%I:%M %p", "%H%M%p", "%H:%M%p",
        "%I%M", "%I:%M"
    ]

    for fmt in formats:
        try:
            clean_time = re.sub(r'[^\dAPM:]', '', time_str)
            return datetime.strptime(clean_time, fmt).time()
        except ValueError:
            continue

    print(f"⚠️ Could not parse time: '{original_time}' (channel: {channel})")
    return None

def get_week_dates():
    """Calculate dates for today and tomorrow only"""
    today = datetime.now()
    today_name = today.strftime("%A")
    
    if today_name == "Sunday":
        tomorrow = today + timedelta(days=1)
        tomorrow_name = "Monday"
    else:
        tomorrow = today + timedelta(days=1)
        tomorrow_name = tomorrow.strftime("%A")
    
    return {
        today_name: today,
        tomorrow_name: tomorrow
    }

def calculate_end_times(programs, channel):
    """Calculate end times for programs based on next program's start time"""
    if not programs:
        return []
    
    programs = sorted(programs, key=lambda p: p['start'])
    
    for i in range(len(programs) - 1):
        programs[i]['end'] = programs[i+1]['start']
    
    # For last program of today
    if programs:
        last_program = programs[-1]
        last_program['end'] = last_program['start'] + timedelta(hours=3)  # 3 hours duration
    
    return programs

# ======================== CHANNEL SCRAPERS ========================
def scrape_channel(driver, channel_name):
    """Generic scraper for all PTV channels (updated for new site structure)"""
    print_channel_header(channel_name)
    try:
        driver.get(CHANNELS[channel_name]['url'])
        time.sleep(random.uniform(2, 5))  # Random wait for human-like behavior
        
        WebDriverWait(driver, 20).until(
            EC.presence_of_element_located((By.CSS_SELECTOR, "#dayofweek a"))
        )
        
        tabs = driver.find_elements(By.CSS_SELECTOR, "#dayofweek a[data-dayofweek]")
        if not tabs:
            print("⚠️ Could not find day tabs")
            return []
        
        all_programs = []
        week_dates = get_week_dates()
        today_name = datetime.now().strftime("%A")
        tomorrow_name = (datetime.now() + timedelta(days=1)).strftime("%A")
        
        # Find today's and tomorrow's tabs
        today_tab = None
        tomorrow_tab = None
        
        for tab in tabs:
            day_name = tab.get_attribute("data-dayofweek").strip()
            if day_name == today_name:
                today_tab = tab
            elif day_name == tomorrow_name:
                tomorrow_tab = tab
        
        if not today_tab:
            print(f"⚠️ Could not find today's tab ({today_name})")
            return []
        
        # ===== Scrape Today's Programs =====
        day_programs = []
        try:
            driver.execute_script("arguments[0].click();", today_tab)
            time.sleep(random.uniform(2, 4))
            
            soup = BeautifulSoup(driver.page_source, "html.parser")
            time_blocks = soup.select("div.col-3[style*='border-bottom']")
            
            for time_block in time_blocks:
                time_span = time_block.select_one("span.rt-meta")
                if not time_span:
                    continue
                
                time_text = time_span.get_text(strip=True).replace(' ', '')
                start_time = parse_time(time_text, channel_name)
                if not start_time:
                    continue
                
                start_datetime = datetime.combine(week_dates[today_name].date(), start_time)
                start_datetime = pytz.timezone('Asia/Karachi').localize(start_datetime)
                
                title_block = time_block.find_next_sibling("div", class_="col-9")
                if title_block:
                    title_h4 = title_block.select_one("h4.post-title")
                    if title_h4:
                        title = title_h4.get_text(strip=True)
                        program = {
                            'channel': channel_name,
                            'day': today_name,
                            'title': format_title(title),
                            'start': start_datetime,
                            'end': None
                        }
                        day_programs.append(program)
        
        except Exception as e:
            print(f"⚠️ Error scraping {today_name}: {str(e)}")
        
       
        # ===== Combine and Calculate =====
        all_programs = day_programs
        
        if all_programs:
            all_programs = calculate_end_times(all_programs, channel_name)
            
            # Print
            if day_programs:
                print_day_results(today_name, day_programs, week_dates)
                   
        print(f"\n✅ {CHANNELS[channel_name]['name']} complete: {len(all_programs)} programs")
        return all_programs
    
    except Exception as e:
        print(f"⚠️ Failed to scrape {CHANNELS[channel_name]['name']}: {str(e)}")
        return []


# ======================== EPG GENERATION ========================
def generate_epg(all_programs):
    """Generate XMLTV format EPG file with all channels"""
    if not all_programs:
        print("⚠️ Warning: No program data provided for EPG generation")
        return False
    
    print("\nStarting EPG XML generation...")
    
    success = False
    program_count = 0
    root = None
    
    output_dir = Path(__file__).parent / "tempest_config" / "epg"
    output_file = output_dir / OUTPUT_FILENAME
    
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
        
        root = ET.Element("tv", {
            "generator-info-name": "PTV Multi-Channel EPG Generator",
            "generator-info-url": ""
        })
        
        # Add channel information
        for channel in CHANNELS.values():
            channel_elem = ET.SubElement(root, "channel", {"id": channel['id']})
            ET.SubElement(channel_elem, "display-name").text = channel['name']
        
        # Add programs
        for program in all_programs:
            try:
                if not all(k in program for k in ['title', 'start', 'end', 'channel']):
                    continue
                    
                channel_id = CHANNELS[program['channel']]['id']
                
                # Convert to GMT for EPG format
                start_gmt = program["start"].astimezone(pytz.utc).strftime("%Y%m%d%H%M%S +0000")
                end_gmt = program["end"].astimezone(pytz.utc).strftime("%Y%m%d%H%M%S +0000")
                
                program_attrs = {
                    "start": start_gmt,
                    "stop": end_gmt,
                    "channel": channel_id
                }
                programme = ET.SubElement(root, "programme", program_attrs)
                ET.SubElement(programme, "title", {'lang': 'en'}).text = program["title"]
                
                program_count += 1
            except Exception as e:
                continue
                
        if program_count > 0:
            xml_string = ET.tostring(root, encoding="utf-8", xml_declaration=True)
            pretty_xml = parseString(xml_string).toprettyxml()
            
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(pretty_xml)
            
            print(f"✅ Successfully generated EPG with {program_count} programs")
            print(f"📁 EPG file saved to {output_file.absolute()}")
            success = True
            
    except Exception as e:
        print(f"⚠️ Error during EPG generation: {str(e)}")
        success = False
        
    return success

def print_program_summary(programs):
    """Print summary of scraped programs"""
    if not programs:
        print("\n⚠️ No programs found or operation failed")
        return
    
    # Group by channel
    channel_programs = defaultdict(list)
    for p in programs:
        channel_programs[p['channel']].append(p)
    
    print("\n📺 Program Summary:")
    for channel, progs in channel_programs.items():
        print(f"\n{CHANNELS[channel]['name']}: {len(progs)} programs")
        for i, p in enumerate(progs[:5], 1):
            day_date = p['start'].strftime('%a %d %b')
            print(f"  {i}. [{day_date}] {p['start'].strftime('%H:%M')} - {p['end'].strftime('%H:%M')} | {p['title']}")
        if len(progs) > 5:
            print(f"  ... and {len(progs)-5} more")

# ======================== MAIN EXECUTION ========================
if __name__ == "__main__":
    start_time = time.time()
    
      
    options = uc.ChromeOptions()
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-gpu')
    options.add_argument('--disable-dev-shm-usage')

    driver = uc.Chrome(options=options)
    try:
        print("\n🚀 PTV Multi-Channel EPG Scraper")
        print(f"📅 Week of: {datetime.now().strftime('%Y-%m-%d')}")
        print("⏳ Starting scraping process...\n")
        
        # Scrape all channels
        all_programs = []
        for channel_name in CHANNELS.keys():
            programs = scrape_channel(driver, channel_name)
            all_programs.extend(programs)
            print(f"\n⭐ {CHANNELS[channel_name]['name']} completed")
        
        # Generate output
        print("\n" + "="*50)
        print("🎉 All Channels Scraped Successfully!")
        print_program_summary(all_programs)
        
        if generate_epg(all_programs):
            print("\n✅ EPG file generated successfully!")
        
        elapsed = time.time() - start_time
        print(f"\n⏱️ Total execution time: {elapsed:.2f} seconds")
        
    except Exception as e:
        print(f"\n🔥 Critical error: {str(e)}")
    finally:
        driver.quit()
    
    # Exit handling
    print("\nScript execution complete")
    timer_thread = threading.Thread(target=countdown_timer)
    timer_thread.daemon = True
    timer_thread.start()
    
    try:
        input("Press Enter to exit now...")
        os._exit(0)
    except:
        pass