"""
Grok Video Provider -- Image-to-Video via Grok's Imagine UI.

Auth: Cookie-based (sso, sso-rw, x-userid JWTs from grok.com).
Flow: Navigate to Imagine → Upload image via UI → Click "Make video" →
      Wait for video elements → Download video from assets CDN.

Uses undetected_chromedriver to bypass Cloudflare anti-bot detection.
The conversations/new endpoint only accepts requests triggered by genuine
UI interactions (button clicks), so we drive the actual Imagine UI.

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 collections.abc import Callable
from pathlib import Path
from typing import Any

import httpx

from providers.ports import VideoProvider

logger = logging.getLogger(__name__)

GROK_BASE = "https://grok.com"
GROK_ASSETS_BASE = "https://assets.grok.com"

USER_AGENT = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36"
)

# How long to wait for video generation (seconds)
VIDEO_GEN_TIMEOUT = 300
VIDEO_POLL_INTERVAL = 3


class GrokVideoProvider(VideoProvider):
    """Generates video clips from images via Grok's Imagine I2V UI.

    Uses undetected_chromedriver to drive the Grok Imagine UI, since the
    /rest/app-chat/conversations/new endpoint rejects programmatic fetch()
    calls with "anti-bot rules" — only UI-initiated requests are accepted.
    """

    name = "grok"

    def __init__(self, sso_token: str = "", sso_rw_token: str = "", user_id: str = ""):
        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", "")
        
        if not all([self.sso_token, self.sso_rw_token, self.user_id]):
            raise EnvironmentError(
                "Grok credentials required. Pass sso_token/sso_rw_token/user_id "
                "or set GROK_SSO_TOKEN/GROK_SSO_RW_TOKEN/GROK_USER_ID env vars."
            )
        
        self._driver = None
        self.client = httpx.Client(timeout=120, follow_redirects=True)

    # ------------------------------------------------------------------
    # Browser lifecycle
    # ------------------------------------------------------------------

    def _ensure_browser(self):
        """Launch undetected Chrome if not already running.

        Uses Xvfb (virtual display) when available so the browser runs
        invisibly without headless mode (which Cloudflare detects).
        Falls back to a visible window if Xvfb is not installed.
        """
        if self._driver is not None:
            return

        import undetected_chromedriver as uc

        # Try to start a virtual display (Xvfb) so Chrome runs invisibly.
        # Headless mode is NOT an option — Cloudflare blocks it.
        self._virtual_display = None
        try:
            from pyvirtualdisplay import Display

            self._virtual_display = Display(visible=False, size=(1400, 900))
            self._virtual_display.start()
            logger.info("Xvfb virtual display started")
        except Exception:
            logger.info("Xvfb not available — Chrome will open a visible window")

        options = uc.ChromeOptions()
        options.add_argument("--window-size=1400,900")
        options.add_argument("--disable-extensions")
        options.add_argument("--no-first-run")
        options.add_argument("--no-default-browser-check")

        logger.info("Launching undetected Chrome browser...")
        self._driver = uc.Chrome(options=options, version_main=144)

        # Navigate to grok.com and set auth cookies
        self._driver.get(f"{GROK_BASE}/")
        time.sleep(2)

        self._driver.add_cookie({
            "name": "sso", "value": self.sso_token,
            "domain": ".grok.com", "path": "/",
            "secure": True, "httpOnly": True,
        })
        self._driver.add_cookie({
            "name": "sso-rw", "value": self.sso_rw_token,
            "domain": ".grok.com", "path": "/",
            "secure": True, "httpOnly": True,
        })
        self._driver.add_cookie({
            "name": "x-userid", "value": self.user_id,
            "domain": ".grok.com", "path": "/",
            "secure": True, "httpOnly": True,
        })

        logger.info("Browser ready with auth cookies")

    # ------------------------------------------------------------------
    # Generate via UI
    # ------------------------------------------------------------------

    def _generate_via_ui(
        self,
        image_path: Path,
        video_prompt: str = "",
        on_progress: Callable[[int], None] | None = None,
    ) -> str:
        """Drive Grok Imagine UI to generate I2V. Returns video URL."""
        from selenium.webdriver.common.by import By
        from selenium.webdriver.common.action_chains import ActionChains
        from selenium.webdriver.support.ui import WebDriverWait
        from selenium.webdriver.support import expected_conditions as EC

        driver = self._driver

        # Navigate to Imagine page
        logger.info("Navigating to Grok Imagine...")
        driver.get(f"{GROK_BASE}/imagine")
        time.sleep(4)

        # Store original window handle and count
        original_window = driver.current_window_handle
        original_window_count = len(driver.window_handles)

        # Upload image via file input
        logger.info("Uploading %s via UI...", image_path.name)
        file_input = driver.find_element(
            By.CSS_SELECTOR, 'input[type="file"][accept="image/*"]'
        )
        file_input.send_keys(str(image_path.resolve()))

        if on_progress:
            on_progress(5)

        # Wait for redirect to /imagine/post/{token} or new tab
        logger.info("Waiting for redirect to post page...")
        post_token = None
        start_wait = time.time()
        
        while time.time() - start_wait < 30:
            time.sleep(1)
            
            # Check if new tab was opened
            if len(driver.window_handles) > original_window_count:
                # Switch to the new tab
                for handle in driver.window_handles:
                    if handle != original_window:
                        driver.switch_to.window(handle)
                        logger.info("Switched to new tab")
                        break
            
            # Check if URL contains /imagine/post/
            current_url = driver.current_url
            match = re.search(r'/imagine/post/([a-f0-9-]+)', current_url)
            if match:
                post_token = match.group(1)
                logger.info("Post page loaded with token: %s", post_token)
                break
        
        if not post_token:
            raise RuntimeError(
                f"Failed to get post token after upload. Current URL: {driver.current_url}"
            )

        # Wait for page to fully load
        time.sleep(3)

        if on_progress:
            on_progress(10)

        # Type the video prompt FIRST, then click Make video
        if video_prompt:
            logger.info("Typing video prompt: %s", video_prompt[:80])
            try:
                # Find the textbox with placeholder "Type to customize video..."
                prompt_input = WebDriverWait(driver, 15).until(
                    EC.presence_of_element_located((
                        By.CSS_SELECTOR,
                        'input[placeholder*="customize video"], '
                        'textarea[placeholder*="customize video"], '
                        '[role="textbox"][aria-label="Make a video"]'
                    ))
                )
                # Clear and type the prompt
                prompt_input.click()
                time.sleep(0.5)
                prompt_input.clear()
                prompt_input.send_keys(video_prompt)
                time.sleep(1)
                logger.info("Video prompt entered successfully")
            except Exception as exc:
                logger.warning("Could not type video prompt via input: %s", exc)
                # Try alternative: use ActionChains to type
                try:
                    body = driver.find_element(By.TAG_NAME, "body")
                    # Find any element that looks like a prompt input
                    inputs = driver.find_elements(
                        By.CSS_SELECTOR,
                        '[contenteditable="true"], input[type="text"], textarea'
                    )
                    for inp in inputs:
                        if "customize" in inp.get_attribute("placeholder") or "":
                            ActionChains(driver).click(inp).pause(0.3).send_keys(video_prompt).perform()
                            logger.info("Video prompt entered via ActionChains")
                            break
                except Exception as exc2:
                    logger.warning("Alternative prompt entry also failed: %s", exc2)

        # Click "Video Options" button to open settings menu and select 720p
        # NOTE: Button has pointer-events: none, so must use ActionChains for real mouse click
        logger.info("Opening Video Options menu to select 720p...")
        try:
            # Find Video Options button and click with ActionChains (real mouse click)
            video_options_btn = WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((
                    By.CSS_SELECTOR,
                    'button[aria-label="Video Options"]'
                ))
            )
            ActionChains(driver).move_to_element(video_options_btn).click().perform()
            time.sleep(2)  # Wait for menu animation
            logger.info("Video Options menu opened")
            
            # Select 720p resolution using ActionChains (real mouse click)
            logger.info("Selecting 720p resolution...")
            try:
                resolution_btn = WebDriverWait(driver, 5).until(
                    EC.presence_of_element_located((
                        By.CSS_SELECTOR,
                        'button[aria-label="720p"]'
                    ))
                )
                ActionChains(driver).move_to_element(resolution_btn).click().perform()
                time.sleep(0.5)
                logger.info("Selected 720p resolution")
            except Exception as e:
                logger.warning("720p button not found in menu: %s", e)
            
        except Exception as e:
            logger.warning("Could not open Video Options menu: %s", e)

        if on_progress:
            on_progress(12)

        # Click "Make video" button (arrow icon after image is uploaded)
        logger.info("Clicking 'Make video' button...")
        try:
            # Use JavaScript to click Make video button - more reliable
            driver.execute_script("""
                // Find button with aria-label="Make video"
                const btn = document.querySelector('button[aria-label="Make video"]');
                if (btn) {
                    btn.click();
                    return;
                }
                // Fallback: find button containing "Make video" text
                const btns = document.querySelectorAll('button');
                for (const b of btns) {
                    if (b.textContent.includes('Make video') || b.textContent.includes('Redo')) {
                        b.click();
                        break;
                    }
                }
            """)
            logger.info("Clicked 'Make video' button")
        except Exception as e:
            logger.warning("Could not click Make video button: %s", e)
        
        time.sleep(2)

        if on_progress:
            on_progress(15)

        # Wait for video generation — poll for <video> elements with src
        # IMPORTANT: Filter by post_token to ensure we get the right video
        logger.info("Waiting for video generation (timeout %ds, token: %s)...", VIDEO_GEN_TIMEOUT, post_token[:8])
        start = time.time()
        video_url = None

        while time.time() - start < VIDEO_GEN_TIMEOUT:
            time.sleep(VIDEO_POLL_INTERVAL)
            elapsed = int(time.time() - start)

            # Check for video elements with src URLs from Grok's CDN
            # Filter by post_token to ensure we get the correct video
            video_srcs = driver.execute_script(r"""
                const srcs = [];
                document.querySelectorAll('video').forEach(v => {
                    if (v.src && v.src.includes('/generated/') && v.src.includes('.mp4')) {
                        srcs.push(v.src);
                    }
                    // Also check source elements
                    v.querySelectorAll('source').forEach(s => {
                        if (s.src && s.src.includes('/generated/') && s.src.includes('.mp4')) {
                            srcs.push(s.src);
                        }
                    });
                });
                return srcs;
            """)

            if video_srcs:
                # Prefer video URL that contains the post token
                for src in video_srcs:
                    if post_token in src:
                        video_url = src
                        logger.info("Found video matching post token!")
                        break
                
                # If no match by token, use the first one (fallback)
                if not video_url:
                    video_url = video_srcs[0]
                    logger.warning("No video matched token, using first available")
                
                if on_progress:
                    on_progress(100)
                logger.info(
                    "Video ready after %ds: %s",
                    elapsed,
                    video_url[:80],
                )
                break

            # Report progress estimate based on elapsed time (~60s typical)
            if on_progress and elapsed > 5:
                estimated = min(95, int(elapsed / 60 * 90))
                on_progress(estimated)

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

        if not video_url:
            raise RuntimeError(
                f"Video generation timed out after {VIDEO_GEN_TIMEOUT}s. "
                f"Post token: {post_token}, Final URL: {driver.current_url}"
            )

        # Clean up cache parameter
        video_url = re.sub(r"\?cache=\d+$", "", video_url)

        # Close extra tabs and return to original window
        try:
            current_handles = driver.window_handles
            if len(current_handles) > 1 and original_window in current_handles:
                # Close all tabs except original
                for handle in current_handles:
                    if handle != original_window:
                        driver.switch_to.window(handle)
                        driver.close()
                # Switch back to original
                driver.switch_to.window(original_window)
                logger.info("Closed extra tabs, returned to original window")
        except Exception as e:
            logger.warning("Could not clean up tabs: %s", e)

        return video_url

    # ------------------------------------------------------------------
    # Download video
    # ------------------------------------------------------------------

    def _download_video(self, video_url: str, output_path: Path) -> Path:
        """Download the generated video from Grok's asset CDN."""
        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": USER_AGENT,
        }

        resp = self.client.get(video_url, headers=headers, timeout=120)
        resp.raise_for_status()

        output_path.write_bytes(resp.content)
        logger.info("Video saved: %s (%d bytes)", output_path, len(resp.content))
        return output_path

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------


    async def generate(
        self,
        image_path: str | Path,
        video_prompt: str,
        output_path: str | Path,
        on_progress: Callable[[int], None] | None = None,
    ) -> Path:
        """Async wrapper for image-to-video generation.
        
        Args:
            image_path:    Path to the source image (e.g. frame_01.png).
            video_prompt:  Motion/video prompt text typed into the Imagine UI
                           input field before triggering video generation.
            output_path:   Where to save the output .mp4.
            on_progress:   Optional callback receiving progress percentage (0-100).

        Returns:
            Path to the saved video file.
        """
        return await asyncio.to_thread(
            self._generate_sync,
            image_path,
            video_prompt,
            output_path,
            on_progress,
        )

    def _generate_sync(
        self,
        image_path: str | Path,
        video_prompt: str,
        output_path: str | Path,
        on_progress: Callable[[int], None] | None = None,
    ) -> Path:
        """Sync implementation of generate."""
        image_path = Path(image_path)
        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("GrokVideo: %s → %s (prompt: %s)", image_path.name, output_path.name, video_prompt[:60] if video_prompt else "none")

        # Ensure browser is ready
        self._ensure_browser()

        # Generate via UI (upload + prompt + Make video + wait)
        video_url = self._generate_via_ui(image_path, video_prompt, on_progress)

        logger.info("GrokVideo: downloading %s", video_url[:80])

        # Download
        return self._download_video(video_url, output_path)

    async def close(self):
        """Async close wrapper."""
        return await asyncio.to_thread(self._close_sync)

    def _close_sync(self):
        """Sync close implementation."""
        if self._driver is not None:
            try:
                self._driver.quit()
            except Exception:
                pass
            self._driver = None
        if getattr(self, "_virtual_display", None) is not None:
            try:
                self._virtual_display.stop()
            except Exception:
                pass
            self._virtual_display = None
        self.client.close()
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self._close_sync()