import requests
from bs4 import BeautifulSoup
from urllib.parse import urlparse, urljoin
from markdownify import markdownify as md
import tempfile
import zipfile
import re
from typing import Tuple
import os
import gradio as gr
from collections import deque
from collections import defaultdict


# ===========================================================
# FORMATING FOR FILENAMES TO PRESERVE URL PATH
# ===========================================================


def url_to_zip_path(url, extension=".md"):
    parsed = urlparse(url)

    # example.com
    domain = parsed.netloc.replace("www.", "")

    # about/team/
    path = parsed.path.strip("/")

    if path == "":
        # Homepage
        return os.path.join(domain, "index" + extension)

    # Remove extension if present
    path = re.sub(r"\.[^.]+$", "", path)

    return os.path.join(domain, path + extension)


# ===========================================================
# 🌐 WEBSITE CRAWLER
# ===========================================================

def crawl_site_for_links(start_url: str, max_pages: int = 50, max_depth: int = 2):
    """
    Recursively crawl a website and collect:
    • Internal HTML pages
    • PDF files
    
    We stay inside the same domain for safety.
    """

    visited = set()
    html_links = set()
    pdf_links = set()

    parsed_base = urlparse(start_url)
    domain = parsed_base.netloc

    queue = deque([(start_url, 0)])
    session = requests.Session()
    session.headers.update({
        "User-Agent": "Mozilla/5.0"
    })

    while queue and len(visited) < max_pages:
        current_url, depth = queue.popleft()

        if current_url in visited or depth > max_depth:
            continue

        visited.add(current_url)

        try:
            response = session.get(current_url, timeout=10)

            if "text/html" not in response.headers.get("Content-Type", ""):
                continue

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

            for a in soup.find_all("a", href=True):
                href = a["href"]
                full_url = urljoin(current_url, href)
                parsed = urlparse(full_url)

                if parsed.netloc != domain:
                    continue

                if full_url.lower().endswith(".pdf"):
                    pdf_links.add(full_url)
                elif not href.startswith(("#", "javascript:", "mailto:", "tel:")):
                    html_links.add(full_url)
                    if full_url not in visited:
                        queue.append((full_url, depth + 1))

        except Exception:
            continue

    return html_links, pdf_links


# ===========================================================
# 📦 EXTRACTION ENGINE
# ===========================================================

def extract_all_content_as_zip(url: str, max_links: int, max_depth: int) -> Tuple[str, str]:
    """
    Main function:
    • Crawls the site
    • Converts pages to Markdown
    • Downloads PDFs
    • Packs everything into a ZIP file
    """

    try:
        if not url.startswith(("http://", "https://")):
            url = "https://" + url

        html_links, pdf_links = crawl_site_for_links(url, max_links, max_depth)

        if not html_links and not pdf_links:
            return "❌ No internal pages or PDFs found.", None

        with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as temp_zip:
            zip_path = temp_zip.name

        session = requests.Session()
        session.headers.update({"User-Agent": "Mozilla/5.0"})

        html_ok = 0
        pdf_ok = 0

        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zip_file:

            # ---- HTML → Markdown ----
            for i, link_url in enumerate(html_links, 1):
                try:
                    resp = session.get(link_url, timeout=10)
                    soup = BeautifulSoup(resp.content, "html.parser")

                    for tag in soup(["script","style","nav","footer","header","aside"]):
                        tag.decompose()

                    main_content = (
                        soup.find("main")
                        or soup.find("article")
                        or soup.find("body")
                    )

                    markdown_text = md(str(main_content))
                    title = soup.find("title")
                    if title:
                        markdown_text = f"# {title.text.strip()}\n\n{markdown_text}"

                    filename = url_to_zip_path(link_url, ".md")
                    zip_file.writestr(filename, markdown_text)
                    html_ok += 1

                except Exception:
                    pass

            # ---- PDFs ----
            for j, pdf_url in enumerate(pdf_links, 1):
                try:
                    resp = session.get(pdf_url, timeout=20)
                    pdf_path = url_to_zip_path(pdf_url, ".pdf")
                    zip_file.writestr(pdf_path, resp.content)
                    pdf_ok += 1
                except Exception:
                    pass

        message = f"""
✅ Extraction completed!

• HTML pages saved as Markdown: {html_ok}
• PDFs downloaded: {pdf_ok}

You can now download the ZIP file below.
"""
        return message, zip_path

    except Exception as e:
        return f"❌ Error: {str(e)}", None



# ===========================================================
# 🗺️ SITEMAP DISCOVERY
# ===========================================================

def discover_sitemaps_from_robots(root_url):
    import requests

    robots_url = root_url.rstrip("/") + "/robots.txt"
    sitemaps = []

    try:
        r = requests.get(robots_url, timeout=10)
        if r.status_code != 200:
            return []

        for line in r.text.splitlines():
            if line.lower().startswith("sitemap:"):
                sitemaps.append(line.split(":", 1)[1].strip())

    except Exception:
        pass

    return sitemaps
    
