import requests
from bs4 import BeautifulSoup
import csv
import time
import argparse
from urllib.parse import urljoin
import os
import re

# Base URL
BASE_URL = "https://www.mcmc.gov.my"
SEARCH_TEMPLATE = "https://www.mcmc.gov.my/en/legal/registers/register-of-apparatus-assignments-search?type={station_type}"

# Default output directory for CSV files
CSV_OUTPUT_DIR = "/var/www/mypapit.net/public_html/mypapit.net-orig/pub/callsign"

# HTTP headers to mimic a real browser
HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/124.0.0.0 Safari/537.36",
    "Accept-Language": "en-US,en;q=0.9",
}

# The station types we can scrape callsigns from
STATION_TYPES = {
    "AARadio": "Amateur Radio Station",
    "AAAero": "Aircraft Station",
    "AAMaritime": "Ship Station"
}

def fetch_page(url: str, session: requests.Session, method: str = "GET", data: dict = None) -> str:
    """Download a page and return its HTML content with retries using session."""
    for attempt in range(3):
        try:
            if method == "POST":
                response = session.post(url, headers=HEADERS, data=data, timeout=30)
            else:
                response = session.get(url, headers=HEADERS, timeout=30)
            response.raise_for_status()
            return response.text
        except requests.RequestException as e:
            print(f"  [Warning] Attempt {attempt + 1} failed: {e}")
            if attempt < 2:
                time.sleep(5)
            else:
                raise

def parse_table(html: str) -> list[dict]:
    """Extract rows from the result table on the page."""
    soup = BeautifulSoup(html, "html.parser")
    table = soup.find("table")
    if not table:
        return []
    
    # Read headers dynamically
    thead = table.find("thead")
    if not thead:
        return []
    header_cells = thead.find_all("th")
    headers = [cell.get_text(strip=True) for cell in header_cells]
    
    rows = []
    tbody = table.find("tbody")
    if not tbody:
        return []
        
    for tr in tbody.find_all("tr"):
        cells = tr.find_all("td")
        if len(cells) != len(headers):
            continue
        row = {headers[i]: cells[i].get_text(strip=True) for i in range(len(headers))}
        rows.append(row)
    return rows

def parse_hidden_inputs(html: str) -> dict:
    """Extract all hidden input fields from the HTML."""
    soup = BeautifulSoup(html, "html.parser")
    inputs = {}
    for inp in soup.find_all("input", type="hidden"):
        name = inp.get("name")
        if name:
            inputs[name] = inp.get("value", "")
    return inputs

def parse_pagination_info(html: str) -> tuple[str, int]:
    """Extract event target and maximum page number from the pagination links."""
    soup = BeautifulSoup(html, "html.parser")
    # Default event target if none found
    default_target = 'p$lt$WebPartZone12$ZContent$pageplaceholder$p$lt$WebPartZone7$ZRegisterApparatusData$UniversalPager$pagerElem'
    event_target = default_target
    max_page = 1
    
    ul = soup.find("ul", class_="pagination")
    if ul:
        a_links = ul.find_all("a", href=re.compile(r"__doPostBack"))
        for a in a_links:
            href = a["href"]
            # Extract event target and page number
            m = re.search(r"__doPostBack\('([^']*)',\s*'(\d+)'\)", href)
            if m:
                target, page_num = m.group(1), int(m.group(2))
                event_target = target
                if page_num > max_page:
                    max_page = page_num
    return event_target, max_page

def scrape_station_type(station_type: str, max_pages: int, delay: float, output_csv_dir: str) -> int:
    """Scrape a specific station type from MCMC and save to CSV incrementally."""
    display_name = STATION_TYPES[station_type]
    print(f"\n=== Starting Scraping for {display_name} ({station_type}) ===")
    
    session = requests.Session()
    url = SEARCH_TEMPLATE.format(station_type=station_type)
    
    # Page 1 (GET request)
    print(f"Page 1: Fetching initial page {url}")
    try:
        html = fetch_page(url, session, method="GET")
    except Exception as e:
        print(f"  [Error] Failed to fetch page 1: {e}")
        return 0
        
    rows = parse_table(html)
    if not rows:
        print("  No records found or page structure changed.")
        return 0
        
    print(f"  Successfully extracted {len(rows)} records.")
    
    # Establish CSV file name and output directory
    os.makedirs(output_csv_dir, exist_ok=True)
    csv_filename = os.path.join(output_csv_dir, f"mcmc_{station_type.lower()}_assignments.csv")
    
    total_scraped = len(rows)
    
    # Write Page 1 immediately and keep file handle open for appending
    fieldnames = list(rows[0].keys())
    with open(csv_filename, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)
        f.flush()
        
        # Get pagination info and other hidden fields
        hidden_inputs = parse_hidden_inputs(html)
        event_target, parsed_max_pages = parse_pagination_info(html)
        
        limit_pages = min(max_pages, parsed_max_pages)
        print(f"  Total pages available: {parsed_max_pages}. Scraping up to page {limit_pages}.")
        
        for page in range(2, limit_pages + 1):
            # Implement delay between requests
            print(f"  Waiting {delay} seconds...")
            time.sleep(delay)
            
            print(f"Page {page}/{limit_pages}: Fetching via POST to {url}")
            
            # Prepare postback data
            post_data = hidden_inputs.copy()
            post_data["__EVENTTARGET"] = event_target
            post_data["__EVENTARGUMENT"] = str(page)
            
            try:
                html = fetch_page(url, session, method="POST", data=post_data)
                rows = parse_table(html)
                
                if not rows:
                    print("  No records found on this page. Stopping.")
                    break
                    
                print(f"  Successfully extracted {len(rows)} records.")
                writer.writerows(rows)
                f.flush()
                total_scraped += len(rows)
                
                # Update hidden inputs and target for next iteration
                hidden_inputs = parse_hidden_inputs(html)
                event_target, _ = parse_pagination_info(html)
                
            except Exception as e:
                print(f"  [Error] Error on page {page}: {e}")
                break
                
    print(f"Scraped details for {display_name} saved directly to: {csv_filename}")
    return total_scraped

def main():
    parser = argparse.ArgumentParser(description="Scrape MCMC Apparatus Assignments for Amateur Radio, Aircraft, and Ship Stations.")
    parser.add_argument("--type", choices=["AARadio", "AAAero", "AAMaritime", "all"], default="AARadio",
                        help="The type of station assignments to scrape (default: AARadio)")
    parser.add_argument("--pages", type=int, default=1000,
                        help="Maximum number of pages to scrape per station type (default: 1000)")
    parser.add_argument("--delay", type=float, default=2.0,
                        help="Delay in seconds between page requests (default: 2.0)")
    parser.add_argument("--outdir", type=str, default=CSV_OUTPUT_DIR,
                        help=f"Output directory for CSV files (default: {CSV_OUTPUT_DIR})")
    
    args = parser.parse_args()
    
    types_to_scrape = [args.type] if args.type != "all" else list(STATION_TYPES.keys())
    
    total_scraped = 0
    for idx, stype in enumerate(types_to_scrape):
        if idx > 0:
            print(f"Waiting {args.delay} seconds before switching station type...")
            time.sleep(args.delay)
        total_scraped += scrape_station_type(stype, args.pages, args.delay, args.outdir)
        
    print(f"\nScrape job completed. Total records parsed/scraped: {total_scraped}")

if __name__ == "__main__":
    main()
