import requests
import pandas as pd
import time
import os
import concurrent.futures
import threading
import asyncio
import aiohttp
from datasets import Dataset, DatasetDict
import datetime
from tqdm import tqdm
import json
from typing import List, Dict, Tuple, Optional, Any, Set


class RateLimiter:
    """Smart rate limiter for GitHub API requests that respects rate limit headers."""

    def __init__(self, calls_per_second=1):
        self.calls_per_second = 1  # Default to conservative rate
        self.last_call = 0
        self.lock = threading.Lock()
        self.remaining_calls = 60  # Default GitHub unauthenticated limit
        self.reset_time = time.time() + 3600  # Default reset after an hour
        self.authenticated = False
        self.core_limit = {}  # Track core API limits
        self.search_limit = {}  # Track search API limits
        self.semaphore = asyncio.Semaphore(10)  # Limit concurrent requests

        # Track resource-specific limits
        self.resource_limits = {
            "core": {"remaining": 60, "limit": 60, "reset": time.time() + 3600},
            "search": {"remaining": 10, "limit": 10, "reset": time.time() + 3600},
        }

    def wait(self):
        with self.lock:
            current_time = time.time()

            # If we're past the reset time, reset our counter
            if current_time > self.reset_time:
                self.remaining_calls = 60 if not self.authenticated else 5000

            # If we're out of calls, wait until reset
            if self.remaining_calls <= 5:  # Keep a small buffer
                wait_time = max(0, self.reset_time - current_time)
                if wait_time > 0:
                    print(
                        f"Rate limit reached. Waiting {wait_time:.1f} seconds for reset..."
                    )
                    time.sleep(wait_time + 2)  # Add 2 second buffer

            # Standard rate limiting
            time_since_last_call = current_time - self.last_call
            time_to_wait = max(0, 1 / self.calls_per_second - time_since_last_call)

            if time_to_wait > 0:
                time.sleep(time_to_wait)

            self.last_call = time.time()
            self.remaining_calls -= 1

    async def async_wait(self, resource="core"):
        """Async version of wait for use with aiohttp"""
        async with self.semaphore:
            # Use a context manager for the lock to ensure it's always released
            lock_acquired = False
            try:
                self.lock.acquire()
                lock_acquired = True

                current_time = time.time()

                # Get resource-specific limits
                res_limit = self.resource_limits.get(
                    resource,
                    {"remaining": 60, "limit": 60, "reset": time.time() + 3600},
                )

                # If we're past the reset time, reset our counter
                if current_time > res_limit["reset"]:
                    res_limit["remaining"] = res_limit["limit"]

                # If we're out of calls, wait until reset
                if res_limit["remaining"] <= 5:  # Keep a small buffer
                    wait_time = max(0, res_limit["reset"] - current_time)
                    if wait_time > 0:
                        print(
                            f"Rate limit reached for {resource}. Waiting {wait_time:.1f} seconds for reset..."
                        )
                        # Release the lock during the wait
                        self.lock.release()
                        lock_acquired = False

                        # Cap the wait time to avoid extremely long waits
                        wait_time = min(wait_time, 300)  # Max 5 minutes wait
                        await asyncio.sleep(wait_time + 2)  # Add 2 second buffer

                        self.lock.acquire()
                        lock_acquired = True

                # Standard rate limiting
                time_since_last_call = current_time - self.last_call
                time_to_wait = max(0, 1 / self.calls_per_second - time_since_last_call)

                if time_to_wait > 0:
                    # Release the lock during the wait
                    self.lock.release()
                    lock_acquired = False

                    await asyncio.sleep(time_to_wait)

                    self.lock.acquire()
                    lock_acquired = True

                self.last_call = time.time()
                res_limit["remaining"] -= 1
                self.resource_limits[resource] = res_limit
            finally:
                # Always release the lock if we acquired it
                if lock_acquired:
                    self.lock.release()

    def update_from_headers(self, headers):
        """Update rate limit info from GitHub API response headers"""
        with self.lock:
            # Check if we're authenticated based on rate limits
            if "X-RateLimit-Limit" in headers:
                limit = int(headers["X-RateLimit-Limit"])
                if limit > 60:
                    self.authenticated = True

            # Update remaining calls
            if "X-RateLimit-Remaining" in headers:
                self.remaining_calls = int(headers["X-RateLimit-Remaining"])

                # Print warning if we're running low
                if self.remaining_calls < 20:
                    print(
                        f"⚠️ Only {self.remaining_calls} API calls remaining before rate limit"
                    )

            # Update reset time
            if "X-RateLimit-Reset" in headers:
                self.reset_time = int(headers["X-RateLimit-Reset"])

            # Track which API we're using (core vs search)
            if "X-RateLimit-Resource" in headers:
                resource = headers["X-RateLimit-Resource"]
                if resource == "core":
                    self.core_limit = {
                        "limit": int(headers.get("X-RateLimit-Limit", 0)),
                        "remaining": int(headers.get("X-RateLimit-Remaining", 0)),
                        "reset": int(headers.get("X-RateLimit-Reset", 0)),
                    }
                    self.resource_limits["core"] = self.core_limit
                elif resource == "search":
                    self.search_limit = {
                        "limit": int(headers.get("X-RateLimit-Limit", 0)),
                        "remaining": int(headers.get("X-RateLimit-Remaining", 0)),
                        "reset": int(headers.get("X-RateLimit-Reset", 0)),
                    }
                    self.resource_limits["search"] = self.search_limit

    def update_from_aiohttp_headers(self, headers):
        """Update rate limit info from aiohttp response headers"""
        with self.lock:
            # Check if we're authenticated based on rate limits
            if "X-RateLimit-Limit" in headers:
                limit = int(headers["X-RateLimit-Limit"])
                if limit > 60:
                    self.authenticated = True

            # Update remaining calls
            if "X-RateLimit-Remaining" in headers:
                self.remaining_calls = int(headers["X-RateLimit-Remaining"])

            # Update reset time
            if "X-RateLimit-Reset" in headers:
                self.reset_time = int(headers["X-RateLimit-Reset"])

            # Track which API we're using (core vs search)
            if "X-RateLimit-Resource" in headers:
                resource = headers["X-RateLimit-Resource"]
                if resource == "core":
                    self.core_limit = {
                        "limit": int(headers.get("X-RateLimit-Limit", 0)),
                        "remaining": int(headers.get("X-RateLimit-Remaining", 0)),
                        "reset": int(headers.get("X-RateLimit-Reset", 0)),
                    }
                    self.resource_limits["core"] = self.core_limit
                elif resource == "search":
                    self.search_limit = {
                        "limit": int(headers.get("X-RateLimit-Limit", 0)),
                        "remaining": int(headers.get("X-RateLimit-Remaining", 0)),
                        "reset": int(headers.get("X-RateLimit-Reset", 0)),
                    }
                    self.resource_limits["search"] = self.search_limit


