#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A set of functions for performing querying of the database
"""
import json
import requests
from google.cloud import bigquery
import pandas as pd
import bilbystats as bs
from typing import Optional, List, Dict, Tuple
from datetime import datetime, timedelta
import aiohttp
import asyncio

def gbq(query: str, project: str = 'bilbyai-prod'):
    """
    Execute a SQL query in Google BigQuery and return the results as a DataFrame.

    This function initializes a BigQuery client using the specified project and 
    executes the provided SQL query. The result is returned as a pandas DataFrame.
    In case of failure, the function catches exceptions and prints an error message.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    query : str
        The SQL query to be executed in BigQuery.
    project : str
        The Google Cloud project ID under which the query will be executed.

    ---------------------------------------------------------------------------
    OUTPUT:
    df : pd.DataFrame or None
        A pandas DataFrame containing the results of the query if successful; 
        None if an error occurs during execution.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Initialize the BigQuery client
    # It will automatically use the credentials from gcloud auth application-default login
    client = bigquery.Client(project=project)

    try:
        # Execute the query
        query_job = client.query(query)

        # Wait for the job to complete and get results
        results = query_job.result()

        # Convert to pandas DataFrame for easier handling
        df = results.to_dataframe()

        return df

    except Exception as e:
        print(f"Error executing query: {str(e)}")
        return None


