import os
import logging
from google import genai
from google.genai import types
from google.adk.tools import ToolContext
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()

# Configure logging
logger = logging.getLogger(__name__)

# ============================================================================
# Helper Functions for Asset Version Management
# ============================================================================

def get_next_version_number(tool_context: ToolContext, asset_name: str) -> int:
    """Get the next version number for a given asset name."""
    asset_versions = tool_context.state.get("asset_versions", {})
    current_version = asset_versions.get(asset_name, 0)
    next_version = current_version + 1
    return next_version


def update_asset_version(tool_context: ToolContext, asset_name: str, version: int, filename: str) -> None:
    """Update the version tracking for an asset."""
    if "asset_versions" not in tool_context.state:
        tool_context.state["asset_versions"] = {}
    if "asset_filenames" not in tool_context.state:
        tool_context.state["asset_filenames"] = {}
    
    tool_context.state["asset_versions"][asset_name] = version
    tool_context.state["asset_filenames"][asset_name] = filename
    
    # Maintain a list of all versions for this asset
    asset_history_key = f"{asset_name}_history"
    if asset_history_key not in tool_context.state:
        tool_context.state[asset_history_key] = []
    tool_context.state[asset_history_key].append({"version": version, "filename": filename})


def create_versioned_filename(asset_name: str, version: int, file_extension: str = "png") -> str:
    """Create a versioned filename for an asset."""
    return f"{asset_name}_v{version}.{file_extension}"


def get_asset_versions_info(tool_context: ToolContext) -> str:
    """Get information about all asset versions in the session."""
    asset_versions = tool_context.state.get("asset_versions", {})
    if not asset_versions:
        return "No renovation renderings have been created yet."
    
    info_lines = ["Current renovation renderings:"]
    for asset_name, current_version in asset_versions.items():
        history_key = f"{asset_name}_history"
        history = tool_context.state.get(history_key, [])
        total_versions = len(history)
        latest_filename = tool_context.state.get("asset_filenames", {}).get(asset_name, "Unknown")
        info_lines.append(f"  • {asset_name}: {total_versions} version(s), latest is v{current_version} ({latest_filename})")
    
    return "\n".join(info_lines)


def get_reference_images_info(tool_context: ToolContext) -> str:
    """Get information about all reference images (current room/inspiration) uploaded in the session."""
    reference_images = tool_context.state.get("reference_images", {})
    if not reference_images:
        return "No reference images have been uploaded yet."
    
    info_lines = ["Available reference images (current room photos & inspiration):"]
    for filename, info in reference_images.items():
        version = info.get("version", "Unknown")
        image_type = info.get("type", "reference")
        info_lines.append(f"  • {filename} ({image_type} v{version})")
    
    return "\n".join(info_lines)


async def load_reference_image(tool_context: ToolContext, filename: str):
    """Load a reference image artifact by filename."""
    try:
        loaded_part = await tool_context.load_artifact(filename)
        if loaded_part:
            logger.info(f"Successfully loaded reference image: {filename}")
            return loaded_part
        else:
            logger.warning(f"Reference image not found: {filename}")
            return None
    except Exception as e:
        logger.error(f"Error loading reference image {filename}: {e}")
        return None


def get_latest_reference_image_filename(tool_context: ToolContext) -> str:
    """Get the filename of the most recently uploaded reference image."""
    return tool_context.state.get("latest_reference_image")


# ============================================================================
# Pydantic Input Models
# ============================================================================

class GenerateRenovationRenderingInput(BaseModel):
    prompt: str = Field(..., description="A detailed description of the renovated space to generate. Include room type, style, colors, materials, fixtures, lighting, and layout.")
    aspect_ratio: str = Field(default="16:9", description="The desired aspect ratio, e.g., '1:1', '16:9', '9:16'. Default is 16:9 for room photos.")
    asset_name: str = Field(default="renovation_rendering", description="Base name for the rendering (will be versioned automatically). Use descriptive names like 'kitchen_modern_farmhouse' or 'bathroom_spa'.")
    current_room_photo: str = Field(default=None, description="Optional: filename of the current room photo to use as reference for layout/structure.")
    inspiration_image: str = Field(default=None, description="Optional: filename of an inspiration image to guide the style. Use 'latest' for most recent upload.")