def get_github_token():
    """Get GitHub token from environment or file"""
    # First try environment variable
    token = os.environ.get("GITHUB_TOKEN")

    # Then try from a token file
    if not token:
        token_paths = [
            os.path.expanduser("~/.github_token"),
            os.path.expanduser("~/.github/token"),
            os.path.expanduser("~/.config/github/token"),
        ]

        for token_file in token_paths:
            if os.path.exists(token_file):
                try:
                    with open(token_file, "r") as f:
                        token = f.read().strip()
                        if token:
                            break
                except:
                    pass

    # Clean the token if it exists
    if token:
        # Remove any "Bearer " or "token " prefix if present
        if token.lower().startswith(("bearer ", "token ")):
            token = token.split(" ", 1)[1]
        # Remove any quotes
        token = token.strip("\"'")

    return token


def get_github_topics(limit=30):
    """
    Fetch a list of popular Python-related GitHub topics.

    Parameters:
    - limit: Maximum number of topics to retrieve
    - github_token: GitHub personal access token (recommended to avoid rate limits)
    - rate_limiter: RateLimiter instance to control request rate

    Returns:
    - Dataset with topics column containing list of Python-related topics
    """
    # Setup API request headers
    # Get token if not provided
    github_token = get_github_token()
    rate_limiter = None
    headers = {"Accept": "application/vnd.github.v3+json"}

    # Get token if not provided
    if not github_token:
        github_token = get_github_token()

    if github_token:
        headers["Authorization"] = f"Bearer {github_token}"
        print("Using GitHub authentication token")
    else:
        print(
            "WARNING: No GitHub token found. API rate limits will be severely restricted."
        )

    rate_limiter = RateLimiter(calls_per_second=0.5)

    base_url = "https://api.github.com/search/topics"

    all_topics = []
    page = 1
    per_page = min(100, limit)  # GitHub API allows max 100 per page
    remaining = limit

    print(f"Fetching up to {limit} Python-related GitHub topics...")

    # Try different search queries if one fails
    search_queries = [
        "python",  # Simple search for python
        "topic:python",  # Explicit topic search
        "language:python",  # Language-based search
        "python framework",  # Python frameworks
        "python library",  # Python libraries
    ]

    for query in search_queries:
        if len(all_topics) >= limit:
            break

        print(f"Trying search query: {query}")

        while remaining > 0:
            # Search for Python-related topics
            params = {
                "q": query,
                "per_page": min(per_page, remaining),
                "page": page,
            }
            if rate_limiter:
                rate_limiter.wait()

            response = requests.get(base_url, headers=headers, params=params)

            # Print response details for debugging
            print(f"Response status: {response.status_code}")

            # Update rate limiter with response headers
            if rate_limiter and response.headers:
                rate_limiter.update_from_headers(response.headers)

            if response.status_code == 403:
                reset_time = int(
                    response.headers.get("X-RateLimit-Reset", time.time() + 3600)
                )
                wait_time = max(0, reset_time - time.time())
                print(f"Rate limit exceeded. Reset in {wait_time:.1f} seconds.")
                if (
                    wait_time > 0 and wait_time < 3600
                ):  # Only wait if reasonable (<1 hour)
                    print(f"Waiting for rate limit reset...")
                    time.sleep(wait_time + 2)  # Add buffer
                    continue
                else:
                    print("Rate limit wait time too long. Please use a GitHub token.")
                    break
            elif response.status_code != 200:
                print(f"Error fetching topics: Status {response.status_code}")
                print(f"Response: {response.text}")
                break

            data = response.json()
            items = data.get("items", [])

            print(f"Found {len(items)} items in response")

            if not items:
                break

            # Less strict filtering - include more Python-related topics
            for item in items:
                topic_name = item.get("name", "").lower()
                if topic_name and topic_name not in [t.lower() for t in all_topics]:
                    all_topics.append(item["name"])

                    # Print each topic we're adding
                    print(f"Adding topic: {item['name']}")

                    if len(all_topics) >= limit:
                        break

            fetched = len(items)
            remaining -= fetched

            if (
                fetched < per_page or len(all_topics) >= limit
            ):  # Less than requested or we have enough
                break

            page += 1
            time.sleep(1)  # Be nice to the API

        # Reset for next query
        page = 1

    # If we still don't have topics, add some common Python topics manually
    if not all_topics:
        print("No Python topics found via API. Using fallback list.")
        all_topics = [
            "python",
            "django",
            "flask",
            "fastapi",
            "pandas",
            "numpy",
            "pytorch",
            "tensorflow",
            "scikit-learn",
            "matplotlib",
            "data-science",
            "machine-learning",
            "web-development",
            "api",
            "automation",
            "scraping",
            "nlp",
            "deep-learning",
        ]
        # Only take up to the limit
        all_topics = all_topics[:limit]

    # Create dataset with topics column
    dataset = Dataset.from_dict({"topics": all_topics})  # Each topic as a separate row

    print(f"Retrieved {len(all_topics)} Python-related topics")
    return dataset


