import os
import asyncio
from playwright.async_api import async_playwright
import google.generativeai as genai
from groq import Groq
from bs4 import BeautifulSoup

class BugHunterEngine:
    def __init__(self, api_key, provider="Gemini", model_name="gemini-1.5-pro-preview-0514"):
        self.provider = provider
        self.api_key = api_key
        self.model_name = model_name
        
        if provider == "Gemini":
            genai.configure(api_key=api_key)
            self.gemini_model = genai.GenerativeModel(model_name)
        elif provider == "Groq":
            self.groq_client = Groq(api_key=api_key)

        # Load the skill file
        skill_path = os.path.join(os.path.dirname(__file__), "ADVANCED_QA_AGENT_SKILL.md")
        with open(skill_path, "r", encoding="utf-8") as f:
            self.skill_content = f.read()

    def call_ai(self, prompt):
        """Unified method to call Gemini or Groq."""
        if self.provider == "Gemini":
            response = self.gemini_model.generate_content(prompt)
            return response.text
        else:
            chat_completion = self.groq_client.chat.completions.create(
                messages=[{"role": "system", "content": self.skill_content}, {"role": "user", "content": prompt}],
                model="llama3-70b-8192",
            )
            return chat_completion.choices[0].message.content

    async def prune_dom(self, page_source):
        soup = BeautifulSoup(page_source, "lxml")
        for tag in soup(["script", "style", "svg", "path", "iframe"]):
            tag.decompose()
        interactive_tags = ["a", "button", "input", "select", "textarea", "h1", "h2", "h3", "h4", "h5", "h6", "p"]
        pruned_html = ""
        for tag in soup.find_all(interactive_tags):
            attrs = {k: v for k, v in tag.attrs.items() if k in ["id", "class", "href", "name", "type", "role"]}
            tag.attrs = attrs
            pruned_html += str(tag) + "\n"
        return pruned_html[:30000]

    async def run_test(self, url, progress=None):
        logs = []
        screenshots = []

        async with async_playwright() as p:
            # Use --no-sandbox for Docker/HuggingFace stability
            browser = await p.chromium.launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])
            page = await browser.new_page(viewport={"width": 1280, "height": 800})
            
            if progress: progress(0.1, desc="🌍 Opening Website...")
            await page.goto(url, wait_until="networkidle")
            
            # Initial Screenshot
            shot_path = "evidence_initial.png"
            await page.screenshot(path=shot_path)
            screenshots.append(shot_path)
            
            dom = await self.prune_dom(await page.content())
            
            if progress: progress(0.5, desc="🧠 AI Agent Exploring...")
            
            # Simulated Agent Thought (Compiling for speed in this demo)
            final_prompt = f"""
{self.skill_content}

URL: {url}
HTML SNAPSHOT:
{dom}

ACT AS AN EXPERT QA AGENT.
Generate a SUPERIOR SBTM v2 Exploratory Test Report in Markdown.
You MUST follow the reporting protocol from the skill file.
IDENTIFY BUGS relevant to: SFDIPOT, HICCUPPS, A11y, and Security (IDOR/JWT).

Return ONLY the Markdown report.
"""
            report_md = self.call_ai(final_prompt)
            
            await browser.close()
            return report_md, screenshots