def load_sitemap_sections(site_url):

    import requests
    import xml.etree.ElementTree as ET
    from collections import defaultdict
    from urllib.parse import urlparse

    if not site_url.startswith(("http://", "https://")):
        site_url = "https://" + site_url

    parsed = urlparse(site_url)
    root = f"{parsed.scheme}://{parsed.netloc}"

    session = requests.Session()
    session.headers.update({"User-Agent": "Mozilla/5.0"})

    # ----------------------------
    # STEP 1: discover sitemap URLs
    # ----------------------------
    sitemap_urls = []

    # from robots.txt
    sitemap_urls += discover_sitemaps_from_robots(root)

    # fallback guesses
    sitemap_urls += [
        root + "/sitemap.xml",
        root + "/sitemap_index.xml"
    ]

    all_urls = set()

    # ----------------------------
    # STEP 2: parse sitemaps
    # ----------------------------
    for sm_url in sitemap_urls:

        try:
            r = session.get(sm_url, timeout=10)
            if r.status_code != 200:
                continue

            content = r.content

            # handle gzip sitemap
            if sm_url.endswith(".gz"):
                import gzip
                content = gzip.decompress(content)

            xml = ET.fromstring(content)

            # ---- CASE A: sitemap index ----
            locs = xml.findall(".//{*}loc")

            sub_sitemaps = [l.text for l in locs]

            if len(sub_sitemaps) > 0 and "sitemap" in sm_url:
                # treat as index
                for sub in sub_sitemaps:
                    try:
                        r2 = session.get(sub, timeout=10)
                        xml2 = ET.fromstring(r2.content)
                        for u in xml2.findall(".//{*}loc"):
                            if u.text:
                                all_urls.add(u.text)
                    except:
                        pass
            else:
                # normal sitemap
                for u in locs:
                    if u.text:
                        all_urls.add(u.text)

        except Exception:
            continue

    # ----------------------------
    # STEP 3: build meaningful prefixes
    # ----------------------------
    prefixes = {}

    for u in all_urls:
        path = urlparse(u).path.strip("/")

        if not path:
            continue

        parts = path.split("/")

        # build hierarchical prefixes (IMPORTANT IMPROVEMENT)
        if len(parts) >= 2:
            prefix = "/" + "/".join(parts[:2]) + "/"
        else:
            prefix = "/" + parts[0] + "/"

        prefixes[prefix] = prefixes.get(prefix, 0) + 1

    ordered = sorted(prefixes.items(), key=lambda x: -x[1])

    return gr.update(
        choices=[p[0] for p in ordered],
        value=[]
    )

    

# ===========================================================
# 🎨 GRADIO WEB APP (GRADIO 6 SAFE)
# ===========================================================

def run_extraction(url, max_links, depth):
    return extract_all_content_as_zip(url, int(max_links), int(depth))


with gr.Blocks(title="Website Content Extractor") as app:

    gr.Markdown("""
# Website Content & PDF Extractor

Download the **text and PDFs from a website** and package everything into a ZIP file.
""")

    gr.Markdown("---")

    # HOW TO USE SECTION (replaces Box)
    with gr.Group():
        gr.Markdown("## How to use")

        gr.Markdown("""
1️⃣ Enter a website homepage  
2️⃣ Choose how deep to crawl  
3️⃣ Click **Start Extraction**  
4️⃣ Download your ZIP file  

""")

    gr.Markdown("---")

    url_input = gr.Textbox(
        label="Website URL",
        placeholder="https://example.com"
    )

    with gr.Row():
        max_links_input = gr.Slider(
            10, 1000, value=50, step=10,
            label="Maximum pages to scan",
            info="Higher = more content but slower"
        )

        depth_input = gr.Slider(
            1, 3, value=2, step=1,
            label="Crawl depth",
            info="How many clicks away from homepage"
        )
        

    # ----------------------------------------------------
    # SITEMAP MODULE
    # ----------------------------------------------------

    with gr.Group():
    
        gr.Markdown("## Additional pages from sitemap")
    
        load_sitemap_btn = gr.Button("Load sitemap")
    
        sitemap_dropdown = gr.Dropdown(
            choices=[],
            multiselect=True,
            label="Select sitemap sections",
            interactive=True
        )


    run_btn = gr.Button("🚀 Start Extraction", variant="primary")

    status_output = gr.Textbox(label="Status")
    file_output = gr.File(label="Download ZIP")

    # -------------------------------
    # Sitemap button callback
    # -------------------------------
    load_sitemap_btn.click(
        fn=load_sitemap_sections,
        inputs=url_input,
        outputs=sitemap_dropdown
    )

    run_btn.click(
        fn=run_extraction,
        inputs=[url_input, max_links_input, depth_input],
        outputs=[status_output, file_output]
    )


# ===========================================================
# 🚀 ENTRY POINT
# ===========================================================

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(),
        ssr_mode=False
    )