import os
import sqlite3
import requests
import re
import concurrent.futures  # Import for concurrency
import multiprocessing
import time  # For rate limiting
import logging  # For logging
from bs4 import BeautifulSoup
from urllib.parse import urlparse, urljoin
import asyncio # Import asyncio if you intend to use key checker functions later

# Import key checking functions from the other file
from deepseek_key_checker import APIKey, check_keys, pretty_print_deepseek_keys, fetch_keys_from_db, update_key_status_in_db

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

DATABASE_NAME = "findings_db.sqlite3"  # Changed database name to 'findings_db' to be more general
COOKIES_DATABASE_NAME = "cookies.db"

# WARNING: Consider using environment variables or secure configuration for API token in production
GITHUB_API_TOKEN = os.environ.get("GITHUB_API_TOKEN") # Use environment variable - recommended
# GITHUB_API_TOKEN = "YOUR_GITHUB_API_TOKEN_HERE" # If you must hardcode for testing (INSECURE!)
GITHUB_API_HEADERS = {"User-Agent": "Mozilla/5.0"}
if GITHUB_API_TOKEN:
    GITHUB_API_HEADERS["Authorization"] = f"token {GITHUB_API_TOKEN}"

RATE_LIMIT_DELAY = 1  # Reduced initial delay, can adjust dynamically if needed - important for concurrency to be effective
SEARCH_RATE_LIMIT_DELAY = 2 # Separate rate limit for search to be less aggressive


