# type: ignore
from agents import Agent, RunContextWrapper, function_tool
from model import get_model
import json
import os
from langchain_core.output_parsers import JsonOutputParser
from google import genai

client = genai.Client()

post_schema = """
{
  "meta": {
    "platform": "...", // Instagram | LinkedIn | Twitter | TikTok
    "aspectRatio": "...", // 1:1 | 16:9 | 9:16 | 4:5
    "style": "...", // Modern | Minimal | Bold | Corporate | Playful
    "presetName": "..." // Optional: unique name for this design template
  },
  "brand": {
    "company": "...", // Company or brand name
    "logo": "Include | Exclude",
    "colors": ["#HEX1", "#HEX2", "#HEX3"] // Primary brand colors
  },
  "content": {
    "headline": "...", // Main text (max 8 words)
    "subtext": "...", // Optional supporting line (max 15 words)
    "tone": "Professional | Inspirational | Energetic | Friendly",
    "cta": "..." // Optional call-to-action text
  },
  "visuals": {
    "primarySubject": "...", // Main visual element: product, illustration, icon, infographic
    "elements": ["..."], // Supporting elements: Icons, Charts, Abstract shapes
    "composition": "...", // Centered | Rule of Thirds | Balanced layout
    "lighting": {
      "source": "...", // Soft window, neon, ambient, flat, etc.
      "direction": "...", // Backlight, top-down, side light
      "quality": "Soft | Hard | Diffused | Dramatic"
    },
    "mood": "...", // Energetic, Playful, Minimal, Corporate, etc.
  },
  "design": {
    "typography": "...", // Font style: Bold Sans, Minimal Serif, etc.
    "contrast": "High | Medium | Low",
    "background": "Solid color | Gradient | Abstract | Image",
    "spacing": "..." // Optional: padding/margins guidance
  },
  "colorPalette": {
    "primaryColors": [
      { "name": "...", "hex": "...", "percentage": "..." }
    ],
    "accentColors": [
      { "name": "...", "hex": "...", "percentage": "..." }
    ]
  },
  "finishing": {
    "quality": "High-resolution | Print-ready | Web-optimized",
    "effects": "Subtle shadows | Soft gradients | None"
  }
}
"""

@function_tool
def generate_designSpec_from_brief(design_brief: str) -> dict:
    """Generates a detailed design specification in JSON format from a high-level design brief."""
    schema_template = """
    You are a Professional Social Media Design JSON Generator.  
    Your task:  
    - Read the provided detailed brief describing the desired social media post.  
    - Analyze all aspects: platform, target audience, goal, tone, content, visuals, brand, design style, composition, lighting, color palette, and finishing requirements.  
    - Generate a complete JSON output strictly following this professional schema:  
    {post_schema}
    Rules:
    1. Output only valid JSON; do not include explanations, notes, or extra text.
    2. If any information is missing in the brief, make reasonable professional assumptions.
    3. Use design principles: composition, visual hierarchy, typography, contrast, and color theory.
    4. Include optional fields (like CTA) if relevant based on the brief.
    design_brief:
    {design_brief}
    """
    
    prompt = schema_template.format(design_brief=design_brief, post_schema=post_schema)
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt,
    )
    data = JsonOutputParser().invoke(response.text)
    return data

@function_tool
async def generate_post_image(design_brief: str, designSpec: str):
    """Generates a social media post image based on a design brief and specification."""
    try:
        result = await fal_client.subscribe_async(
            "fal-ai/nano-banana",
            arguments={
                "prompt": image_generate_prompt(design_brief, designSpec),
                "num_images": 1,
                "output_format": "jpeg"
            },
            with_logs=True,
        )
        return result['images']
    except Exception as e:
        print(f"An error occurred during image generation: {e}")

def image_generate_prompt(design_brief, designSpec):
    return f"""
Your goal is to generate attention-grabbing, brand-aligned images that clearly communicate the given message, following design best practices and platform-specific guidelines.

Design Brief:
{design_brief}

Design Specification:
{designSpec}

Design Guidelines:
- Style & Typography: Modern, clean, bold, high contrast; legible fonts, emphasize headlines.
- Composition & Elements: Balanced layout, strong hierarchy, relevant icons/illustrations, avoid clutter.
- Colors & Branding: Brand-relevant palettes, platform-appropriate tones.
- Aspect Ratios & Text: 1:1 (IG/LinkedIn), 16:9 (Twitter/YouTube), 9:16 (Stories/Reels); headline 4-10 words, concise subtext.
- CTA, Lighting & Mood: Place CTA clearly; follow schema hints for mood, lighting, and subtle effects.
- Quality & Output: Ultra-sharp, professional, no watermarks or artifacts; output format per schema (JPG, PNG, etc.).
"""

def media_agent_prompt(context: RunContextWrapper, agent: Agent):
    return """
You are Media Agent, a professional and specialized in creating social media for post. 

Your task:
1. Receive a high-level user brief describing a social media post idea.
2. Generate a detailed DesignSpec (JSON structured specification) from the brief using 'generate_designSpec_from_brief', including platform, style, content, visuals, colors, typography, composition, lighting, mood, and finishing requirements.
3. Using the generated DesignSpec, create a high-quality, brand-aligned social media image using 'generate_post_image' tool, (Don't change the schema use same as generated)

Be concise, professional, and strictly follow the structured DesignSpec and design guidelines provided.
"""

media_agent = Agent(
    name="media_agent",
    instructions=media_agent_prompt,
    model=get_model('gemini-2.0-flash'),
    tools=[
        generate_post_image,
        generate_designSpec_from_brief,
    ]
)