class EditRenovationRenderingInput(BaseModel):
    artifact_filename: str = Field(default=None, description="The filename of the rendering artifact to edit. If not provided, uses the last generated rendering.")
    prompt: str = Field(..., description="The prompt describing the desired changes (e.g., 'make cabinets darker', 'add pendant lights', 'change floor to hardwood').")
    asset_name: str = Field(default=None, description="Optional: specify asset name for the new version (defaults to incrementing current asset).")
    reference_image_filename: str = Field(default=None, description="Optional: filename of a reference image to guide the edit. Use 'latest' for most recent upload.")


# ============================================================================
# Image Generation Tool
# ============================================================================

async def generate_renovation_rendering(tool_context: ToolContext, inputs: GenerateRenovationRenderingInput) -> str:
    """
    Generates a photorealistic rendering of a renovated space based on the design plan.
    
    This tool uses Gemini 3 Pro's image generation capabilities to create visual 
    representations of renovation plans. It can optionally use current room photos 
    and inspiration images as references.
    """
    if "GEMINI_API_KEY" not in os.environ and "GOOGLE_API_KEY" not in os.environ:
        raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY environment variable not set.")

    logger.info("Starting renovation rendering generation")
    try:
        client = genai.Client()
        
        # Handle inputs that might come as dict instead of Pydantic model
        if isinstance(inputs, dict):
            inputs = GenerateRenovationRenderingInput(**inputs)
        
        # Handle reference images (current room photo or inspiration)
        reference_images = []
        
        if inputs.current_room_photo:
            current_photo_part = await load_reference_image(tool_context, inputs.current_room_photo)
            if current_photo_part:
                reference_images.append(current_photo_part)
                logger.info(f"Using current room photo: {inputs.current_room_photo}")
        
        if inputs.inspiration_image:
            if inputs.inspiration_image == "latest":
                insp_filename = get_latest_reference_image_filename(tool_context)
            else:
                insp_filename = inputs.inspiration_image
            
            if insp_filename:
                inspiration_part = await load_reference_image(tool_context, insp_filename)
                if inspiration_part:
                    reference_images.append(inspiration_part)
                    logger.info(f"Using inspiration image: {insp_filename}")
        
        # Build the enhanced prompt using SLC formula (Subject, Lighting, Camera)
        base_rewrite_prompt = f"""
        Create an ultra-detailed, photorealistic prompt for generating a professional interior design photograph.
        
        Original description: {inputs.prompt}
        
        **CRITICAL REQUIREMENT - PRESERVE EXACT LAYOUT:**
        The generated image MUST maintain the EXACT same room layout, structure, and spatial arrangement described in the prompt:
        - Keep all windows, doors, skylights in their exact positions
        - Keep all cabinets, counters, appliances in their exact positions
        - Keep the same room dimensions and proportions
        - Keep the same camera angle/perspective
        - ONLY change surface finishes: paint colors, cabinet colors, countertop materials, flooring, backsplash, hardware, and decorative elements
        - DO NOT move, add, or remove any structural elements or major fixtures
        
        **Use the SLC Formula for Photorealism:**
        
        1. **SUBJECT (S)** - Be highly specific about details and textures:
           - Describe exact materials with rich adjectives (e.g., "smooth matte white shaker-style cabinets", "honed Carrara marble countertops with subtle grey veining")
           - Include texture details (e.g., "brushed nickel hardware", "wide-plank oak flooring with natural grain")
           - Specify finishes precisely (e.g., "satin finish", "polished", "matte", "textured")
        
        2. **LIGHTING (L)** - Describe lighting conditions that create mood and realism:
           - Natural light sources (e.g., "soft morning sunlight streaming through windows", "golden hour warm glow")
           - Artificial lighting (e.g., "warm LED under-cabinet lighting", "pendant lights casting gentle shadows")
           - Light quality (e.g., "diffused natural light", "dramatic side lighting", "even ambient illumination")
           - Shadows and highlights (e.g., "subtle shadows adding depth", "highlights on polished surfaces")
        
        3. **CAMERA (C)** - Include professional photography specifications:
           - Camera type: "shot on professional DSLR" or "architectural photography camera"
           - Resolution: "8K resolution", "ultra high definition", "HDR"
           - Perspective: specific angle (e.g., "wide-angle lens from doorway", "eye-level perspective", "slightly elevated view")
           - Depth of field: "sharp focus throughout" or "shallow depth of field with background blur"
           - Quality keywords: "professional interior design photography", "magazine quality", "architectural digest style"
        
        **Additional Requirements:**
        - Maintain existing spatial layout and dimensions exactly as described
        - Include specific color names and codes when mentioned
        - Add atmospheric details (e.g., "clean, inviting atmosphere", "modern luxury feel")
        - Specify the aspect ratio: {inputs.aspect_ratio}
        
        **Output Format:** Create a single, flowing paragraph that reads like a professional photography brief. 
        Start with the camera/technical specs, then describe the subject with rich detail, then lighting conditions.
        Include keywords: "photorealistic", "8K", "HDR", "professional interior photography", "architectural photography".
        Emphasize that the layout must remain unchanged - only surface finishes are modified.
        """
        
        if reference_images:
            base_rewrite_prompt += "\n\n**Reference Image Layout:** The reference image shows the EXACT layout that must be preserved. Match the camera angle, room structure, window/door positions, and furniture/appliance placement EXACTLY. Only change the surface finishes and colors. Analyze the lighting in the reference image and replicate it."
        
        # Get enhanced prompt
        rewritten_prompt_response = client.models.generate_content(
            model="gemini-3-pro-preview", 
            contents=base_rewrite_prompt
        )
        rewritten_prompt = rewritten_prompt_response.text
        logger.info(f"Enhanced prompt: {rewritten_prompt}")

        model = "gemini-3-pro-image-preview"
        
        # Build content parts
        content_parts = [types.Part.from_text(text=rewritten_prompt)]
        content_parts.extend(reference_images)

        contents = [
            types.Content(
                role="user",
                parts=content_parts,
            ),
        ]
        
        generate_content_config = types.GenerateContentConfig(
            response_modalities=[
                "IMAGE",
                "TEXT",
            ],
        )

        # Generate versioned filename
        version = get_next_version_number(tool_context, inputs.asset_name)
        artifact_filename = create_versioned_filename(inputs.asset_name, version)
        logger.info(f"Generating rendering with artifact filename: {artifact_filename} (version {version})")

        # Generate the image
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if (
                chunk.candidates is None
                or chunk.candidates[0].content is None
                or chunk.candidates[0].content.parts is None
            ):
                continue
            
            if chunk.candidates[0].content.parts[0].inline_data and chunk.candidates[0].content.parts[0].inline_data.data:
                inline_data = chunk.candidates[0].content.parts[0].inline_data
                
                # Create a Part object from the inline data
                # The inline_data already contains the mime_type from the API response
                image_part = types.Part(inline_data=inline_data)
                
                try:
                    # Save the image as an artifact
                    version = await tool_context.save_artifact(
                        filename=artifact_filename, 
                        artifact=image_part
                    )
                    
                    # Update version tracking
                    update_asset_version(tool_context, inputs.asset_name, version, artifact_filename)
                    
                    # Store in session state
                    tool_context.state["last_generated_rendering"] = artifact_filename
                    tool_context.state["current_asset_name"] = inputs.asset_name
                    
                    logger.info(f"Saved rendering as artifact '{artifact_filename}' (version {version})")
                    
                    return f"✅ Renovation rendering generated successfully!\n\nThe rendering has been saved and is available in the artifacts panel. Artifact name: {inputs.asset_name} (version {version}).\n\nNote: The image is stored as an artifact and can be accessed through the session artifacts, not as a direct image link."
                    
                except Exception as e:
                    logger.error(f"Error saving artifact: {e}")
                    return f"Error saving rendering as artifact: {e}"
            else:
                # Log any text responses
                if hasattr(chunk, 'text') and chunk.text:
                    logger.info(f"Model response: {chunk.text}")
                
        return "No rendering was generated. Please try again with a more detailed prompt."
        
    except Exception as e:
        logger.error(f"Error in generate_renovation_rendering: {e}")
        return f"An error occurred while generating the rendering: {e}"


