
import json
from openai import OpenAI
from tqdm import tqdm
import base64
import uuid
import mimetypes
import os
import requests
from datetime import datetime
import asyncio

from multiprocessing import Pool, cpu_count
from jimeng import generate_from_image,generate_from_prompt
import re
from PIL import Image, ImageDraw, ImageFont
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--model_name', type=str, default='gemini-2.5-flash', help='Model name for CCS content generation')
parser.add_argument('--ccs_content_save_dir', type=str, default='./ccs_content', help='')
parser.add_argument('--ccs_image_save_dir', type=str, default='./generated_dataset7', help='')
parser.add_argument('--labels', type=str, default='imagenet_labels4.txt')
parser.add_argument('--ccs_content', type=str, default='None')

# parser.add_argument('--save_dir', type=str, default='.', help='')
# parser.add_argument('--prompt_type', type=str, default='wo', help='wo / normal / sbs')
args = parser.parse_args()

with open("env.json",'r') as f:
    api_keys = json.load(f)

    client = OpenAI(
        base_url="",
        api_key=api_keys['SEEDREAM_API_KEY'],
        timeout=120
    )

def encode_image(image_path):
    if image_path.startswith("http"):
        user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
        request_kwargs = {
            "headers": {"User-Agent": user_agent},
            "stream": True,
        }

        # Send a HTTP request to the URL
        response = requests.get(image_path, **request_kwargs)
        response.raise_for_status()
        content_type = response.headers.get("content-type", "")

        extension = mimetypes.guess_extension(content_type)
        if extension is None:
            extension = ".download"

        fname = str(uuid.uuid4()) + extension
        download_path = os.path.abspath(os.path.join("downloads", fname))

        with open(download_path, "wb") as fh:
            for chunk in response.iter_content(chunk_size=512):
                fh.write(chunk)

        image_path = download_path

    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"
}



def get_generation_instruction_prompt(prompt):
    prompt = "Keep the whole image unchanged, " + prompt
    return prompt

def get_init_image_prompt_v2(setting_dict,scene):
    client = OpenAI(
        base_url="",
        api_key=api_keys['SEEDREAM_API_KEY'],
        timeout=120
    )
    print(setting_dict)
    object = setting_dict["object"] + "(" + setting_dict["condition"] + ")"
    if scene == 0:
        messages = [
            {
                'role': 'user',
                'content': f"""
       Create a prompt for the object {object}. I will use this prompt to generate images via Stable Diffusion.
Note: Emphasize the object and its modified features in parentheses (adjust the content in parentheses to reflect the altered characteristics of the object in the image). The generated image must meet the requirements described in the parentheses while ensuring realism.
Additionally: Do not use depth-of-field keywords. The {object} must be at a certain distance from the camera and occupy only a small portion of the frame (make sure the prompt specifies this!). Based on the object and its features, specify a suitable, complex scene in which the object is only part of the environment, and emphasize natural realism.
Note: Generate only the prompt, without extra explanations. If this is a trademark, generate the corresponding product/object.
Example:
Input: Guitar (6 strings clear visible from this view point)
Ultra-realistic photograph, showing an acoustic guitar with five strings clearly visible, placed in a cozy, vintage music studio. The guitar leans against a retro leather armchair, surrounded by other musical elements: a standing microphone, a glowing tube amplifier, sheet music scattered on the wooden floor, and headphones hanging on a nearby stool. Soft natural sunlight streams through large windows with thin curtains, illuminating dust particles in the air and highlighting the complex details of the guitar’s spruce top, rosewood fretboard, and nickel strings. The room has exposed brick walls, a shelf with vinyl records, and a Persian rug, creating a warm, nostalgic atmosphere, professional lighting, sharp focus, extremely detailed, realistic style, capturing the essence of a creative space.
                    """
            }
        ]
        response = client.chat.completions.create(
            model=args.model_name,
            messages=messages
        )
        response = response.choices[0].message.content
    # print('#PROMPT',response)
    return response