def init_db():
    """Initializes the SQLite database for found findings (keys, strings, etc.)."""
    conn = sqlite3.connect(DATABASE_NAME)
    c = conn.cursor()
    c.execute("""
        CREATE TABLE IF NOT EXISTS findings (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            url TEXT UNIQUE,
            finding_type TEXT,  -- e.g., 'api_key', 'go_ty_string'
            finding_value TEXT,
            context_url TEXT,
            checked INTEGER DEFAULT 0,    -- Added for checker
            balance TEXT,               -- Added for checker
            available INTEGER DEFAULT 0,  -- Added for checker
            rate_limited INTEGER DEFAULT 0, -- Added for checker
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.commit()
    conn.close()
    logging.info(f"Database '{DATABASE_NAME}' initialized.")


def init_cookies_db():
    """Initializes the SQLite database for cookies (optional)."""
    conn = sqlite3.connect(COOKIES_DATABASE_NAME)
    c = conn.cursor()
    c.execute("""
        CREATE TABLE IF NOT EXISTS cookies (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT UNIQUE,
            value TEXT
        )
    """)
    conn.commit()
    conn.close()
    logging.info(f"Cookies database '{COOKIES_DATABASE_NAME}' initialized.")


def load_netscape_cookies(cookie_file):
    """Loads cookies from a Netscape format cookie file."""
    cookies = {}
    try:
        with open(cookie_file, "r") as f:
            for line in f:
                if line.startswith("#") or not line.strip():
                    continue
                parts = line.strip().split("\t")
                if len(parts) >= 7:
                    cookies[parts[5]] = parts[6]
                    save_cookie(parts[5], parts[6])
        logging.info(f"Cookies loaded from '{cookie_file}'.")
    except FileNotFoundError:
        logging.warning(f"Cookie file '{cookie_file}' not found. Proceeding without cookies.")
    return cookies


def save_cookie(name, value):
    """Saves a cookie to the cookies database (optional)."""
    conn = sqlite3.connect(COOKIES_DATABASE_NAME)
    c = conn.cursor()
    try:
        c.execute("INSERT OR IGNORE INTO cookies (name, value) VALUES (?, ?)", (name, value))
        conn.commit()
        logging.debug(f"Cookie '{name}' saved to cookies database.")
    except sqlite3.Error as e:
        logging.error(f"Database error saving cookie: {e}")
    finally:
        conn.close()


def search_github_repos(query, cookies, max_pages=10):
    """Searches GitHub code and extracts repository URLs."""
    base_url = "https://github.com/search?q={}&type=code&p={}"
    headers = {"User-Agent": "Mozilla/5.0"}
    repo_urls = set()

    for page in range(1, max_pages + 1):
        url = base_url.format(query, page)
        logging.info(f"Searching page {page} for repos: {url}")
        try:
            response = requests.get(url, headers=headers, cookies=cookies, timeout=10)
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            logging.error(f"Request error on page {page}: {e}")
            break

        soup = BeautifulSoup(response.text, "html.parser")
        found_on_page = 0

        for link in soup.find_all("a", href=True):
            href = link["href"]
            if "/blob/" in href:
                repo_url_part = href.split("/blob/")[0]
                repo_url = f"https://github.com{repo_url_part}"  # Corrected to get the repo URL not blob URL
                if repo_url not in repo_urls:
                    repo_urls.add(repo_url)
                    found_on_page += 1

        logging.info(f"Page {page}: Found {found_on_page} new repo URLs, Total unique repos: {len(repo_urls)}")
        if found_on_page == 0:
            logging.info(f"No new repo URLs on page {page}. Stopping search.")
            break
        time.sleep(SEARCH_RATE_LIMIT_DELAY) # Use separate rate limit for search

    return list(repo_urls)


def crawl_repository_for_findings(repo_url, key_pattern=r"sk-[A-Za-z0-9]{32}", string_to_find="go ty"):
    """Crawls a GitHub repository and extracts API keys and specific strings."""
    owner_repo_match = re.search(r"github\.com/([^/]+)/([^/]+)", repo_url)
    if not owner_repo_match:
        logging.warning(f"Invalid repository URL: {repo_url}")
        return {"api_keys": [], "go_ty_strings": []}  # Return dict to hold both types of findings

    owner, repo = owner_repo_match.groups()
    api_url = f"https://api.github.com/repos/{owner}/{repo}/git/trees/HEAD?recursive=1"

    findings = {"api_keys": [], "go_ty_strings": []} # Dictionary to store findings
    try:
        response = requests.get(api_url, headers=GITHUB_API_HEADERS, timeout=20)
        response.raise_for_status()
        repo_content = response.json()

        if 'truncated' in repo_content and repo_content['truncated']:
            logging.warning(f"Repository tree for {repo_url} is truncated. May not get all files.")

        files = [item for item in repo_content.get('tree', []) if item['type'] == 'blob']
        logging.info(f"Crawling repository: {repo_url}, Files to check: {len(files)}")

        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor: # Increased max_workers for concurrency within repo
            file_futures = [executor.submit(process_file, file_info, owner, repo, repo_url, key_pattern, string_to_find) for file_info in files]
            for future in concurrent.futures.as_completed(file_futures):
                file_findings = future.result() # Get findings from each file processing
                findings["api_keys"].extend(file_findings["api_keys"])
                findings["go_ty_strings"].extend(file_findings["go_ty_strings"])


    except requests.exceptions.RequestException as e:
        logging.error(f"Error accessing repository tree for {repo_url}: {e}")
        return findings  # Return findings even if repo crawl fails partially

    logging.info(f"Found {len(findings['api_keys'])} API keys and {len(findings['go_ty_strings'])} 'go ty' strings in repository: {repo_url}")
    return findings


def process_file(file_info, owner, repo, repo_url, key_pattern, string_to_find):
    """Processes a single file to find API keys and strings."""
    file_path = file_info['path']
    raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/HEAD/{file_path}"
    file_findings = {"api_keys": [], "go_ty_strings": []} # Findings for this file

    try:
        raw_response = requests.get(raw_url, timeout=15)
        raw_response.raise_for_status()
        file_content = raw_response.text

        # Find API Keys
        api_keys = re.findall(key_pattern, file_content)
        for key in api_keys:
            if save_finding_repo_crawl(repo_url, raw_url, key, "api_key"):  # Save as 'api_key' type
                file_findings["api_keys"].append(key)

        # Find "go ty" strings (case-insensitive)
        go_ty_matches = re.findall(re.escape(string_to_find), file_content, re.IGNORECASE) # Escape string and use case-insensitive search
        for match in go_ty_matches:
            if save_finding_repo_crawl(repo_url, raw_url, string_to_find, "go_ty_string"): # Save as 'go_ty_string' type
                file_findings["go_ty_strings"].append(string_to_find) # Storing the search string itself, not the match object

        time.sleep(RATE_LIMIT_DELAY) # Keep rate limit delay, but concurrency makes it less impactful on total time

    except requests.exceptions.RequestException as e:
        logging.warning(f"Error fetching raw content for {raw_url}: {e}")
    return file_findings # Return findings for this file


def save_finding_repo_crawl(repo_url, raw_url, finding_value, finding_type):
    """Saves a finding (key or string) found during repository crawl to the database."""
    conn = sqlite3.connect(DATABASE_NAME)
    c = conn.cursor()
    try:
        c.execute(
            "INSERT OR IGNORE INTO findings (url, finding_type, finding_value, context_url) VALUES (?, ?, ?, ?)",
            (raw_url, finding_type, finding_value, repo_url),
        )
        saved = c.rowcount > 0
        if saved:
            logging.debug(f"Finding '{finding_value}' (type: {finding_type}) from URL '{raw_url}' (repo: {repo_url}) saved to database.")
        return saved
    except sqlite3.Error as e:
        logging.error(f"Database error saving finding '{finding_value}' (type: {finding_type}) from URL '{raw_url}' (repo: {repo_url}): {e}")
        return False
    finally:
        conn.close()


def main():
    """Main function to execute GitHub search and repository crawl for findings."""
    init_db()
    # init_cookies_db()
    cookies = load_netscape_cookies("../scraping/cookies.txt")
    query = "deepseek AND /sk-[A-Za-z0-9]{32}/"

    # 1. Search GitHub and get repository URLs
    repo_search_results = search_github_repos(query, cookies)
    logging.info(f"Found {len(repo_search_results)} repository URLs from search.")

    # 2. Crawl repositories for findings (keys and "go ty") - CONCURRENT CRAWLING
    total_api_keys_crawled = 0
    total_go_ty_strings_crawled = 0
    all_api_keys = []  # To keep track of all API keys found
    all_go_ty_strings = [] # To keep track of all "go ty" strings found

    with concurrent.futures.ThreadPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor: # Use ThreadPoolExecutor for concurrent crawling
        futures = [executor.submit(crawl_repository_for_findings, repo_url, string_to_find="go ty") for repo_url in repo_search_results] # Submit crawl tasks

        for future in concurrent.futures.as_completed(futures): # Process results as they become available
            repo_findings = future.result()
            total_api_keys_crawled += len(repo_findings["api_keys"])
            total_go_ty_strings_crawled += len(repo_findings["go_ty_strings"])
            all_api_keys.extend(repo_findings["api_keys"])
            all_go_ty_strings.extend(repo_findings["go_ty_strings"])


    logging.info(f"Total API keys found during repository crawling: {total_api_keys_crawled}")
    logging.info(f"Total 'go ty' strings found during repository crawling: {total_go_ty_strings_crawled}")

    unique_api_keys_found = list(set(all_api_keys)) # Get unique keys for checking

    unique_go_ty_strings_found = len(set(all_go_ty_strings)) # Technically always 1, but kept for consistency if searching for multiple strings later

    print("-" * 50)
    print(f"Finding Extraction Summary:")
    print(f"  Total repositories crawled: {len(repo_search_results)}")
    print(f"  Total API keys found: {total_api_keys_crawled}")
    print(f"  Unique API keys found: {len(unique_api_keys_found)}") # Print count of unique keys
    print(f"  Total 'go ty' strings found: {total_go_ty_strings_crawled}")
    print(f"  Unique 'go ty' strings found (string searched): {unique_go_ty_strings_found}") # Clarified what unique strings means
    print("-" * 50)

    # 3. Check and Validate Unique API Keys (using functions from deepseek_key_checker.py)
    if unique_api_keys_found:
        logging.info("Starting API key validation...")
        api_key_objects = [APIKey(key) for key in unique_api_keys_found] # Create APIKey objects
        checked_keys = asyncio.run(check_keys(api_key_objects)) # Run async key check
        update_key_status_in_db(checked_keys, DATABASE_NAME) # Update status in DB, pass DB_NAME
        pretty_print_deepseek_keys(checked_keys) # Print formatted key check results
    else:
        logging.info("No unique API keys found to check.")


    logging.info("Script finished.")


if __name__ == "__main__":
    main()