# ============================================================================
# Image Editing Tool
# ============================================================================

async def edit_renovation_rendering(tool_context: ToolContext, inputs: EditRenovationRenderingInput) -> str:
    """
    Edits an existing renovation rendering based on feedback or refinements.
    
    This tool allows iterative improvements to the rendered image, such as 
    changing colors, materials, lighting, or layout elements.
    """
    if "GEMINI_API_KEY" not in os.environ and "GOOGLE_API_KEY" not in os.environ:
        raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY environment variable not set.")

    logger.info("Starting renovation rendering edit")

    try:
        client = genai.Client()
        
        # Handle inputs that might come as dict instead of Pydantic model
        if isinstance(inputs, dict):
            inputs = EditRenovationRenderingInput(**inputs)
        
        # Get artifact_filename from session state if not provided
        artifact_filename = inputs.artifact_filename
        if not artifact_filename:
            artifact_filename = tool_context.state.get("last_generated_rendering")
            if not artifact_filename:
                return "❌ No artifact_filename provided and no previous rendering found in session. Please generate a rendering first using generate_renovation_rendering."
            logger.info(f"Using last generated rendering from session: {artifact_filename}")
        
        # Validate filename format - check for common hallucination patterns
        if "_v0." in artifact_filename:
            # Version 0 doesn't exist - the first version is always v1
            logger.warning(f"Invalid version v0 detected in filename: {artifact_filename}")
            corrected_filename = artifact_filename.replace("_v0.", "_v1.")
            logger.info(f"Attempting corrected filename: {corrected_filename}")
            artifact_filename = corrected_filename
        
        # Load the existing rendering
        logger.info(f"Loading artifact: {artifact_filename}")
        loaded_image_part = None
        try:
            loaded_image_part = await tool_context.load_artifact(artifact_filename)
        except Exception as e:
            logger.error(f"Error loading artifact: {e}")
        
        # If loading failed, try to find the most recent version of this asset
        if not loaded_image_part:
            # Extract base asset name and try to find any existing version
            base_name = artifact_filename.split('_v')[0] if '_v' in artifact_filename else artifact_filename.replace('.png', '')
            asset_filenames = tool_context.state.get("asset_filenames", {})
            
            # Check if we have any version of this asset
            if base_name in asset_filenames:
                fallback_filename = asset_filenames[base_name]
                logger.info(f"Attempting fallback to known artifact: {fallback_filename}")
                try:
                    loaded_image_part = await tool_context.load_artifact(fallback_filename)
                    if loaded_image_part:
                        artifact_filename = fallback_filename
                        logger.info(f"Successfully loaded fallback artifact: {fallback_filename}")
                except Exception as e:
                    logger.error(f"Fallback load also failed: {e}")
            
            # Last resort: try the last generated rendering
            if not loaded_image_part:
                last_rendering = tool_context.state.get("last_generated_rendering")
                if last_rendering and last_rendering != artifact_filename:
                    logger.info(f"Attempting last resort fallback to: {last_rendering}")
                    try:
                        loaded_image_part = await tool_context.load_artifact(last_rendering)
                        if loaded_image_part:
                            artifact_filename = last_rendering
                            logger.info(f"Successfully loaded last resort artifact: {last_rendering}")
                    except Exception as e:
                        logger.error(f"Last resort load also failed: {e}")
        
        if not loaded_image_part:
            available_renderings = get_asset_versions_info(tool_context)
            return f"❌ Could not find rendering artifact: {inputs.artifact_filename}\n\n{available_renderings}\n\nPlease use one of the available artifact filenames, or generate a new rendering first."

        # Handle reference image if specified
        reference_image_part = None
        if inputs.reference_image_filename:
            if inputs.reference_image_filename == "latest":
                ref_filename = get_latest_reference_image_filename(tool_context)
            else:
                ref_filename = inputs.reference_image_filename
            
            if ref_filename:
                reference_image_part = await load_reference_image(tool_context, ref_filename)
                if reference_image_part:
                    logger.info(f"Using reference image for editing: {ref_filename}")

        model = "gemini-3-pro-image-preview"

        # Build content parts
        content_parts = [loaded_image_part, types.Part.from_text(text=inputs.prompt)]
        if reference_image_part:
            content_parts.append(reference_image_part)

        contents = [
            types.Content(
                role="user",
                parts=content_parts,
            ),
        ]
        
        generate_content_config = types.GenerateContentConfig(
            response_modalities=[
                "IMAGE",
                "TEXT",
            ],
        )

        # Determine asset name and generate versioned filename
        if inputs.asset_name:
            asset_name = inputs.asset_name
        else:
            current_asset_name = tool_context.state.get("current_asset_name")
            if current_asset_name:
                asset_name = current_asset_name
            else:
                # Extract from filename
                base_name = artifact_filename.split('_v')[0] if '_v' in artifact_filename else "renovation_rendering"
                asset_name = base_name
        
        version = get_next_version_number(tool_context, asset_name)
        edited_artifact_filename = create_versioned_filename(asset_name, version)
        logger.info(f"Editing rendering with artifact filename: {edited_artifact_filename} (version {version})")

        # Edit the image
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if (
                chunk.candidates is None
                or chunk.candidates[0].content is None
                or chunk.candidates[0].content.parts is None
            ):
                continue
            
            if chunk.candidates[0].content.parts[0].inline_data and chunk.candidates[0].content.parts[0].inline_data.data:
                inline_data = chunk.candidates[0].content.parts[0].inline_data
                
                # Create a Part object from the inline data
                # The inline_data already contains the mime_type from the API response
                edited_image_part = types.Part(inline_data=inline_data)
                
                try:
                    # Save the edited image as an artifact
                    version = await tool_context.save_artifact(
                        filename=edited_artifact_filename, 
                        artifact=edited_image_part
                    )
                    
                    # Update version tracking
                    update_asset_version(tool_context, asset_name, version, edited_artifact_filename)
                    
                    # Store in session state
                    tool_context.state["last_generated_rendering"] = edited_artifact_filename
                    tool_context.state["current_asset_name"] = asset_name
                    
                    logger.info(f"Saved edited rendering as artifact '{edited_artifact_filename}' (version {version})")
                    
                    return f"✅ Rendering edited successfully!\n\nThe updated rendering has been saved and is available in the artifacts panel. Artifact name: {asset_name} (version {version}).\n\nNote: The image is stored as an artifact and can be accessed through the session artifacts, not as a direct image link."
                    
                except Exception as e:
                    logger.error(f"Error saving edited artifact: {e}")
                    return f"Error saving edited rendering as artifact: {e}"
            else:
                # Log any text responses
                if hasattr(chunk, 'text') and chunk.text:
                    logger.info(f"Model response: {chunk.text}")
                
        return "No edited rendering was generated. Please try again."
        
    except Exception as e:
        logger.error(f"Error in edit_renovation_rendering: {e}")
        return f"An error occurred while editing the rendering: {e}"


