import asyncio
import json
import os
import io
from pathlib import Path
from bs4 import BeautifulSoup
from playwright.async_api import async_playwright
import google.genai as genai

# Setup environment
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
if not OPENAI_API_KEY:
    print("Warning: OPENAI_API_KEY not found. Asset generation will use placeholders.")

# Initialize GenAI Client
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "AIzaSyCT_rx0NusUJWoqtT7uxHAKEfHo129SJb8")
client = genai.Client(api_key=GEMINI_API_KEY)

SYSTEM_PROMPT_CRITIC = """ROLE: You are the Chief Creative Officer (CCO) at an Awwwards-winning digital agency. Your job is to elevate functional B2B/B2C landing pages into breathtaking, immersive experiences.

CONTEXT: You are reviewing a generated website draft. You are provided with:
1. A high-res full-page screenshot of the rendered page.
2. The raw HTML/Tailwind CSS string.

CRITERIA FOR PREMIUM POLISH:
1. Spatial Rhythm & Glassmorphism: Ensure blur layers (`backdrop-filter`) have correct opacity, gradients have dramatic contrast, and typography features massive scale differences.
2. Kinematics (Motion): Static pages are dead. Identify elements requiring scroll-triggered fade-ups, staggered reveals, or parallax depth.
3. Rich Media: Identify sterile color blocks (like the Hero) that require cinematic looping background videos, authentic lifestyle photography, or bespoke 3D assets.

YOUR TASK: 
Do NOT rewrite the HTML. Output a strict JSON object mapping structural enhancements, requested rich media, and specific CSS/attribute injections. 

EXPECTED JSON SCHEMA:
{
  "visual_score": 85, 
  "critique_summary": "Hero feels flat. Needs cinematic background video and a stagger-fade-up on the typography.",
  "asset_requests": [
    {
      "target_selector": "#hero-section", 
      "type": "video_background | image_background | 3d_floating_asset", 
      "search_query": "cinematic dark moody corporate architecture loop",
      "overlay_class": "bg-black/60 backdrop-blur-sm"
    }
  ],
  "dom_mutations": [
    {
      "target_selector": "h1.main-title", 
      "action": "add_classes", 
      "value": "opacity-0 translate-y-10 data-animate-reveal delay-200"
    }
  ],
  "injected_styles": "/* Complex @keyframes or custom CSS variables to append to <style> */"
}
"""

async def render_dom_to_image(html_content: str) -> bytes:
    """Uses headless browser to render the DOM for Multimodal Critique."""
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        # 1440x900 viewport to simulate standard desktop experiential viewing
        page = await browser.new_page(viewport={"width": 1440, "height": 900})
        await page.set_content(html_content, wait_until="networkidle")
        await page.wait_for_timeout(2000) # Allow CSS/Fonts/Media to paint
        screenshot = await page.screenshot(full_page=True)
        await browser.close()
        return screenshot

async def fetch_rich_media(asset_request: dict) -> str:
    """Mock implementation for the Media Buyer Engine."""
    print(f"    [Media Buyer] Sourcing {asset_request['type']} for query: '{asset_request.get('search_query', '')}'")
    
    # In a full implementation, this hits Pexels API, Unsplash API, or DALL-E 3
    if asset_request['type'] == 'video_background':
        return "https://videos.pexels.com/video-files/853889/853889-hd_1920_1080_25fps.mp4" # Placeholder
    elif asset_request['type'] == 'image_background':
        return "https://images.unsplash.com/photo-1584622650111-993a426fbf0a?auto=format&fit=crop&q=80&w=2000"
    elif asset_request['type'] == '3d_floating_asset':
        return "https://ui-avatars.com/api/?name=3D&background=random" # Placeholder
        
    return ""