async def get_repositories_async(
    search_query, max_repos, headers, rate_limiter, max_retries=3
):
    """Get repositories matching the search query using async."""
    params = {
        "q": search_query,
        "sort": "stars",
        "order": "desc",
        "per_page": min(100, max_repos),  # GitHub API max is 100 per page
    }

    url = "https://api.github.com/search/repositories"
    all_repos = []

    async with aiohttp.ClientSession() as session:
        # Get first page
        await rate_limiter.async_wait("search")

        for retry in range(max_retries):
            try:
                async with session.get(
                    url, headers=headers, params=params, timeout=30
                ) as response:
                    # Update rate limiter with response headers
                    rate_limiter.update_from_aiohttp_headers(response.headers)

                    if response.status == 403:
                        # Check if this is a rate limit issue
                        remaining = response.headers.get("X-RateLimit-Remaining")
                        if remaining and int(remaining) == 0:
                            reset_time = int(
                                response.headers.get(
                                    "X-RateLimit-Reset", time.time() + 3600
                                )
                            )
                            wait_time = max(0, reset_time - time.time())
                            print(
                                f"Rate limit exceeded when searching repos. Reset in {wait_time:.1f} seconds."
                            )
                            if (
                                wait_time > 0 and wait_time < 3600
                            ):  # Only wait if reasonable (<1 hour)
                                print(f"Waiting for rate limit reset...")
                                await asyncio.sleep(wait_time + 2)  # Add buffer
                                continue  # Try again after waiting
                            else:
                                print(
                                    "Rate limit wait time too long. Please use a GitHub token."
                                )
                                return []
                        else:
                            # This might be a secondary rate limit or abuse detection
                            retry_after = int(response.headers.get("Retry-After", 60))
                            print(
                                f"GitHub API temporary restriction. Waiting {retry_after} seconds before retry."
                            )
                            await asyncio.sleep(retry_after)
                            continue
                    # Handle other error responses
                    elif response.status == 401:
                        print("Authentication failed. Check your GitHub token.")
                        return []
                    elif response.status == 404:
                        print(f"Resource not found: {url}")
                        return []
                    elif response.status == 422:
                        text = await response.text()
                        print(f"Query error: {text}")
                        # Try with a simpler query
                        if "stars:>" in search_query:
                            simplified_query = search_query.split("stars:>")[0].strip()
                            print(f"Trying with simplified query: {simplified_query}")
                            return await get_repositories_async(
                                simplified_query,
                                max_repos,
                                headers,
                                rate_limiter,
                                max_retries - 1,
                            )
                        return []
                    elif response.status != 200:
                        text = await response.text()
                        print(f"Error searching repositories: Status {response.status}")
                        print(f"Response: {text}")
                        # Wait before retry
                        await asyncio.sleep(10 * (retry + 1))
                        continue

                    # Success!
                    result = await response.json()
                    repos = result.get("items", [])
                    all_repos.extend(repos)

                    # If we need more repos and there are multiple pages, fetch them in parallel
                    if len(all_repos) < max_repos and len(repos) == 100:
                        # Check if there are more pages
                        link_header = response.headers.get("Link", "")
                        next_urls = []

                        # Parse Link header to find next URLs
                        if 'rel="next"' in link_header:
                            parts = link_header.split(",")
                            for part in parts:
                                if 'rel="next"' in part:
                                    url_part = part.split(";")[0].strip()
                                    next_url = url_part[1:-1]  # Remove < and >
                                    next_urls.append(next_url)
                                    break

                        if next_urls:
                            print(
                                f"Fetched {len(all_repos)} repos, getting more to reach {max_repos}..."
                            )

                            # Fetch next pages in parallel
                            tasks = []
                            for next_url in next_urls[
                                : min(5, max_repos // 100)
                            ]:  # Limit to 5 pages max
                                tasks.append(
                                    fetch_next_page(
                                        session, next_url, headers, rate_limiter
                                    )
                                )

                            if tasks:
                                next_pages_results = await asyncio.gather(*tasks)
                                for page_repos in next_pages_results:
                                    all_repos.extend(page_repos)
                                    if len(all_repos) >= max_repos:
                                        break

                    return all_repos[:max_repos]  # Ensure we don't exceed max_repos

            except asyncio.TimeoutError:
                print(
                    f"Request timed out (attempt {retry+1}/{max_retries}). Retrying..."
                )
                await asyncio.sleep(5)
            except Exception as e:
                print(
                    f"Exception when searching repositories (attempt {retry+1}/{max_retries}): {e}"
                )
                await asyncio.sleep(5)

    print("Failed to fetch repositories after multiple attempts")
    return all_repos[:max_repos] if all_repos else []


async def fetch_next_page(session, url, headers, rate_limiter):
    """Fetch a single page of repository results."""
    await rate_limiter.async_wait("search")

    try:
        async with session.get(url, headers=headers, timeout=30) as response:
            # Update rate limiter
            rate_limiter.update_from_aiohttp_headers(response.headers)

            if response.status == 200:
                result = await response.json()
                return result.get("items", [])
            else:
                print(f"Error fetching next page: Status {response.status}")
                return []
    except Exception as e:
        print(f"Error fetching next page: {e}")
        return []


def get_repositories(search_query, max_repos, headers, rate_limiter, max_retries=3):
    """Synchronous wrapper for get_repositories_async."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(
            get_repositories_async(
                search_query, max_repos, headers, rate_limiter, max_retries
            )
        )
    finally:
        loop.close()


async def get_issues_page_async(
    repo_info, page, per_page, headers, rate_limiter, max_retries=3
):
    """Get a single page of issues for a repository using async."""
    owner, repo_name = repo_info

    # Only get closed issues as they're more likely to have solutions
    params = {
        "state": "closed",
        "per_page": per_page,
        "page": page,
        "sort": "comments",
        "direction": "desc",
    }

    url = f"https://api.github.com/repos/{owner}/{repo_name}/issues"

    async with aiohttp.ClientSession() as session:
        for retry in range(max_retries):
            try:
                # Wait for rate limiter
                await rate_limiter.async_wait("core")

                async with session.get(
                    url, headers=headers, params=params, timeout=30
                ) as response:
                    # Update rate limiter with response headers
                    rate_limiter.update_from_aiohttp_headers(response.headers)

                    # Handle rate limiting
                    if response.status == 403:
                        # Check if this is a rate limit issue
                        remaining = response.headers.get("X-RateLimit-Remaining")
                        if remaining and int(remaining) == 0:
                            reset_time = int(
                                response.headers.get(
                                    "X-RateLimit-Reset", time.time() + 3600
                                )
                            )
                            wait_time = max(0, reset_time - time.time())
                            print(
                                f"Rate limit exceeded for {owner}/{repo_name}. Reset in {wait_time:.1f} seconds."
                            )
                            if (
                                wait_time > 0 and wait_time < 3600
                            ):  # Only wait if reasonable (<1 hour)
                                print(f"Waiting for rate limit reset...")
                                await asyncio.sleep(wait_time + 2)  # Add buffer
                                continue  # Try again after waiting
                            else:
                                print(
                                    "Rate limit wait time too long. Please use a GitHub token."
                                )
                                return owner, repo_name, []
                        else:
                            # This might be a secondary rate limit or abuse detection
                            retry_after = int(response.headers.get("Retry-After", 60))
                            print(
                                f"GitHub API temporary restriction for {owner}/{repo_name}. Waiting {retry_after} seconds."
                            )
                            await asyncio.sleep(retry_after)
                            continue
                    # Handle other error responses
                    elif response.status == 401:
                        print(
                            f"Authentication failed for {owner}/{repo_name}. Check your GitHub token."
                        )
                        return owner, repo_name, []
                    elif response.status == 404:
                        print(f"Repository not found or no access: {owner}/{repo_name}")
                        return owner, repo_name, []
                    elif response.status != 200:
                        text = await response.text()
                        print(
                            f"Error fetching issues for {owner}/{repo_name} (page {page}): Status {response.status}"
                        )
                        print(f"Response: {text}")
                        # Wait before retry
                        await asyncio.sleep(5 * (retry + 1))
                        continue

                    # Success!
                    return owner, repo_name, await response.json()

            except asyncio.TimeoutError:
                print(
                    f"Request timed out for {owner}/{repo_name} (attempt {retry+1}/{max_retries}). Retrying..."
                )
                await asyncio.sleep(5)
            except Exception as e:
                print(
                    f"Exception when fetching issues for {owner}/{repo_name} (page {page}) (attempt {retry+1}/{max_retries}): {e}"
                )
                await asyncio.sleep(5)

        print(f"Failed to fetch issues for {owner}/{repo_name} after multiple attempts")
        return owner, repo_name, []


def get_issues_page(repo_info, page, per_page, headers, rate_limiter, max_retries=3):
    """Synchronous wrapper for get_issues_page_async."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(
            get_issues_page_async(
                repo_info, page, per_page, headers, rate_limiter, max_retries
            )
        )
    finally:
        loop.close()


