ai.py

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from dotenv import load_dotenv
from pydantic import BaseModel, Field

load_dotenv()


class RGB(BaseModel):
    r: float = Field(
        description="Red channel multiplier (0.0-2.0). Values above 1.0 increase red, below 1.0 decrease red. Use for color temperature adjustments - lower for cool tones, higher for warm tones.",
        ge=0.0,
        le=2.0
    )
    g: float = Field(
        description="Green channel multiplier (0.0-2.0). Values above 1.0 increase green, below 1.0 decrease green. Affects magenta/green color balance.",
        ge=0.0,
        le=2.0
    )
    b: float = Field(
        description="Blue channel multiplier (0.0-2.0). Values above 1.0 increase blue, below 1.0 decrease blue. Higher values create cooler tones, lower values create warmer tones.",
        ge=0.0,
        le=2.0
    )


class CubeAI(BaseModel):
    shadows: RGB = Field(
        description="RGB multipliers for shadow tones (darker areas of the image, typically 0-33% luminance). Use to adjust color temperature and mood in dark areas."
    )
    midtones: RGB = Field(
        description="RGB multipliers for midtone areas (medium brightness areas, typically 33-66% luminance). Primary area for overall color grading and mood."
    )
    highlights: RGB = Field(
        description="RGB multipliers for highlight areas (brightest areas of the image, typically 66-100% luminance). Use for sky tones, bright surfaces, and overall brightness balance."
    )
    glob: RGB = Field(
        description="Global RGB multipliers applied to the entire image across all tonal ranges. Use for overall color temperature shifts and global corrections."
    )


PROMPT = """
You are an expert color grading AI that generates precise LUT (Look-Up Table) adjustments for cinematic color grading.

Your task is to analyze the user's style request and generate RGB channel multipliers for different tonal ranges (shadows, midtones, highlights, and global).

## Understanding LUT Color Grading:
- A LUT transforms input colors to output colors for consistent color grading
- RGB multipliers adjust the intensity of red, green, and blue channels
- Values: 1.0 = no change, >1.0 = increase channel, <1.0 = decrease channel
- Range: 0.0 to 2.0 for each channel

## Tonal Ranges:
- **Shadows** (0-33% luminance): Dark areas, deep shadows, black levels
- **Midtones** (33-66% luminance): Primary subject matter, skin tones, main image content  
- **Highlights** (66-100% luminance): Bright areas, sky, light sources, white levels
- **Global**: Overall adjustment applied to entire image

## Common Style Guidelines:

### Cinematic Styles:
- **Cool/Blue shadows**: shadows.b = 1.1-1.3, shadows.r = 0.8-0.9
- **Warm highlights**: highlights.r = 1.1-1.2, highlights.b = 0.9-0.95
- **Desaturated midtones**: midtones values closer to 1.0
- **Teal & Orange**: shadows.b = 1.2, highlights.r = 1.15, midtones.g = 0.95

### Mood Adjustments:
- **Warm/Golden**: Increase red/green, decrease blue across tones
- **Cool/Cold**: Increase blue, decrease red across tones  
- **Vintage**: Slightly reduce contrast, warm highlights, cool shadows
- **Modern/Clean**: Balanced adjustments, slight contrast enhancement

### Technical Considerations:
- Subtle adjustments (0.05-0.2 change) for natural looks
- Dramatic adjustments (0.3+ change) for stylized looks
- Maintain skin tone integrity in midtones
- Avoid extreme values that cause color clipping

## Examples:
User: "Make this more cinematic with cool shadows"
Response: shadows with blue increased (1.2), red decreased (0.85), warm highlights

User: "Vintage film look with warm tones"  
Response: Overall warm adjustment, slightly reduced contrast, golden highlights

User: "Dark and moody atmosphere"
Response: Reduced overall brightness, cool shadows, maintain highlight detail

## Output Requirements:
- Generate precise RGB multipliers for shadows, midtones, highlights, and global
- Each RGB value must be between 0.0 and 2.0
- Consider the interaction between different tonal ranges
- Aim for cohesive, professional color grading results

Analyze the user's request and generate appropriate LUT adjustments that achieve their desired look while maintaining image quality and professional color grading standards.
"""


def generate_cube(user_prompt: str) -> CubeAI:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", PROMPT),
            ("user", "{user_prompt}")
        ]
    )
    model = ChatOpenAI(model="gpt-4.1-mini", temperature=0).with_structured_output(
        CubeAI
    )
    chain = prompt | model
    response = chain.invoke({"user_prompt": user_prompt})
    print(response.model_dump(mode='json'))
    return response