# ============================================================================
# Utility Tools
# ============================================================================

async def list_renovation_renderings(tool_context: ToolContext) -> str:
    """Lists all renovation renderings created in this session."""
    return get_asset_versions_info(tool_context)


async def list_reference_images(tool_context: ToolContext) -> str:
    """Lists all reference images (current room photos & inspiration) uploaded in this session."""
    return get_reference_images_info(tool_context)


async def save_uploaded_image_as_artifact(
    tool_context: ToolContext,
    image_data: str,
    artifact_name: str,
    image_type: str = "current_room"
) -> str:
    """
    Saves an uploaded image as a named artifact for later use in editing.
    
    This tool is used when the Visual Assessor detects an uploaded image
    and wants to make it available for the Project Coordinator to edit.
    
    Args:
        tool_context: The tool context
        image_data: Base64 encoded image data or image bytes
        artifact_name: Name to save the artifact as (e.g., "current_room_1", "inspiration_1")
        image_type: Type of image ("current_room" or "inspiration")
    
    Returns:
        Success message with the artifact filename
    """
    try:
        # Create a Part from the image data
        # Note: This assumes image_data is already in the right format
        # In practice, we'll get this from the message content
        
        # Save as artifact
        await tool_context.save_artifact(
            filename=artifact_name,
            artifact=image_data
        )
        
        # Track in state
        if "uploaded_images" not in tool_context.state:
            tool_context.state["uploaded_images"] = {}
        
        tool_context.state["uploaded_images"][artifact_name] = {
            "type": image_type,
            "filename": artifact_name
        }
        
        if image_type == "current_room":
            tool_context.state["current_room_artifact"] = artifact_name
        elif image_type == "inspiration":
            tool_context.state["inspiration_artifact"] = artifact_name
        
        logger.info(f"Saved uploaded image as artifact: {artifact_name}")
        
        return f"✅ Image saved as artifact: {artifact_name} (type: {image_type}). This can now be used for editing."
        
    except Exception as e:
        logger.error(f"Error saving uploaded image: {e}")
        return f"❌ Error saving uploaded image: {e}"

