"""
Grok Camoufox Provider - Image/Video generation via UI automation with anti-detect browser.

Uses Camoufox (Firefox fork with C++ level fingerprint spoofing) for maximum stealth.
This is significantly more robust against bot detection than standard Playwright.

Key advantages over grok_playwright_provider.py:
- 0% headless detection rate (vs 67%+ with Playwright)
- Fingerprint spoofing at C++ level (not JavaScript injection)
- Built-in humanize behavior
- GeoIP-based timezone/locale matching

Auth: Cookie-based (sso, sso-rw, x-userid JWTs from grok.com).
Flow: Navigate to Imagine → Type prompt → Wait for images → Click "Make video" → Download.

Env vars:
  GROK_SSO_TOKEN       - 'sso' cookie value (JWT)
  GROK_SSO_RW_TOKEN    - 'sso-rw' cookie value (JWT)
  GROK_USER_ID         - 'x-userid' cookie value (UUID)
"""

from __future__ import annotations

import asyncio
import logging
import os
import re
import time
from pathlib import Path
from typing import Callable, Optional, List

import httpx

logger = logging.getLogger(__name__)

GROK_BASE = "https://grok.com"
VIDEO_GEN_TIMEOUT = 300
IMAGE_GEN_TIMEOUT = 120


class GrokCamoufoxProvider:
    """Generates images and videos via Grok's Imagine UI using Camoufox anti-detect browser."""

    def __init__(
        self,
        sso_token: str = "",
        sso_rw_token: str = "",
        user_id: str = "",
        headless: bool = True,
        use_virtual_display: bool = True,
    ):
        """
        Initialize Grok Camoufox Provider.

        Args:
            sso_token: Grok SSO token (or from GROK_SSO_TOKEN env var)
            sso_rw_token: Grok SSO-RW token (or from GROK_SSO_RW_TOKEN env var)
            user_id: Grok user ID (or from GROK_USER_ID env var)
            headless: Run in headless mode (default True)
            use_virtual_display: Use virtual display for headless (more stealth, default True)
        """
        self.sso_token = sso_token or os.getenv("GROK_SSO_TOKEN", "")
        self.sso_rw_token = sso_rw_token or os.getenv("GROK_SSO_RW_TOKEN", "")
        self.user_id = user_id or os.getenv("GROK_USER_ID", "")
        self.headless = headless
        self.use_virtual_display = use_virtual_display

        if not all([self.sso_token, self.sso_rw_token, self.user_id]):
            logger.warning(
                "Grok credentials not fully set (sso_token/sso_rw_token/user_id). "
                "Browser operations will fail until valid credentials are provided."
            )

        self._browser = None
        self._page = None
        self.client = httpx.Client(timeout=120, follow_redirects=True)

    async def _ensure_browser(self):
        """Launch Camoufox browser if not already running."""
        if self._page is not None:
            return

        try:
            from camoufox.async_api import AsyncCamoufox
        except ImportError:
            raise ImportError(
                "Camoufox not installed. Install with: pip install -U camoufox[geoip]"
            )

        # Configure Camoufox for maximum stealth
        camoufox_config = {
            "headless": "virtual" if self.headless and self.use_virtual_display else self.headless,
            "geoip": True,  # Auto-detect timezone/locale from IP
            "humanize": True,  # Enable human-like behavior
            "screen": {"max_width": 1920, "max_height": 1080},
            "locale": "en-US",
        }

        logger.info("Starting Camoufox browser (headless=%s, virtual=%s)", 
                    self.headless, self.use_virtual_display)

        self._browser = await AsyncCamoufox(**camoufox_config).__aenter__()
        self._page = await self._browser.new_page()

        # Add auth cookies
        await self._page.context.add_cookies([
            {
                "name": "sso",
                "value": self.sso_token,
                "domain": ".grok.com",
                "path": "/",
                "secure": True,
                "httpOnly": True,
            },
            {
                "name": "sso-rw",
                "value": self.sso_rw_token,
                "domain": ".grok.com",
                "path": "/",
                "secure": True,
                "httpOnly": True,
            },
            {
                "name": "x-userid",
                "value": self.user_id,
                "domain": ".grok.com",
                "path": "/",
                "secure": True,
                "httpOnly": True,
            },
        ])

        logger.info("Camoufox browser ready with Grok authentication")

    async def generate_image(
        self,
        prompt: str,
        output_dir: str | Path,
        num_images: int = 4,
        on_progress: Optional[Callable[[int], None]] = None,
    ) -> List[Path]:
        """
        Generate images using Grok Imagine UI.

        Args:
            prompt: Text prompt for image generation
            output_dir: Directory to save generated images
            num_images: Number of images to download (max available)
            on_progress: Optional callback for progress updates

        Returns:
            List of paths to saved images
        """
        await self._ensure_browser()
        page = self._page
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        logger.info("Generating images: %s", prompt[:60])

        # Navigate to main chat (works better than /imagine for headless)
        await page.goto(f"{GROK_BASE}/", timeout=60000)
        await page.wait_for_load_state("domcontentloaded")
        await asyncio.sleep(3)

        if on_progress:
            on_progress(5)

        # Click on the input area
        input_area = page.locator('div[contenteditable="true"]').first
        await input_area.click(force=True, timeout=10000)
        await asyncio.sleep(0.3)

        # Type the prompt using keyboard (with "Generate an image of" prefix)
        full_prompt = f"Generate an image of {prompt}" if not prompt.lower().startswith("generate") else prompt
        await page.keyboard.type(full_prompt, delay=30)
        await asyncio.sleep(0.5)

        if on_progress:
            on_progress(10)

        # Press Enter to submit
        await page.keyboard.press('Enter')

        logger.info("Prompt submitted, waiting for images...")

        # Wait for images to appear
        start = time.time()
        image_urls = []

        while time.time() - start < IMAGE_GEN_TIMEOUT:
            await asyncio.sleep(3)
            elapsed = int(time.time() - start)

            # Look for generated images (check for /generated/ in src)
            found_images = await page.evaluate("""
                () => {
                    return Array.from(document.querySelectorAll('img'))
                        .map(img => img.src)
                        .filter(src => src.includes('/generated/') && src.includes('.jpg'));
                }
            """)

            if found_images and len(found_images) >= 2:
                # Deduplicate
                image_urls = list(dict.fromkeys(found_images))[:num_images]
                logger.info("Found %d images after %ds", len(image_urls), elapsed)
                break

            if on_progress and elapsed > 5:
                progress = min(90, 10 + int(elapsed / IMAGE_GEN_TIMEOUT * 80))
                on_progress(progress)

            if elapsed % 10 == 0:
                logger.info("Waiting for images... (%ds)", elapsed)

        if not image_urls:
            raise RuntimeError(f"Image generation timed out after {IMAGE_GEN_TIMEOUT}s")

        if on_progress:
            on_progress(95)

        # Download images
        saved_paths = []
        for i, url in enumerate(image_urls):
            output_path = output_dir / f"grok_image_{i+1}.png"
            await self._download_file(url, output_path)
            saved_paths.append(output_path)
            logger.info("Saved: %s", output_path)

        if on_progress:
            on_progress(100)

        return saved_paths

    async def generate_video(
        self,
        image_path: str | Path,
        video_prompt: str = "",
        output_path: str | Path = None,
        on_progress: Optional[Callable[[int], None]] = None,
    ) -> Path:
        """
        Generate video from image using Grok Imagine I2V.

        Grok Imagine UI (as of 2026-03):
        - Bottom bar has: Image | Video | AspectRatio | 480p/720p | 6s/10s
        - "Type to imagine" contenteditable prompt field
        - "Upload image" button (top-right) + hidden file input
        - Submit button (aria-label="Submit")

        Flow: Video mode → 9:16 → 720p → 6s → upload image → type prompt → Submit → wait .mp4

        Args:
            image_path: Path to source image
            video_prompt: Optional motion prompt
            output_path: Where to save the video
            on_progress: Optional callback for progress updates

        Returns:
            Path to saved video file
        """
        await self._ensure_browser()
        page = self._page
        image_path = Path(image_path)

        if output_path is None:
            output_path = image_path.with_suffix(".mp4")
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")

        logger.info("Generating video from: %s", image_path.name)

        # ── 1. Navigate to Imagine ────────────────────────────────────
        await page.goto(f"{GROK_BASE}/imagine", timeout=60000)
        await page.wait_for_load_state("domcontentloaded")
        await asyncio.sleep(4)

        if on_progress:
            on_progress(5)

        # ── 2. Switch to Video mode (role="radio" button) ────────────
        logger.info("Selecting Video mode...")
        try:
            video_btn = page.locator('button[role="radio"]:has-text("Video")').first
            await video_btn.click(timeout=8000)
            await asyncio.sleep(2)
            logger.info("Video mode selected")
        except Exception as e:
            logger.warning("Could not click Video button: %s — trying text fallback", e)
            await page.locator('button:has-text("Video")').first.click(timeout=5000)
            await asyncio.sleep(2)

        if on_progress:
            on_progress(8)

        # ── 3. Select 9:16 aspect ratio ──────────────────────────────
        # Dropdown items are DIV[role="menuitem"], NOT buttons.
        logger.info("Selecting 9:16 aspect ratio...")
        try:
            ar_btn = page.locator('button[aria-label="Aspect Ratio"]').first
            await ar_btn.click(timeout=5000)
            await asyncio.sleep(1.5)
            nine_sixteen = page.locator('[role="menuitem"]:has-text("9:16")').first
            await nine_sixteen.click(timeout=5000)
            await asyncio.sleep(1)
            logger.info("Aspect ratio 9:16 selected")
        except Exception as e:
            logger.warning("Could not select 9:16 aspect ratio: %s", e)

        if on_progress:
            on_progress(10)

        # ── 4. Select 720p resolution ────────────────────────────────
        logger.info("Selecting 720p resolution...")
        try:
            res_btn = page.locator('button[role="radio"]:has-text("720p")').first
            await res_btn.click(timeout=5000)
            await asyncio.sleep(0.5)
            logger.info("720p selected")
        except Exception as e:
            logger.warning("Could not select 720p: %s", e)

        # ── 5. Select 6s duration ────────────────────────────────────
        logger.info("Selecting 6s duration...")
        try:
            dur_btn = page.locator('button[role="radio"]:has-text("6s")').first
            await dur_btn.click(timeout=5000)
            await asyncio.sleep(0.5)
            logger.info("6s duration selected")
        except Exception as e:
            logger.warning("Could not select 6s: %s", e)

        if on_progress:
            on_progress(12)

        # ── 6. Upload source image via hidden file input ─────────────
        logger.info("Uploading image: %s", image_path.name)
        file_input = page.locator('input[type="file"][accept="image/*"]').first
        await file_input.set_input_files(str(image_path.resolve()))
        await asyncio.sleep(4)
        logger.info("Image uploaded")

        if on_progress:
            on_progress(18)

        # ── 7. Type video prompt ─────────────────────────────────────
        if video_prompt:
            logger.info("Typing video prompt: %s", video_prompt[:80])
            try:
                prompt_el = page.locator('div[contenteditable="true"]').first
                await prompt_el.click(timeout=8000)
                await asyncio.sleep(0.3)
                await page.keyboard.type(video_prompt, delay=20)
                await asyncio.sleep(0.5)
                logger.info("Video prompt entered")
            except Exception as e:
                logger.warning("Could not type video prompt: %s", e)

        if on_progress:
            on_progress(22)

        # ── 8. Click Submit ──────────────────────────────────────────
        logger.info("Clicking Submit...")
        try:
            submit_btn = page.locator('button[aria-label="Submit"]').first
            await submit_btn.click(timeout=10000)
            logger.info("Submit clicked")
        except Exception as e:
            logger.warning("Submit button click failed: %s — trying Enter key", e)
            await page.keyboard.press("Enter")

        await asyncio.sleep(3)

        if on_progress:
            on_progress(25)

        # ── 9. Wait for redirect to /imagine/post/{uuid} ────────────
        post_token = None
        logger.info("Waiting for redirect to /imagine/post/...")
        start_redirect = time.time()
        while time.time() - start_redirect < 30:
            await asyncio.sleep(1)
            current_url = page.url
            match = re.search(r'/imagine/post/([a-f0-9-]+)', current_url)
            if match:
                post_token = match.group(1)
                logger.info("Post page loaded: %s", post_token)
                break

        if not post_token:
            logger.warning("No redirect to post page (URL: %s). Will still poll for video.", page.url)

        if on_progress:
            on_progress(30)

        # ── 10. Wait for video generation on post page ───────────────
        logger.info("Waiting for video generation (timeout %ds)...", VIDEO_GEN_TIMEOUT)
        start = time.time()
        video_url = None

        while time.time() - start < VIDEO_GEN_TIMEOUT:
            await asyncio.sleep(5)
            elapsed = int(time.time() - start)

            video_srcs = await page.evaluate("""
                () => {
                    const srcs = [];
                    document.querySelectorAll('video').forEach(v => {
                        if (v.src && v.src.includes('.mp4')) srcs.push(v.src);
                        v.querySelectorAll('source').forEach(s => {
                            if (s.src && s.src.includes('.mp4')) srcs.push(s.src);
                        });
                    });
                    return srcs;
                }
            """)

            if video_srcs:
                if post_token:
                    for src in video_srcs:
                        if post_token in src:
                            video_url = src
                            break
                if not video_url:
                    video_url = video_srcs[0]

                logger.info("Video ready after %ds: %s", elapsed, video_url[:100])
                break

            if on_progress and elapsed > 10:
                progress = min(95, 30 + int(elapsed / VIDEO_GEN_TIMEOUT * 65))
                on_progress(progress)

            if elapsed % 15 == 0:
                logger.info("Still generating video... (%ds elapsed)", elapsed)

        if not video_url:
            try:
                await page.screenshot(path="/tmp/grok_video_timeout.png")
                logger.error("Timeout screenshot: /tmp/grok_video_timeout.png (URL: %s)", page.url)
            except Exception:
                pass
            raise RuntimeError(
                f"Video generation timed out after {VIDEO_GEN_TIMEOUT}s. "
                f"Post token: {post_token}, Final URL: {page.url}"
            )

        if on_progress:
            on_progress(98)

        video_url = re.sub(r"\?.*$", "", video_url)

        if video_url.startswith("blob:"):
            logger.info("Downloading blob video via browser...")
            video_bytes = await page.evaluate("""
                async (url) => {
                    const resp = await fetch(url);
                    const buf = await resp.arrayBuffer();
                    return Array.from(new Uint8Array(buf));
                }
            """, video_url)
            output_path.write_bytes(bytes(video_bytes))
        else:
            await self._download_file(video_url, output_path)

        logger.info("Video saved: %s", output_path)

        if on_progress:
            on_progress(100)

        return output_path

    async def generate(
        self,
        image_path: str | Path,
        video_prompt: str = "",
        output_path: str | Path = None,
        on_progress: Optional[Callable[[int], None]] = None
    ) -> Path:
        """Alias for generate_video — matches GrokVideoProvider interface."""
        return await self.generate_video(image_path, video_prompt, output_path, on_progress)

    async def _download_file(self, url: str, output_path: Path):
        """Download file from URL using authenticated session."""
        cookie_header = "; ".join([
            f"sso={self.sso_token}",
            f"sso-rw={self.sso_rw_token}",
            f"x-userid={self.user_id}",
        ])

        headers = {
            "Cookie": cookie_header,
            "Referer": f"{GROK_BASE}/",
            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36",
        }

        resp = self.client.get(url, headers=headers, timeout=120)
        resp.raise_for_status()
        output_path.write_bytes(resp.content)

    async def close(self):
        """Close browser and cleanup."""
        if self._browser:
            try:
                await self._browser.__aexit__(None, None, None)
            except Exception as e:
                logger.warning("Error closing Camoufox browser: %s", e)
            self._browser = None
            self._page = None
        self.client.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.close()