def wayback(url: str):
    """
    Retrieve archived snapshots of a given URL from the Wayback Machine.

    This function queries the Internet Archive's CDX API to fetch historical 
    snapshots of a specified URL. It returns a list of snapshot metadata 
    including timestamp, readable date, status code, MIME type, size in bytes, 
    and both standard and raw Wayback URLs.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    url : str
        The URL for which to retrieve archived snapshots from the Wayback Machine.

    ---------------------------------------------------------------------------
    OUTPUT:
    snapshots : list of dict
        A list of dictionaries, each containing metadata for a snapshot:
            - timestamp : str
                The original timestamp of the snapshot in YYYYMMDDHHMMSS format.
            - date : str
                Human-readable date and time derived from the timestamp.
            - status_code : str
                HTTP status code returned by the archived page (only 200 included).
            - mime_type : str
                MIME type of the archived content.
            - size_bytes : int
                Size of the archived snapshot in bytes.
            - wayback_url : str
                URL to view the snapshot on the Wayback Machine.
            - raw_url : str
                URL to the raw archived content.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    cdx_url = "http://web.archive.org/cdx/search/cdx"

    params = {
        'url': url,
        'output': 'json',
        'filter': 'statuscode:200'  # Only successful captures
    }

    try:
        response = requests.get(cdx_url, params=params, timeout=30)
        response.raise_for_status()

        data = response.json()

        if not data:
            print("No snapshots found!")
            return []

        # First row contains field names
        fields = data[0]
        snapshots = []

        for row in data[1:]:
            snapshot = dict(zip(fields, row))

            # Convert timestamp to readable date
            timestamp = snapshot['timestamp']
            readable_date = datetime.strptime(
                timestamp, '%Y%m%d%H%M%S').strftime('%Y-%m-%d %H:%M:%S')

            # Generate Wayback URLs
            wayback_url = f"https://web.archive.org/web/{timestamp}/{snapshot['original']}"
            raw_url = f"https://web.archive.org/web/{timestamp}id_/{snapshot['original']}"

            snapshots.append({
                'timestamp': timestamp,
                'date': readable_date,
                'status_code': snapshot['statuscode'],
                'mime_type': snapshot['mimetype'],
                'size_bytes': int(snapshot.get('length', 0)),
                'wayback_url': wayback_url,
                'raw_url': raw_url
            })

        return snapshots

    except requests.RequestException as e:
        print(f"Error fetching snapshots: {e}")
        return []

async def wayback_async(
    url: str,
    limit: Optional[int] = None,
    collapse: Optional[str] = None,
    timeout: int = 30,
    http = 0
) -> Tuple[List[Dict], bool]:
    """
    Retrieve archived snapshots of a given URL from the Wayback Machine (async version).
    
    This function queries the Internet Archive's CDX API to fetch historical 
    snapshots of a specified URL. It returns a list of snapshot metadata 
    including timestamp, readable date, status code, MIME type, size in bytes, 
    and both standard and raw Wayback URLs.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    url : str
        The URL for which to retrieve archived snapshots from the Wayback Machine.
    limit : int, optional
        Maximum number of snapshots to retrieve from the API.
    collapse : str, optional
        Collapse results by field (e.g., 'timestamp:8' for daily snapshots).
    timeout : int, default=30
        Request timeout in seconds.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    snapshots : list of dict
        A list of dictionaries, each containing metadata for a snapshot:
            - timestamp : str
                The original timestamp of the snapshot in YYYYMMDDHHMMSS format.
            - date : str
                Human-readable date and time derived from the timestamp.
            - status_code : str
                HTTP status code returned by the archived page (only 200 included).
            - mime_type : str
                MIME type of the archived content.
            - size_bytes : int
                Size of the archived snapshot in bytes.
            - wayback_url : str
                URL to view the snapshot on the Wayback Machine.
            - raw_url : str
                URL to the raw archived content.
    failed : bool
        True if the request failed, False if successful.
    
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if http == 1:
        cdx_url = "http://web.archive.org/cdx/search/cdx:81"
    else: 
        cdx_url = "https://web.archive.org/cdx/search/cdx"

    params = {
        'url': url,
        'output': 'json',
        'filter': 'statuscode:200'
    }
    
    if limit:
        params['limit'] = limit
    if collapse:
        params['collapse'] = collapse
    
    try:
        timeout_obj = aiohttp.ClientTimeout(total=timeout)
        async with aiohttp.ClientSession(timeout=timeout_obj) as session:
            async with session.get(cdx_url, params=params) as response:
                response.raise_for_status()
                data = await response.json()
                
                if not data:
                    print("No snapshots found!")
                    return [], False
                
                # First row contains field names
                fields = data[0]
                
                # Pre-calculate field indices for performance
                timestamp_idx = fields.index('timestamp')
                statuscode_idx = fields.index('statuscode')
                mimetype_idx = fields.index('mimetype')
                length_idx = fields.index('length')
                original_idx = fields.index('original')
                
                snapshots = [
                    {
                        'timestamp': row[timestamp_idx],
                        'date': datetime.strptime(
                            row[timestamp_idx], '%Y%m%d%H%M%S'
                        ).strftime('%Y-%m-%d %H:%M:%S'),
                        'status_code': row[statuscode_idx],
                        'mime_type': row[mimetype_idx],
                        'size_bytes': int(row[length_idx] or 0),
                        'wayback_url': f"https://web.archive.org/web/{row[timestamp_idx]}/{row[original_idx]}",
                        'raw_url': f"https://web.archive.org/web/{row[timestamp_idx]}id_/{row[original_idx]}"
                    }
                    for row in data[1:]
                ]
                
                return snapshots, False
                
    except aiohttp.ClientError as e:
        #print(f"Error fetching snapshots for {url}: {e}")
        return [], True
    except Exception as e:
        #print(f"Unexpected error for {url}: {e}")
        return [], True


async def wayback_multiple_async(urls: List[str], http = 0) -> Tuple[Dict[str, List[Dict]], List[str]]:
    """
    Retrieve snapshots for multiple URLs concurrently.
    
    ARGUMENTS:
    urls : list of str
        List of URLs to fetch snapshots for.
    **kwargs : optional
        Additional arguments to pass to wayback_async (limit, collapse, timeout).
    
    OUTPUT:
    results : dict
        Dictionary mapping each URL to its list of snapshots.
    failed_urls : list of str
        List of URLs that failed to fetch snapshots.
    """
    tasks = [wayback_async(url, http = http) for url in urls]
    results = await asyncio.gather(*tasks)
    
    # Separate successful results from failed URLs
    snapshots_dict = {}
    failed_urls = []
    
    for url, (snapshots, failed) in zip(urls, results):
        if failed:
            failed_urls.append(url)
        else:
            snapshots_dict[url] = snapshots
    
    return snapshots_dict, failed_urls

