from __future__ import annotations

import csv
import json
import re
import time
from dataclasses import asdict, dataclass
from typing import Any, Dict, List, Literal, Optional, Tuple
from urllib.parse import urljoin, urlparse
from urllib import robotparser

from pydantic import BaseModel

from a2a_pack import (
    A2AAgent,
    Pricing,
    Resources,
    RunContext,
    WorkspaceAccess,
    WorkspaceMode,
    skill,
)


# -----------------------------
# Public config/auth models
# -----------------------------


class SeleniumbaseWebsiteScraperConfig(BaseModel):
    # Reserved for future agent-level defaults
    default_headless: bool = True
    default_wait_seconds: float = 2.0


# -----------------------------
# Result models (for clarity)
# -----------------------------


@dataclass
class ScrapeFiles:
    json_path: Optional[str] = None
    csv_path: Optional[str] = None
    screenshots: Tuple[str, ...] = ()


@dataclass
class ScrapeResult:
    ok: bool
    items: List[Dict[str, Any]]
    pages_visited: int
    robots_respected: bool
    blocked: bool = False
    block_reason: str = ""
    warnings: Tuple[str, ...] = ()
    files: ScrapeFiles = ScrapeFiles()


# -----------------------------
# Agent implementation
# -----------------------------


class SeleniumbaseWebsiteScraper(A2AAgent[SeleniumbaseWebsiteScraperConfig, BaseModel]):
    name = "seleniumbase-website-scraper"
    description = (
        "Compliant website scraping via SeleniumBase/browser automation. "
        "Respects robots.txt, same-origin pagination, and safety rails."
    )
    version = "0.1.0"

    config_model = SeleniumbaseWebsiteScraperConfig

    # Deterministic/no-LLM agent: do not consume caller LLM credentials.
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes=(
            "Deterministic SeleniumBase scraper. No LLM usage; caller never billed for LLM."
        ),
    )

    # Browser automation benefits from more resources than the tiny default.
    resources = Resources(cpu="2", memory="2Gi", max_runtime_seconds=600)

    # Allow writing outputs (JSON/CSV/screenshots) into the caller workspace.
    workspace_access = WorkspaceAccess.dynamic(
        max_files=200,
        allowed_modes=(WorkspaceMode.READ_ONLY, WorkspaceMode.READ_WRITE_OVERLAY),
        require_reason=False,
    )

    tools_used = ("seleniumbase", "beautifulsoup4")

    # ---------------
    # Helper methods
    # ---------------

    @staticmethod
    def _robots_allows(url: str, user_agent: Optional[str]) -> Tuple[bool, List[str]]:
        ua = user_agent or "*"
        parsed = urlparse(url)
        if not parsed.scheme or not parsed.netloc:
            return False, ["invalid_url"]
        robots_url = f"{parsed.scheme}://{parsed.netloc}/robots.txt"
        warnings: List[str] = []
        rp = robotparser.RobotFileParser()
        try:
            rp.set_url(robots_url)
            rp.read()
        except Exception as exc:  # noqa: BLE001
            warnings.append(f"robots_fetch_error:{type(exc).__name__}")
            # If robots cannot be fetched, be conservative: allow but record warning
            return True, warnings
        try:
            allowed = bool(rp.can_fetch(ua, url))
        except Exception as exc:  # noqa: BLE001
            warnings.append(f"robots_parse_error:{type(exc).__name__}")
            allowed = True
        crawl_delay = getattr(rp, "crawl_delay", None)
        if callable(crawl_delay):
            try:
                delay = crawl_delay(ua)  # type: ignore[call-arg]
                if delay:
                    warnings.append(f"robots_crawl_delay:{delay}")
            except Exception:
                pass
        return allowed, warnings

    @staticmethod
    def _blocked_by_content(page_html: str) -> Tuple[bool, str]:
        hay = page_html.lower()
        patterns = {
            "captcha": r"captcha|are you human|i am not a robot|recaptcha",
            "access_denied": r"access denied|error 403|forbidden|blocked",
            "login_required": r"login required|please log in|sign in to continue",
            "paywall": r"subscribe now|membership required|paywall|premium only",
            "bot_detected": r"unusual traffic|automated queries|bot detected",
        }
        for reason, rx in patterns.items():
            if re.search(rx, hay, re.IGNORECASE):
                return True, reason
        return False, ""

    @staticmethod
    def _same_origin(u1: str, u2: str) -> bool:
        p1, p2 = urlparse(u1), urlparse(u2)
        return (p1.scheme, p1.netloc) == (p2.scheme, p2.netloc)

    # ----------------------
    # Public skills
    # ----------------------

    @skill(
        description=(
            "Scrape a website with SeleniumBase given a URL and CSS selectors. "
            "Honors robots.txt, supports same-origin pagination, optional screenshots, and JSON/CSV outputs."
        )
    )
    async def scrape_website(
        self,
        ctx: RunContext[BaseModel],
        url: str,
        selectors: Dict[str, str],
        pagination_selector: Optional[str] = None,
        item_selector: Optional[str] = None,
        max_pages: int = 1,
        headless: bool = True,
        wait_seconds: float = 2.0,
        user_agent: Optional[str] = None,
        rate_limit_seconds: float = 1.0,
        same_origin_only: bool = True,
        output_format: Literal["json", "csv"] = "json",
        screenshot: bool = False,
        out_basename: str = "scrape",
    ) -> Dict[str, Any]:
        # Robots.txt check up front
        allowed, robots_warnings = self._robots_allows(url, user_agent)
        if not allowed:
            result = ScrapeResult(
                ok=False,
                blocked=True,
                block_reason="robots_disallow",
                items=[],
                pages_visited=0,
                robots_respected=True,
                warnings=tuple(robots_warnings),
                files=ScrapeFiles(),
            )
            return asdict(result)

        await ctx.emit_progress("starting SeleniumBase scrape")

        # Prepare and run scraper inside the workspace sandbox so outputs persist
        script = self._render_seleniumbase_script(
            start_url=url,
            selectors=selectors,
            pagination_selector=pagination_selector,
            item_selector=item_selector,
            max_pages=max_pages,
            headless=headless,
            wait_seconds=wait_seconds,
            user_agent=user_agent,
            rate_limit_seconds=rate_limit_seconds,
            same_origin_only=same_origin_only,
            output_format=output_format,
            screenshot=screenshot,
            out_basename=out_basename,
        )
        exec_result = await ctx.workspace_python(
            script,
            image="python:3.11-slim",
            timeout_seconds=min(max(60.0, 10.0 * max_pages), 540.0),
            memory_mib=1536,
            cpus=2,
        )

        warnings: List[str] = []
        if exec_result.exit_code != 0:
            warnings.append(
                f"scraper_exit_nonzero:{exec_result.exit_code}: {exec_result.output[:300]}"
            )
        try:
            payload = json.loads(exec_result.output.strip().splitlines()[-1])
        except Exception:
            payload = {
                "ok": False,
                "blocked": False,
                "block_reason": "",
                "items": [],
                "pages_visited": 0,
                "robots_respected": True,
                "warnings": warnings + ["no_json_payload"],
                "files": asdict(ScrapeFiles()),
            }
        else:
            # merge sandbox warnings with ours
            if isinstance(payload.get("warnings"), list):
                payload["warnings"] = list(set(payload["warnings"] + warnings))
            else:
                payload["warnings"] = warnings

        return payload

    @skill(
        description=(
            "Iteratively attempt to scrape toward a goal without LLMs. "
            "Heuristically refines selectors and pagination (max 2 iterations) within same-origin and robots rules."
        )
    )
    async def iterate_scrape_goal(
        self,
        ctx: RunContext[BaseModel],
        goal: str,
        start_url: str,
        min_items: int = 25,
        max_steps: int = 2,
        headless: bool = True,
        wait_seconds: float = 2.0,
        user_agent: Optional[str] = None,
        rate_limit_seconds: float = 1.0,
        output_format: Literal["json", "csv"] = "json",
        screenshot: bool = False,
        out_basename: str = "iter",
    ) -> Dict[str, Any]:
        allowed, robots_warnings = self._robots_allows(start_url, user_agent)
        if not allowed:
            return asdict(
                ScrapeResult(
                    ok=False,
                    blocked=True,
                    block_reason="robots_disallow",
                    items=[],
                    pages_visited=0,
                    robots_respected=True,
                    warnings=tuple(robots_warnings),
                )
            )

        await ctx.emit_progress("heuristic discovery of selectors")

        # Minimal heuristic candidates for item container and fields
        container_candidates = [
            "article",
            "li",
            "div.card",
            "div.post",
            "div.result",
            "div.item",
            "div.entry",
            "div.product",
        ]
        field_candidates = {
            "title": ["h1", "h2", "h3", "a[title]", "a"],
            "link@href": ["a[href]"],
            "description": ["p", ".description", ".summary"],
        }

        steps = max(1, min(max_steps, 2))
        best_payload: Dict[str, Any] = {}
        best_count = 0
        for step in range(steps):
            item_selector = container_candidates[min(step, len(container_candidates) - 1)]
            selectors: Dict[str, str] = {}
            for key, opts in field_candidates.items():
                selectors[key] = opts[min(step, len(opts) - 1)]

            payload = await self.scrape_website(
                ctx,
                url=start_url,
                selectors=selectors,
                pagination_selector='a[rel="next"], a.next, a[aria-label="Next"], a:contains("Next"), a:contains("›")',
                item_selector=item_selector,
                max_pages=10,
                headless=headless,
                wait_seconds=wait_seconds,
                user_agent=user_agent,
                rate_limit_seconds=rate_limit_seconds,
                same_origin_only=True,
                output_format=output_format,
                screenshot=screenshot,
                out_basename=f"{out_basename}_s{step+1}",
            )
            items = payload.get("items") or []
            if len(items) >= min_items:
                payload.setdefault("warnings", []).append(
                    f"goal_reached_in_step:{step+1}"
                )
                return payload
            if len(items) > best_count:
                best_payload, best_count = payload, len(items)

        # return best effort
        if best_payload:
            best_payload.setdefault("warnings", []).append(
                "min_items_not_reached_best_effort"
            )
            return best_payload
        # fallback empty struct
        return asdict(
            ScrapeResult(
                ok=False,
                items=[],
                pages_visited=0,
                robots_respected=True,
                blocked=False,
                warnings=("no_items_found", *tuple(robots_warnings)),
            )
        )

    # ----------------------
    # Sandbox script factory
    # ----------------------

    def _render_seleniumbase_script(
        self,
        *,
        start_url: str,
        selectors: Dict[str, str],
        pagination_selector: Optional[str],
        item_selector: Optional[str],
        max_pages: int,
        headless: bool,
        wait_seconds: float,
        user_agent: Optional[str],
        rate_limit_seconds: float,
        same_origin_only: bool,
        output_format: str,
        screenshot: bool,
        out_basename: str,
    ) -> str:
        # Inline Python to run inside the sandbox with the workspace mounted
        # at /workspace. Keep imports local to avoid affecting agent card load.
        sel_json = json.dumps(selectors)
        pag_sel = json.dumps(pagination_selector) if pagination_selector else "None"
        item_sel = json.dumps(item_selector) if item_selector else "None"
        ua = json.dumps(user_agent) if user_agent else "None"
        script = f"""
import json, os, sys, time, csv, re
from urllib.parse import urlparse, urljoin
from urllib import robotparser

# Install runtime deps lightweightly if missing
try:
    import seleniumbase  # type: ignore
except Exception:
    import subprocess, sys
    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '--no-cache-dir', 'seleniumbase==4.28.4'])
    import seleniumbase  # type: ignore

items = []
warnings = []
pages_visited = 0
blocked = False
block_reason = ''

start_url = {json.dumps(start_url)}
selectors = {sel_json}
pagination_selector = {pag_sel}
item_selector = {item_sel}
max_pages = int({max_pages})
headless = bool({str(headless)})
wait_seconds = float({wait_seconds})
user_agent = {ua}
rate_limit_seconds = float({rate_limit_seconds})
same_origin_only = bool({str(same_origin_only)})
output_format = {json.dumps(output_format)}
screenshot = bool({str(screenshot)})
out_basename = {json.dumps(out_basename)}

origin = '{{}}://{{}}'.format(urlparse(start_url).scheme, urlparse(start_url).netloc)

# Robots.txt re-check in sandbox (best effort)
rp = robotparser.RobotFileParser()
try:
    rp.set_url(origin + '/robots.txt')
    rp.read()
    if not rp.can_fetch(user_agent or '*', start_url):
        print(json.dumps({{'ok': False, 'blocked': True, 'block_reason': 'robots_disallow', 'items': [], 'pages_visited': 0, 'robots_respected': True, 'warnings': ['sandbox_robots_disallow'], 'files': {{'json_path': None, 'csv_path': None, 'screenshots': []}}}}))
        sys.exit(0)
except Exception as e:
    warnings.append('robots_sandbox_warning')

# Start SeleniumBase driver (undetected-chromedriver helps with bot walls)
from seleniumbase import Driver  # type: ignore

try:
    driver = Driver(uc=True, headless=headless, agent=user_agent)
except Exception as e:
    warnings.append('driver_init_error:' + e.__class__.__name__)
    print(json.dumps({{'ok': False, 'blocked': False, 'block_reason': '', 'items': [], 'pages_visited': 0, 'robots_respected': True, 'warnings': warnings, 'files': {{'json_path': None, 'csv_path': None, 'screenshots': []}}}}))
    sys.exit(0)

current_url = start_url

try:
    for page in range(max_pages):
        driver.get(current_url)
        time.sleep(max(0.0, wait_seconds))
        pages_visited += 1
        html = driver.page_source or ''
        low = html.lower()
        # Blocker heuristics
        for key, rx in {{
            'captcha': r'captcha|are you human|i am not a robot|recaptcha',
            'access_denied': r'access denied|error 403|forbidden|blocked',
            'login_required': r'login required|please log in|sign in to continue',
            'paywall': r'subscribe now|membership required|paywall|premium only',
            'bot_detected': r'unusual traffic|automated queries|bot detected',
        }}.items():
            if re.search(rx, low, re.IGNORECASE):
                blocked = True
                block_reason = key
                break
        if blocked:
            break

        def pick_attr(selector: str) -> tuple[str, str | None]:
            # Support 'selector@attr' to pick element attribute
            if '@' in selector:
                base, attr = selector.split('@', 1)
                return base.strip(), attr.strip() or None
            return selector, None

        # Extraction
        page_items = []
        try:
            from selenium.webdriver.common.by import By  # type: ignore
            if item_selector:
                containers = driver.find_elements(By.CSS_SELECTOR, item_selector)[:200]
                for c in containers:
                    row: dict[str, str] = {{}}
                    for key, sel in selectors.items():
                        base_sel, attr = pick_attr(sel)
                        try:
                            el = c.find_element(By.CSS_SELECTOR, base_sel)
                            val = (el.get_attribute(attr) if attr else el.text).strip()
                        except Exception:
                            val = ''
                        row[key] = val
                    # compact empty rows
                    if any(v for v in row.values()):
                        page_items.append(row)
            else:
                # No item container: single snapshot using first match of each selector
                row: dict[str, str] = {{}}
                for key, sel in selectors.items():
                    base_sel, attr = pick_attr(sel)
                    try:
                        el = driver.find_element(By.CSS_SELECTOR, base_sel)
                        val = (el.get_attribute(attr) if attr else el.text).strip()
                    except Exception:
                        val = ''
                    row[key] = val
                if any(v for v in row.values()):
                    page_items.append(row)
        except Exception as e:
            warnings.append('extract_error:' + e.__class__.__name__)
        items.extend(page_items)

        # Screenshot
        if screenshot:
            try:
                os.makedirs('outputs', exist_ok=True)
                shot_path = f"outputs/{{out_basename}}_p{{page+1}}.png"
                driver.save_screenshot(shot_path)
            except Exception as e:
                warnings.append('screenshot_error')

        # Pagination
        if pagination_selector:
            try:
                from selenium.webdriver.common.by import By  # type: ignore
                next_candidates = driver.find_elements(By.CSS_SELECTOR, pagination_selector)
                next_href = None
                for el in next_candidates:
                    href = el.get_attribute('href')
                    if href:
                        if same_origin_only:
                            if urlparse(href).netloc and urlparse(href).netloc != urlparse(start_url).netloc:
                                continue
                        next_href = href
                        break
                if not next_href and next_candidates:
                    try:
                        next_candidates[0].click()
                        time.sleep(max(0.0, wait_seconds))
                        current_url = driver.current_url
                    except Exception:
                        break
                else:
                    if not next_href:
                        break
                    if same_origin_only and urlparse(next_href).netloc and urlparse(next_href).netloc != urlparse(start_url).netloc:
                        break
                    current_url = urljoin(current_url, next_href)
                time.sleep(max(0.0, rate_limit_seconds))
                continue
            except Exception:
                break
        else:
            break
finally:
    try:
        driver.quit()
    except Exception:
        pass

# Persist outputs
os.makedirs('outputs', exist_ok=True)
json_path = f"outputs/{{out_basename}}.json"
csv_path = f"outputs/{{out_basename}}.csv"
try:
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(items, f, ensure_ascii=False, indent=2)
except Exception as e:
    warnings.append('json_write_error')
    json_path = None

try:
    if items:
        keys = sorted({{k for row in items for k in row.keys()}})
        with open(csv_path, 'w', encoding='utf-8', newline='') as f:
            w = csv.DictWriter(f, fieldnames=keys)
            w.writeheader()
            for row in items:
                w.writerow(row)
    else:
        csv_path = None
except Exception:
    warnings.append('csv_write_error')
    csv_path = None

# Collect screenshot paths
shots = []
try:
    for name in os.listdir('outputs'):
        if name.startswith(out_basename + '_p') and name.endswith('.png'):
            shots.append('outputs/' + name)
except Exception:
    pass

payload = {{
    'ok': (not blocked),
    'items': items,
    'pages_visited': pages_visited,
    'robots_respected': True,
    'blocked': blocked,
    'block_reason': block_reason,
    'warnings': warnings,
    'files': {{'json_path': json_path, 'csv_path': csv_path, 'screenshots': shots}},
}}
print(json.dumps(payload))
"""
        return script