from crewai.tools import BaseTool
from typing import Literal
import requests
from bs4 import BeautifulSoup
import fitz  # PyMuPDF
from urllib.parse import urlparse
import os


class WebScrapingToolBS4(BaseTool):
    name: Literal["web_scraping_tool"] 
    description: str = (
        "Scrapes text and media from a webpage or PDF. "
        "Returns a dictionary containing: page_url, title, content, img_url, "
        "video_url, audio_url, pdf_url."
    )

    def extract_pdf_text(self, pdf_url):
        try:
            response = requests.get(pdf_url, timeout=15)
            temp_path = "temp_scraped.pdf"

            with open(temp_path, "wb") as f:
                f.write(response.content)

            text = ""
            with fitz.open(temp_path) as doc:
                for page in doc:
                    text += page.get_text()

            os.remove(temp_path)
            return text.strip()

        except Exception as e:
            return f"Error extracting PDF text: {str(e)}"

    def _run(self, url: str) -> dict:
        """Synchronous execution of the scraping tool."""
        try:
            parsed = urlparse(url)

            # ========== PDF Mode ==========
            if parsed.path.lower().endswith(".pdf"):
                text = self.extract_pdf_text(url)
                return {
                    "page_url": url,
                    "title": os.path.basename(parsed.path),
                    "content": text,
                    "img_url": [],
                    "video_url": [],
                    "audio_url": [],
                    "pdf_url": [url],
                }

            # ========== HTML Mode ==========
            response = requests.get(url, timeout=15)
            soup = BeautifulSoup(response.text, "html.parser")

            title = soup.title.string if soup.title else "Untitled"

            paragraphs = [p.get_text(" ", strip=True) for p in soup.find_all("p")]
            content = "\n".join(paragraphs)

            images = [img["src"] for img in soup.find_all("img", src=True)]
            videos = [v["src"] for v in soup.find_all("video", src=True)]
            audios = [a["src"] for a in soup.find_all("audio", src=True)]
            pdfs = [
                a["href"]
                for a in soup.find_all("a", href=True)
                if a["href"].lower().endswith(".pdf")
            ]

            return {
                "page_url": url,
                "title": title,
                "content": content,
                "img_url": images,
                "video_url": videos,
                "audio_url": audios,
                "pdf_url": pdfs,
            }

        except Exception as e:
            return {"error": str(e), "page_url": url}

    async def _arun(self, url: str) -> dict:
        """Async version (calls sync internally)."""
        return self._run(url)