async def get_issue_comments_async(
    repo_info, issue_number, headers, rate_limiter, max_retries=3
):
    """Get comments for a specific issue using async."""
    owner, repo_name = repo_info
    url = f"https://api.github.com/repos/{owner}/{repo_name}/issues/{issue_number}/comments"

    print(f"Fetching comments for issue #{issue_number} from {owner}/{repo_name}")

    async with aiohttp.ClientSession() as session:
        # Wait for rate limiter
        await rate_limiter.async_wait("core")

        async with session.get(url, headers=headers, timeout=30) as response:
            # Update rate limiter with response headers
            rate_limiter.update_from_aiohttp_headers(response.headers)

            print(f"Comments response status: {response.status}")

            if response.status == 403:
                # Handle rate limiting
                remaining = response.headers.get("X-RateLimit-Remaining")
                if remaining and int(remaining) == 0:
                    reset_time = int(
                        response.headers.get("X-RateLimit-Reset", time.time() + 3600)
                    )
                    wait_time = max(0, reset_time - time.time())
                    print(
                        f"Rate limit exceeded when fetching comments. Reset in {wait_time:.1f} seconds."
                    )
                    return []
            elif response.status != 200:
                print(
                    f"Error fetching comments for issue #{issue_number}: Status {response.status}"
                )
                response_text = await response.text()
                print(f"Response: {response_text}")
                return []

            # Get the comments
            comments = await response.json()
            print(f"Fetched {len(comments)} comments for issue #{issue_number}")

            if len(comments) > 0:
                print(f"First comment keys: {list(comments[0].keys())}")

            # For each comment, get its reactions
            for i, comment in enumerate(comments):
                if "reactions" not in comment:
                    print(f"Warning: Comment {i} has no reactions field")
                    continue

                comment["reactions_data"] = comment["reactions"]
            return comments


