#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A set of functions for scraping websites
"""

import requests
import re
import json
import bilbystats as bs
import asyncio
import pandas as pd
from joblib import Parallel, delayed
from tqdm import tqdm

def scrape(url, code = None, proxies = None):
    if proxies:
        response = requests.get(url, proxies=proxies)
    else:
        response = requests.get(url)
        
    if code == "utf-8":
        response.encoding = "utf-8"

    if response.status_code == 403:
        headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/119.0 Safari/537.36",
    "Accept-Language": "en-US,en;q=0.9",
    }

        response = requests.get(url, headers=headers)
        out = response.text
    elif response.status_code == 200:
        out = response.text
    else:
        print(f"Failed to retrieve the page. Status code: {response.status_code}")

    return out


def parse_html(html_content, element_str, normalize = False):
    from bs4 import BeautifulSoup

    # If you have the HTML as a string
    soup = BeautifulSoup(html_content, 'html.parser')

    # Extract using CSS selectors
    elements = soup.select(element_str)
    extracted_text = [element.get_text(strip=True) for element in elements]

    if normalize:
        extracted_text = [normalize_spaces(t) for t in extracted_text]

    return extracted_text

def normalize_spaces(text):
    # Replace all Unicode whitespace with a normal space
    return re.sub(r"\s+", " ", text)

def clean_html_content(
    html: str,
) -> tuple[str, int, int]:
    """
    Extract clean text from HTML.

    Args:
        html: Raw HTML content

    Returns:
        Tuple of (cleaned_text, character_count, word_count)
    """
    from bs4 import BeautifulSoup

    soup = BeautifulSoup(html, "html.parser")

    # Remove script, style, nav, header, footer elements
    for element in soup(["script", "style", "nav", "header", "footer", "aside"]):
        element.decompose()

    # Get text
    text = soup.get_text(separator="\n", strip=True)

    # Clean up multiple newlines
    lines = [line.strip() for line in text.split("\n") if line.strip()]

    cleaned_text = "\n".join(lines)

    return cleaned_text

def get_title(html):
    from bs4 import BeautifulSoup

    soup = BeautifulSoup(html, 'html.parser')
    
    # 1. Try Open Graph title (usually most accurate for articles)
    og_title = soup.find('meta', property='og:title')
    if og_title and og_title.get('content'):
        return og_title['content'].strip()
    
    # 2. Try Twitter title
    twitter_title = soup.find('meta', {'name': 'twitter:title'})
    if twitter_title and twitter_title.get('content'):
        return twitter_title['content'].strip()
    
    # 3. Try JSON-LD structured data
    json_ld = soup.find('script', type='application/ld+json')
    if json_ld:
        try:
            data = json.loads(json_ld.string)
            if isinstance(data, dict):
                if 'headline' in data:
                    return data['headline'].strip()
                if 'name' in data:
                    return data['name'].strip()
        except:
            pass
    
    # 4. Try H1 tag
    h1 = soup.find('h1')
    if h1:
        return h1.get_text().strip()
    
    # 5. Fall back to title tag
    title_tag = soup.find('title')
    if title_tag:
        return clean_title(title_tag.get_text().strip())
    
    return None

def clean_title(title):
    """Remove site name and other common suffixes from title tag"""
    # Common separators between article title and site name
    separators = [' - ', ' | ', ' — ', ' :: ', ' » ', ' · ']
    
    for separator in separators:
        if separator in title:
            # Usually the article title comes first
            parts = title.split(separator)
            # Return the longest part (often the actual title)
            return max(parts, key=len).strip()
    
    return title

def get_text(html):
    import trafilatura

    #downloaded = trafilatura.fetch_url(url)
    text = trafilatura.extract(html)
    
    return text

def get_publication_date(html):
    from bs4 import BeautifulSoup
    from dateutil import parser

    soup = BeautifulSoup(html, 'html.parser')
    
    # Check JSON-LD structured data first (most reliable)
    json_ld = soup.find('script', type='application/ld+json')
    if json_ld:
        try:
            data = json.loads(json_ld.string)
            if isinstance(data, dict):
                if 'datePublished' in data:
                    return parser.parse(data['datePublished'])
                # Check for nested article data
                if '@graph' in data:
                    for item in data['@graph']:
                        if 'datePublished' in item:
                            return parser.parse(item['datePublished'])
        except:
            pass
    
    # Check meta tags
    meta_tags = [
        ('property', 'article:published_time'),
        ('property', 'og:updated_time'),
        ('name', 'publish_date'),
        ('name', 'DC.date.issued'),
        ('itemprop', 'datePublished')
    ]
    
    for attr, value in meta_tags:
        meta = soup.find('meta', {attr: value})
        if meta and meta.get('content'):
            try:
                return parser.parse(meta['content'])
            except:
                continue
    
    # Check <time> elements
    time_elem = soup.find('time')
    if time_elem:
        # Try datetime attribute first
        if time_elem.get('datetime'):
            try:
                return parser.parse(time_elem['datetime'])
            except:
                pass
        # Try text content
        try:
            return parser.parse(time_elem.get_text())
        except:
            pass
    
    # Check common class names
    date_classes = ['publish-date', 'published-date', 'post-date', 
                    'entry-date', 'date', 'timestamp', 'article-date']
    
    for class_name in date_classes:
        date_elem = soup.find(class_=re.compile(class_name, re.I))
        if date_elem:
            try:
                # Clean the text and parse
                date_text = date_elem.get_text().strip()
                # Remove common prefixes
                date_text = re.sub(r'^(Published|Posted|Updated|Date)[:\s]*', 
                                  '', date_text, flags=re.I)
                return parser.parse(date_text)
            except:
                continue
    
    # Look for date patterns in the text
    # This is less reliable but can work as a fallback
    date_pattern = r'\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]* \d{1,2},? \d{4}\b'
    date_match = re.search(date_pattern, soup.get_text())
    if date_match:
        try:
            return parser.parse(date_match.group())
        except:
            pass
        
    try:
        a = soup.find("span", class_="date")
        return a.text
    except:
        pass
    
    return None

def get_author(html):
    from bs4 import BeautifulSoup

    soup = BeautifulSoup(html, 'html.parser')
    
    # Check meta tags
    meta_author = soup.find('meta', attrs={'name': 'author'})
    if meta_author:
        return meta_author.get('content')
    
    # Check Open Graph
    og_author = soup.find('meta', property='article:author')
    if og_author:
        return og_author.get('content')
    
    # Check common class names
    for class_name in ['author', 'byline', 'by-line', 'author-name']:
        author_elem = soup.find(class_=class_name)
        if author_elem:
            return author_elem.get_text().strip()
    
    # Check schema.org microdata
    author_elem = soup.find(attrs={'itemprop': 'author'})
    if author_elem:
        return author_elem.get_text().strip()
    
    return None

def get_meta_data(url = None, html = None):
    if not html:
        html = bs.scrape(url)
        
    author = get_author(html)
    date = get_publication_date(html)
    text = get_text(html)
    title = get_title(html)
    
    out = {'title': title, 'text': text, 'publication_date': date, 'author': author, 'url': url}
    
    return out

def inverse_scrape(html, savename, savedir = '/Users/samd/Documents/General_code/Scraping/article_pages/', model = 'anthropic/claude-sonnet-4', prompt = None):
    if not prompt:
        prompt = '''Extract the following fields from this HTML using Python and BeautifulSoup:
- date (publication/updated date in the form YYYY-MM-DD)
- text (main body content, cleaned of HTML tags)
- title (article/page title)
- author (author name)

Requirements:
- Use robust selectors that handle missing elements gracefully
- Return None for fields that cannot be found
- Clean and normalize extracted text (strip whitespace, remove extra newlines)
- Return results as a dictionary with keys: 'date', 'text', 'title', 'author', 'source'

Call the function used to extract the fields: extract_fields. 

Provide only the Python code without explanations. Begin directly with the code - no backticks.'''

    prompt = html + ' ' + prompt

    code = bs.llm_api(prompt, '', model, max_tokens = 10000)
    code = process_python_string(code)
    
    with open(savedir + savename + '.py', "w") as file:
        file.write(code)
        
    return code

async def inverse_scrape_async(htmls, savenames, savedir = '/Users/samd/Documents/General_code/Scraping/article_pages/', model = 'anthropic/claude-sonnet-4', prompt = None):
    if not prompt:
        prompt = '''Extract the following fields from the provided html using Python and BeautifulSoup:
- date (publication/post date)
- text (main body content, cleaned of HTML tags)
- title (article/page title)
- author (author name)

Requirements:
- Use robust selectors that handle missing elements gracefully
- Return None for fields that cannot be found
- Clean and normalize extracted text (strip whitespace, remove extra newlines)
- Return results as a dictionary with keys: 'date', 'text', 'title', 'author', 'source'

Call the function used to extract the fields: extract_fields. 

Provide only the Python code without explanations. Begin directly with the code - no backticks.

Don't include an example of the code being implemented.
'''

    listofcode = await bs.llm_api_async(htmls, prompt, model, max_tokens = 10000)
    
    for i in range(len(listofcode)):
        code = listofcode[i]
        code = process_python_string(code)
    
        with open(savedir + savenames[i] + '.py', "w") as file:
            file.write(code)
        
    return code

def import_module(savename, savedir = '/Users/samd/Documents/General_code/Scraping/article_pages/'):
    import importlib.util  
    spec = importlib.util.spec_from_file_location(savename, savedir + savename + '.py')
    xinhua = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(xinhua)

    return xinhua

def get_proxy_url(url):
    api_key = bs.read_api_key('scraper_ops')
    scraper_ops_url = "https://proxy.scrapeops.io/v1/?api_key=" + api_key + "&url={0}"
    proxy_url = scraper_ops_url.format(url)
    return proxy_url

def get_proxy_urls(urls):
    out = [get_proxy_url(url) for url in urls]
    return out

def extract_links(html):
    #extract_urls
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(html, "html.parser")
    links = [a.get("href") for a in soup.find_all("a", href=True)]
    return links


async def fetch_html(session, url, semaphore):
    async with semaphore:
        try:
            async with session.get(url) as response:
                return await response.text()
        except Exception as e:
            print(f"Error fetching {url}: {e}")
            return None

async def scrape_urls_orig(urls, batch_size=50):
    import aiohttp
    
    semaphore = asyncio.Semaphore(batch_size)
    async with aiohttp.ClientSession() as session:
        tasks = [fetch_html(session, url, semaphore) for url in urls]
        html_pages = await asyncio.gather(*tasks)
        return html_pages
    
async def scrape_urls(urls, batch_size=50):
    import aiohttp
    from tqdm.asyncio import tqdm
    
    semaphore = asyncio.Semaphore(batch_size)
    async with aiohttp.ClientSession() as session:
        tasks = [bs.fetch_html(session, url, semaphore) for url in urls]
        html_pages = await tqdm.gather(
            *tasks,
            desc="Scraping URLs",
            total=len(urls)
        )
        return html_pages
    
def process_python_string(code):
    if code[:9] == "```python":
        code = code[9:]
        
    if code[:10] == "\n```python":
        code = code[10:]
    
    if code[-3:] == "```":
        code = code[:-3]

    if code[-1:] == "\n":
        code = code[:-1]

    if code[-3:] == "```":
        code = code[:-3]

    if code[:1] == "\n":
        code = code[1:]
        
    return code


def browser_scrape(url):
    from selenium import webdriver
    from selenium.webdriver.chrome.options import Options
    from selenium.webdriver.chrome.service import Service
    from webdriver_manager.chrome import ChromeDriverManager
    import time
    """Download HTML using Selenium with visible Chrome browser"""
    
    # Set up Chrome options
    chrome_options = Options()
    chrome_options.add_argument("--start-maximized")  # Start maximized
    chrome_options.add_argument("--disable-blink-features=AutomationControlled")  # Avoid bot detection
    chrome_options.add_experimental_option("excludeSwitches", ["enable-automation"])
    chrome_options.add_experimental_option('useAutomationExtension', False)
    
    # Create driver - this will open a visible browser window
    print("Starting Chrome browser...")
    driver = webdriver.Chrome(
        service=Service(ChromeDriverManager().install()), 
        options=chrome_options
    )
    
    try:
        print(f"\nNavigating to: {url}")
        driver.get(url)
        
        print("Waiting for page to load...")
        time.sleep(5)  # Wait for page to fully load (adjust if needed)
        
        # Get the HTML source
        html = driver.page_source
        
        # Save to file
        filename = "downloaded_page.html"
        with open(filename, "w", encoding="utf-8") as f:
            f.write(html)
        
        print(f"\n✓ HTML downloaded successfully!")
        print(f"✓ Saved to: {filename}")
        print(f"✓ Page title: {driver.title}")
        print(f"✓ HTML length: {len(html):,} characters")
        
        return html
        
    except Exception as e:
        print(f"\n✗ Error occurred: {e}")
        raise
        
    finally:
        print("\nClosing browser...")
        driver.quit()

def browser_scrape_many(urls, max_workers=4):
    from concurrent.futures import ThreadPoolExecutor, as_completed
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(browser_scrape, url): url for url in urls}
        for future in as_completed(futures):
            try:
                results.append(future.result())
            except Exception as e:
                print(f"Error scraping {futures[future]}: {e}")
    return results

def federal_register_json(date: str) -> dict:
    """
    Fetch a Federal Register issue for a given date.
    
    Args:
        date: Date string in 'YYYY-MM-DD' format (e.g., '2026-02-02')
    
    Returns:
        JSON response as a dictionary
    """
    url = f"https://www.federalregister.gov/api/v1/issues/{date}.json"
    response = requests.get(url, headers={"accept": "*/*"})
    response.raise_for_status()
    return response.json()

def federal_register_docs(document_numbers: list[str], fields: list[str]) -> dict:
    ids = ",".join(document_numbers)
    url = f"https://www.federalregister.gov/api/v1/documents/{ids}.json"
    params = {"fields[]": fields}
    response = requests.get(url, params=params, headers={"accept": "*/*"})
    response.raise_for_status()
    return response.json()


def federal_register(date: str) -> pd.DataFrame:
    """
    Fetch a Federal Register issue for a given date.

    Args:
        date: Date string in 'YYYY-MM-DD' format (e.g., '2026-02-02')

    Returns:
        DataFrame with one row per document
    """
    url = f"https://www.federalregister.gov/api/v1/issues/{date}.json"
    response = requests.get(url, headers={"accept": "*/*"})
    response.raise_for_status()
    data = response.json()

    rows = []
    for agency in data["agencies"]:
        for cat in agency.get("document_categories", []):
            for doc in cat.get("documents", []):
                for doc_number in doc.get("document_numbers", []):
                    rows.append({
                        "publication_date": data["meta"]["publication_date"],
                        "agency": agency["name"],
                        "agency_slug": agency["slug"],
                        "document_type": cat["type"],
                        "subject_1": doc.get("subject_1"),
                        "subject_2": doc.get("subject_2"),
                        "document_number": doc_number,
                    })

    return pd.DataFrame(rows)
