from collections import deque
import pandas as pd
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import time
from src.logger import get_logger

logger = get_logger(__name__)


API = "https://pl.wikipedia.org/w/api.php"


def build_session(user_agent: str) -> requests.Session:
    session = requests.Session()
    retries = Retry(
        total=5,
        backoff_factor=2.0,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET"],
        raise_on_status=False,
        respect_retry_after_header=True,
    )
    adapter = HTTPAdapter(max_retries=retries)
    session.mount("https://", adapter)
    session.mount("http://", adapter)
    session.headers.update({"User-Agent": user_agent})
    return session


def get_subcategories(cat_title, session):
    params = {
        "action": "query",
        "list": "categorymembers",
        "cmtitle": cat_title,
        "cmtype": "page|subcat",
        "cmlimit": "max",
        "format": "json",
    }

    subcategories = []
    subpages = []
    while True:
        r = session.get(API, params=params, timeout=30)
        r.raise_for_status()

        data = r.json()
        if "error" in data:
            logger.error(
                f"Error fetching subcategories of {cat_title}: {data['error']}"
            )
            return [], []

        batch_subcategories = [
            c["title"]
            for c in data.get("query", {}).get("categorymembers", [])
            if c["ns"] == 14
        ]
        subcategories.extend(batch_subcategories)

        batch_subpages = [
            c["title"]
            for c in data.get("query", {}).get("categorymembers", [])
            if c["ns"] == 0
        ]
        subpages.extend(batch_subpages)

        if "continue" in data:
            params["cmcontinue"] = data["continue"]["cmcontinue"]
            time.sleep(0.1)
        else:
            break

    return subcategories, subpages


def get_category_tree(root, user_agent, max_depth=5) -> list[tuple[str, str]]:
    logger.info(f"Starting from root category: {root}. Max depth: {max_depth}")

    edges = []

    seen = set([root])
    queue = deque([(root, 0)])
    last = time.time()

    session = build_session(user_agent)
    while queue:
        parent, depth = queue.popleft()
        if max_depth and depth >= max_depth:
            continue

        subcategories, subpages = get_subcategories(parent, session)
        edges.append((parent, subpages, subcategories))

        for child in subcategories:
            if child not in seen:
                seen.add(child)
                queue.append((child, depth + 1))

        if len(edges) % 10 == 0:
            elapsed, last = time.time() - last, time.time()
            logger.info(f"Collected {len(edges)} edges so far... Took {elapsed:.1f}s")

        time.sleep(0.1)

    return edges


if __name__ == "__main__":
    import argparse
    import os

    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=str, help="Output file path", required=True)
    parser.add_argument(
        "--root", type=str, default="Kategoria:Kategorie", help="Root category"
    )
    parser.add_argument(
        "--user_agent",
        type=str,
        default="TODO: replace with your user agent",
        help="User-Agent for Wikipedia API requests",
    )
    parser.add_argument(
        "--max_depth", type=int, default=None, help="Maximum depth to crawl"
    )
    args = parser.parse_args()

    edges = get_category_tree(args.root, args.user_agent, args.max_depth)

    df = pd.DataFrame(edges, columns=["title", "subpages", "subcategories"])
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    use_lines = args.output.endswith(".jsonl")
    df.to_json(
        args.output,
        lines=use_lines,
        indent=2 if not use_lines else None,
        orient="records",
    )