async def wayback_multiple_async_retry(urls: List[str], max_retries: int = 2, http = 0) -> Dict[str, List[Dict]]:
    """
    Retry fetching snapshots for multiple URLs, alternating between HTTP/HTTPS on failures.
    
    ARGUMENTS:
    urls : list of str
        List of URLs to fetch snapshots for.
    max_retries : int, default=2
        Maximum number of retry attempts.
    **kwargs : optional
        Additional arguments to pass to wayback_multiple_async (limit, collapse, timeout).
    
    OUTPUT:
    snapshots_dict : dict
        Dictionary mapping each URL to its list of snapshots.
    """
    urls2do = urls.copy()
    snapshots_dict = {}
    retry_count = 0
    
    while len(urls2do) > 0 and retry_count < max_retries:
        # Fetch snapshots for remaining URLs
        results, failed_urls = await wayback_multiple_async(urls2do, http = http)
        
        # Add successful results to snapshots_dict
        snapshots_dict.update(results)
        
        # Update URLs to retry (these are the failed ones)
        urls2do = failed_urls
        
        # Alternate between HTTPS and HTTP
        http = 1 - http
        retry_count += 1
        
        if urls2do:
            print(f"Retrying {len(urls2do)} failed URLs with {'HTTP' if http else 'HTTPS'}...")
    
    # For any remaining failed URLs, set empty list
    for url in urls2do:
        if url not in snapshots_dict:
            snapshots_dict[url] = []
    
    return snapshots_dict, urls2do

async def wayback_snapshot2url_org(urls, snapshots_dict):
    setofdfs = []
    for i in range(len(urls)):
        print(i)
        url = urls[i]
        out = snapshots_dict[url]
        if len(out) > 0:
            set_of_urls = [a['wayback_url'] for a in out]
            proxy_urls = [bs.get_proxy_url(url) for url in set_of_urls]
            wayback_htmls = await bs.scrape_urls(proxy_urls)
            fun = lambda html: bs.get_meta_data(html=html)['text']
            article_texts = bs.parallelize(wayback_htmls, fun)

            out_df = pd.DataFrame(out)
            out_df['date'] = pd.to_datetime(out_df['date']).dt.date
            d = {"url": url,"date": out_df['date'], "text": article_texts, "wayback_url": out_df['wayback_url'], "wayback_html": wayback_htmls}
            setofdfs.append(pd.DataFrame(d))
        else:
            df = pd.DataFrame(columns=['url', 'date', 'text', 'wayback_url', 'wayback_html'])
            setofdfs.append(df)
        
    wayback_df = pd.concat(setofdfs)
    wayback_df.reset_index(drop=True, inplace=True)
    wayback_df

    return wayback_df

async def scrape_wayback(urls, snapshots):
    urls2scrape = []
    setofdfs = []
    for i in range(len(urls)):
        url = urls[i]
        out = snapshots[url]
        if len(out) > 0:
            set_of_urls = [a['wayback_url'] for a in out]
            proxy_urls = [bs.get_proxy_url(url) for url in set_of_urls]
            urls2scrape += proxy_urls
            out_df = pd.DataFrame(out)
            out_df['wayback_date'] = pd.to_datetime(out_df['date']).dt.date
            out_df['url'] = url
            setofdfs.append(out_df)
    
    combo_df = pd.concat(setofdfs)
    
    print('Scraping wayback urls')
    wayback_htmls = await bs.scrape_urls(urls2scrape)
    
    print('Parsing htmls into article text')
    fun = lambda html: bs.get_meta_data(html=html)['text']
    article_texts = bs.parallelize(wayback_htmls, fun)
    
    d = {"url": combo_df['url'],"wayback_date": combo_df['wayback_date'], "text": article_texts, "wayback_url": combo_df['wayback_url']}
    wayback_df = pd.DataFrame(d)
    wayback_df.reset_index(drop=True, inplace=True)
    
    return wayback_df

def get_dates_between(start_date: str, end_date: str):
    """
    Returns a list of dates between start_date and end_date (inclusive).
    
    Parameters:
        start_date (str): "YYYY-MM-DD"
        end_date (str): "YYYY-MM-DD"
        
    Returns:
        List[str]: Dates in "YYYY-MM-DD" format
    """
    start = datetime.strptime(start_date, "%Y-%m-%d").date()
    end = datetime.strptime(end_date, "%Y-%m-%d").date()
    
    if start > end:
        raise ValueError("start_date must be before or equal to end_date")
    
    dates = []
    current = start
    while current <= end:
        dates.append(current.isoformat())
        current += timedelta(days=1)
    
    return dates