# Sync wrapper for non-async usage
class GrokCamoufoxProviderSync:
    """Synchronous wrapper for GrokCamoufoxProvider."""

    def __init__(
        self,
        sso_token: str = "",
        sso_rw_token: str = "",
        user_id: str = "",
        headless: bool = True,
        use_virtual_display: bool = True,
    ):
        self._provider = GrokCamoufoxProvider(
            sso_token=sso_token,
            sso_rw_token=sso_rw_token,
            user_id=user_id,
            headless=headless,
            use_virtual_display=use_virtual_display,
        )
        self._loop = None

    def _get_loop(self):
        if self._loop is None or self._loop.is_closed():
            self._loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._loop)
        return self._loop

    def generate_image(self, prompt: str, output_dir: str | Path, **kwargs) -> List[Path]:
        return self._get_loop().run_until_complete(
            self._provider.generate_image(prompt, output_dir, **kwargs)
        )

    def generate_video(
        self,
        image_path: str | Path,
        video_prompt: str = "",
        output_path: str | Path = None,
        **kwargs
    ) -> Path:
        return self._get_loop().run_until_complete(
            self._provider.generate_video(image_path, video_prompt, output_path, **kwargs)
        )

    def close(self):
        if self._loop and not self._loop.is_closed():
            self._loop.run_until_complete(self._provider.close())
            self._loop.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    async def test():
        print("=" * 60)
        print("Grok Camoufox Provider Test (Anti-Detect Browser)")
        print("=" * 60)

        async with GrokCamoufoxProvider(headless=True) as provider:
            # Test image generation
            print("\n1. Testing image generation...")
            try:
                images = await provider.generate_image(
                    prompt="a majestic dragon flying over a medieval castle at sunset",
                    output_dir="./test_output",
                    num_images=2,
                    on_progress=lambda p: print(f"   Progress: {p}%")
                )
                print(f"   ✓ Generated {len(images)} images:")
                for img in images:
                    print(f"     - {img}")
            except Exception as e:
                print(f"   ✗ Error: {e}")
                import traceback
                traceback.print_exc()

        print("\n" + "=" * 60)
        print("Test complete")
        print("=" * 60)

    asyncio.run(test())