def find_best_solution(comments):
    """Find the most liked comment as the solution."""
    print(f"Finding best solution from {len(comments) if comments else 0} comments")

    if not comments:
        print("No comments available")
        return None

    # Debug: print the first comment to see its structure
    if comments and len(comments) > 0:
        print(f"First comment: {comments[0]}")
        print(f"Reactions in first comment: {comments[0].get('reactions', {})}")

    # Sort comments by reactions (+1 reactions) in descending order

    try:
        sorted_comments = sorted(
            comments,
            key=lambda c: (
                (
                    c.get("reactions_data", {}).get("heart", 0)
                    if "reactions_data" in c
                    else 0
                ),
            ),
            reverse=True,
        )

        print(f"Sorted {len(sorted_comments)} comments by reactions")

        # Return the most liked comment
        if sorted_comments:
            # Check if this comment has actual content
            best_comment = sorted_comments[0]
            body = best_comment.get("body", "").strip()
            print(f"Best comment has body length: {len(body)}")

            # If the best comment is empty or too short, try the next one
            if not body or len(body) < 10:
                print("Best comment is too short, looking for alternatives")
                for comment in sorted_comments[1:]:
                    body = comment.get("body", "").strip()
                    if body and len(body) >= 10:
                        print(
                            f"Found alternative comment with body length: {len(body)}"
                        )
                        return comment

            return best_comment
        return None
    except Exception as e:
        print(f"Error sorting comments: {e}")
        print(f"Comment types: {[type(c) for c in comments[:5]]}")
        # Return the first comment if sorting fails
        if comments and len(comments) > 0:
            return comments[0]
        return None


async def process_issues_with_solutions(
    owner, repo_name, issues, topic_name, headers, rate_limiter, max_issues=None
):
    """Process issue data into the desired format, including solutions."""
    processed_issues = []
    full_repo_name = f"{owner}/{repo_name}"

    for issue in issues:
        # Skip pull requests
        if "pull_request" in issue:
            continue

        # Skip open issues - we only want closed issues with solutions