async def apply_ast_healing_and_mutations(html_content: str, patch: dict, assets: dict) -> str:
    """Safe DOM manipulation. BeautifulSoup naturally auto-closes broken tags."""
    soup = BeautifulSoup(html_content, 'lxml') 

    # 1. Inject Media Assets safely
    for req in patch.get("asset_requests", []):
        try:
            target = soup.select_one(req["target_selector"])
            if target:
                # Force target to be relative to contain absolute media
                existing_classes = target.get('class', [])
                if 'relative' not in existing_classes:
                    target['class'] = existing_classes + ['relative', 'overflow-hidden']
                    
                if req["type"] == "video_background":
                    video_html = f'<div class="absolute inset-0 z-0 overflow-hidden"><video autoplay loop muted playsinline class="w-full h-full object-cover opacity-60"><source src="{assets[req["target_selector"]]}" type="video/mp4"></video><div class="absolute inset-0 {req.get("overlay_class", "bg-black/60")} mix-blend-multiply"></div></div>'
                    video_tag = BeautifulSoup(video_html, 'html.parser').div
                    target.insert(0, video_tag)
                elif req["type"] == "image_background":
                    img_html = f'<div class="absolute inset-0 z-0 overflow-hidden"><img src="{assets[req["target_selector"]]}" class="w-full h-full object-cover opacity-60"><div class="absolute inset-0 {req.get("overlay_class", "bg-black/60")} mix-blend-multiply"></div></div>'
                    img_tag = BeautifulSoup(img_html, 'html.parser').div
                    target.insert(0, img_tag)
        except Exception as e:
            print(f"    [AST Warning] Failed media injection for {req.get('target_selector', 'unknown')}: {e}")

    # 2. Apply Surgical DOM Mutations (Tailwind / Kinetic attributes)
    for mutation in patch.get("dom_mutations", []):
        try:
            elements = soup.select(mutation["target_selector"])
            for el in elements:
                if mutation["action"] == "add_classes":
                    # BeautifulSoup returns classes as a list, we must combine them safely
                    existing = el.get("class", [])
                    new_classes = mutation["value"].split()
                    # Use a set to prevent duplicates, then cast back to list
                    el["class"] = list(set(existing + new_classes))
        except Exception as e:
            print(f"    [AST Warning] Failed class mutation for {mutation.get('target_selector', 'unknown')}: {e}")

    # 3. Inject Complex CSS (@keyframes, animated gradients)
    if patch.get("injected_styles"):
        style_tag = soup.new_tag("style")
        style_tag.string = patch["injected_styles"]
        if soup.head:
            soup.head.append(style_tag)
        else:
            soup.insert(0, style_tag)

    # Return pure HTML string, BeautifulSoup automatically heals unclosed tags during `str()`
    return str(soup)

async def visual_polish_loop(initial_html: str, max_iterations: int = 2, target_score: int = 92) -> str:
    current_html = initial_html
    
    for iteration in range(max_iterations):
        print(f"\n--- [AIVPL] Starting Polish Iteration {iteration + 1}/{max_iterations} ---")
        
        # 1. Capture Render Reality (The "Eyes")
        print("  [Capturing Render Reality...]")
        screenshot_bytes = await render_dom_to_image(current_html)
        
        # 2. Multimodal Critique
        print("  [Requesting Multimodal Critique from Gemini 3.1 Pro...]")
        
        # We need to construct a multipart request for Gemini
        from google.genai import types
        
        prompt = f"""Review the provided HTML and screenshot against the system instructions.
        
        RAW HTML:
        ```html
        {current_html}
        ```
        """
        
        try:
            response = client.models.generate_content(
                model='gemini-3.1-pro-preview',
                contents=[
                    types.Content(role="user", parts=[
                        types.Part.from_bytes(data=screenshot_bytes, mime_type='image/png'),
                        types.Part.from_text(text=prompt)
                    ])
                ],
                config=types.GenerateContentConfig(
                    system_instruction=SYSTEM_PROMPT_CRITIC,
                    response_mime_type="application/json"
                )
            )
            
            patch_json = json.loads(response.text)
            score = patch_json.get("visual_score", 0)
            print(f"  Visual Score: {score}/100")
            print(f"  Critique: {patch_json.get('critique_summary')}")
            
            if score >= target_score:
                print("✨ Awwwards-Level aesthetics achieved. Breaking loop.")
                break
                
            # 3. Parallel Asset Fetching Engine
            asset_tasks = {req["target_selector"]: fetch_rich_media(req) for req in patch_json.get("asset_requests", [])}
            fetched_assets = {}
            for selector, task in asset_tasks.items():
                fetched_assets[selector] = await task
            
            # 4. Safe AST Code Mutation & Healing
            print("  [Applying AST Mutations...]")
            current_html = await apply_ast_healing_and_mutations(current_html, patch_json, fetched_assets)
            
        except Exception as e:
            print(f"  [Loop Error] Failed during critique/mutation: {e}")
            break

    return current_html

if __name__ == "__main__":
    # Test script stub
    print("AIVPL Core Module loaded.")