def get_rule_v2(object):
    client = OpenAI(
        base_url="",
        api_key=api_keys['OPENAI_API_KEY'],
        timeout=120
    )
    message = [
        {
            "role":"user",
            "content":f"""
I now need the following content: originally, a certain object in an image is real and normal, but now I want to make some modifications to this object so that the overall appearance remains unchanged, yet the details are contrary to reality and inconsistent with the real object. Below, I will first introduce several common modification methods.

    1. Quantity modification: by adding or removing a certain key structure of the object, the modified object becomes unrealistic. Note: you need to modify according to the uniqueness of the object. For example, animals have four legs, so for a dog, changing the number of legs would not be unique. This modification must make the result counter-intuitive!
        For example: ① For starfish: add an arm to make a six-armed starfish. (Explanation: in reality, starfish only have five arms; this modification is not visually obvious but changes a key detail.) ② For a Western fork: change a fork from four tines to five tines. (Explanation: standard forks have four tines, so five is unrealistic, a subtle detail modification.) ③ For a clock: change the clock hands to four. (Explanation: clocks with four hands do not exist in reality.) ④ For a snowflake: change it from six branches to eight branches. (Explanation: due to water molecule arrangement, real snowflakes always have six branches; this modification focuses on branch count.) ⑤ For a guitar: change it from six strings to five strings. (Explanation: standard guitars have six strings; changing string number is a detail modification.)
    2. Color modification: by changing the color of a small part of the object, the new object becomes inconsistent with reality in details. Note: the color change must focus on a small part, not a large area. For example, you cannot turn a zebra entirely black because that is too broad and does not meet the requirement of unconventional details.
        For example: ① For a rainbow: change the innermost color to red. (Explanation: rainbows strictly follow red-orange-yellow-green-blue-indigo-violet from outer to inner; the innermost cannot be red, and only one layer is changed.) ② For a traffic light: change the bottom bulb to blue. (Explanation: traffic lights are red-yellow-green from top to bottom; changing the green to blue is unrealistic.)

    3. Shape modification: some objects have fixed shapes in reality. Changing a structure to another shape may violate common sense, physical laws, or human perception. Note: if the object is large, only one part can be modified, not the whole shape, to comply with detail modification rules.
        For example: ① For a ladybug: change the back spots to triangles. (Explanation: spots are originally round.) ② For a front-loading washing machine: change the window from circular to square. (Explanation: washing machine windows are circular.)

    4. Key structure deletion: some objects have key features that distinguish them from others. By deleting this structure, the overall object remains, but it loses a key detail. The key part must be unique and visually perceivable, but not too large.
        For example: ① For a chicken: remove the comb while keeping everything else. (Explanation: the comb is a unique, small feature.) ② For a clock: remove all hands. (Explanation: clocks without hands do not exist.) ③ For a bicycle: remove the pedals. (Explanation: pedals are key parts for all bicycles.) ④ For a teapot: remove the spout. (Explanation: teapots need a spout to pour.) ⑤ For a tower crane: remove the operator cabin. (Explanation: without it, the crane loses function.)

    5. Key structure replacement: replace a key structure of the object with a structure from another object. The object remains recognizable overall, but details conflict with reality. Note: the replaced and target objects must be somewhat similar; for example, do not replace a dog’s nose with an elephant’s trunk.
        For example: ① For a chicken: replace its claws with duck webbed feet. (Explanation: chickens have no webbed feet, but chickens and ducks are similar types of birds; the modification focuses on detail and violates common sense.) ② For a pen: replace the nib with a pencil tip. (Explanation: similar shapes, but the tip differs significantly in form.)

    6. Logical errors: for objects with structured or rule-based designs, you can create rule violations through modifications. Note: only detail-level errors are allowed. Physical rule violations are also logical errors but must be visible in the image and not too obvious. This includes all illogical modifications.
        For example: ① For a dice: change the top face to seven. (Explanation: dice only have 1–6.) ② For a keyboard: change the A key to X. (Explanation: the keyboard has 26 keys; duplicates are impossible.) ③ For a clock: change 12 to 13, or replace 6, 9, 10 with 14. (Explanation: clocks are divided into 12 parts.) ④ For a punching bag: remove the hanging chain but keep the fixed rod. (Explanation: the bag appears to float, violating common sense; detail-level change.) ⑤ For chainmail: add a zipper. (Explanation: chainmail is ancient armor; zippers are modern, so this is unrealistic and a detail addition.)

The above are six examples, but other unconventional modifications exist, e.g., printing multiple suits on a playing card or adding nine directional markers to a compass. You need to carefully consider the object’s characteristics to devise clever modifications.

Here, you should consider the object: {object}. First analyze its features, including color, shape, key structures, and their quantity. Identify which one or several modification strategies may apply. After careful analysis, list the feasible modification methods for this object. Your answer should follow this format:
[object]: [A one-sentence description of the object, including the mentioned features]. Analysis: [Detailed analysis of which aspects can be modified].
<answer><condition>[The object’s required form in the image, according to the modification rule. Example: for deleting piano black keys: top-down view, black keys clearly visible.]</condition> <rule>[Modification rule: remove, add, or change a feature]</rule><description>[Rationale for this modification]</description><name>[File name to save]</name></answer>
Below are two examples to learn from and imitate:
######################
Button: round, with two rows and two columns of four holes. Analysis: The button’s key feature is the central 2*2 small holes. Quantity modification can be applied by changing the four holes to five, violating common sense but keeping overall appearance.
    1. <answer><condition>Button, top-down view, four central holes</condition> <rule>Keep image overall unchanged, change four holes to five</rule> <description>Buttons in reality have four holes; five holes violate common sense without changing the object’s overall form; aligns with detail modification.</description><name>button_with_5_poles</name></answer>

Violin: pear-shaped body, with f-holes, shoulder rest, four strings. Analysis: f-holes are unique structural details. Shape modification can remove the shoulder rest, a key structure deletion, or remove the f-holes, also a key deletion. Quantity modification can change four strings to three.
    1. <answer><condition>Violin, four strings clearly visible</condition> <rule>Change violin from four to three strings</rule> <description>Standard violins have four strings; three-string violins do not exist, a detail-level quantity modification.</description><name>violin_with_3_strings</name></answer>
    2. <answer><condition>Violin, front view, two f-holes visible</condition> <rule>Remove the f-holes</rule> <description>F-holes are essential for sound; key structure deletion.</description><name>violin_without_f</name></answer>
    3. <answer><condition>Violin, left-side view, shoulder rest visible</condition> <rule>Remove the shoulder rest</rule> <description>The shoulder rest is a characteristic small component; key structure deletion.</description><name>violin_without_shouldering</name></answer>

Requirements:
    1. Not too difficult or forced; modifications must be noticeable to ordinary people. Provide only 4–5 optimal modifications. If the object has few feasible modifications, only output the best one.
    2. Modifications should not be too professional or academic. For detailed objects (animals, tools), modify features characteristic of their class, not niche subcategories. Ordinary people should perceive the anomaly.
    3. Modifications must not result in the object becoming another normal object. For example, removing zebra stripes should not make it look like a regular horse.
    4. Focus on shape modifications that either violate conventions or contradict common sense. Avoid “formulaic” changes, e.g., changing an owl’s eyes to blue is imaginable and not unique.
    5. Modifications must be drawable and sufficiently detailed to guide image changes.
    6. For animals, focus on their species type. For birds of prey, modify features characteristic of birds or raptors, not arbitrary color changes.
Now, based on the above guidance and examples, analyze how the {object} can be modified, and answer in the specified format only.
            """
        }
    ]
    response = client.chat.completions.create(
        model=args.model_name,
        messages=message
    )
    response = response.choices[0].message.content
    _answer_list = re.findall(r'<answer>(.*?)</answer>',response)
    # print('WELCOME TO THE HOTEL CALIFORNIA',_answer_list)
    # print(object,"#######\n",response)
    answer_list = []
    for answer in _answer_list:
        condition = re.findall(r'<condition>(.*?)</condition>',answer)
        rule = re.findall(r'<rule>(.*?)</rule>',answer)
        name = re.findall(r'<name>(.*?)</name>', answer)
        answer_list.append(
            {
                "object":object,
                "condition":condition[0],
                "rule":"Keep the whole image unchanged, " + rule[0],
                "name":name[0] if 0 < len(name) else "null",
                "prompt":get_init_image_prompt_v2(
                    {
                        "object":object,
                        "condition": condition[0],
                        "rule": "Keep the whole image unchanged, " + rule[0],
                        "name": name[0] if 0 < len(name) else "null",
                    },
                    0
                ),
                "use": False
            }
        )
    return answer_list

def fully_automated_generation():
    def load_labels(labels_file):
        with open(labels_file,encoding='utf-8') as reader:
            f = reader.read()
            labels = [line.strip() for line in f.splitlines() if line.strip()]
        return labels
    if args.ccs_content == "None":
        object_list = load_labels(args.labels)
        print(f"Object Len: {len(object_list)}")
        with Pool(processes=int(cpu_count() * 0.9)) as pool:
            json_answer = list(tqdm(pool.imap(get_rule_v2, object_list), total=len(object_list)))
        # json_answer = []
        with open(os.path.join(args.ccs_content_save_dir,f'ccs_content_{args.labels}_{args.model_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.json'), "w", encoding='utf-8') as f:
            json.dump(json_answer,f,ensure_ascii=False,indent=2)
        print('CCS content saved.')
    else:
        with open(args.ccs_content,'r',encoding='utf-8') as f:
            json_answer = json.load(f)
        print("CCS already loaded.")
if __name__ == '__main__':

    fully_automated_generation()
