import http.client
import json
import requests
import os
import base64
from PIL import Image
from io import BytesIO


def format_nested_structure(data, max_str_len=50, indent=4, _level=0):
    indent_str = ' ' * (_level * indent)

    if isinstance(data, dict):
        result = "{\n"
        for key, value in data.items():
            result += f"{indent_str}{json.dumps(key)}: {format_nested_structure(value, max_str_len, indent, _level + 1)},\n"
        result += indent_str + "}"
        return result

    elif isinstance(data, list):
        result = "[\n"
        for i, item in enumerate(data):
            result += f"{indent_str}{format_nested_structure(item, max_str_len, indent, _level + 1)},\n"
        result += indent_str + "]"
        return result

    elif isinstance(data, str):
        # 智能截断长字符串
        if len(data) > max_str_len:
            truncated = data[:max_str_len // 2] + " ... " + data[-max_str_len // 2:]
            return f"'{truncated}' (length: {len(data)})"
        return json.dumps(data)

    else:
        return str(data)

def download_image(url, save_path):
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        with open(save_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)

        print(f"Save to: {save_path}")
        return True

    except requests.exceptions.RequestException as e:
        print(f"Failed !: {e}")
        return False
    except IOError as e:
        print(f"Failed !: {e}")
        return False
def generate_image(text_prompt,save_path,model_name="doubao-seedream-3-0-t2i-250415"):
    if os.path.isfile(save_path):
        return
    conn = http.client.HTTPSConnection("API_URL")
    if "gemini" in model_name:
        payload = json.dumps({
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {
                            "text": text_prompt
                        }
                    ]
                }
            ],
            "generationConfig": {
                "responseModalities": [
                    "IMAGE"
                ]
            }
        })
        headers = {
            'Authorization': '[API_KEY]',
            'Content-Type': 'application/json'
        }
        conn.request("POST",
                     "/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=[API_KEY]",
                     payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode('utf-8'))
        for part in data['candidates'][0]['content']['parts']:
            if "text" in part.keys():
                print(part['text'])
            if "inlineData" in part.keys():
                base64_str = part['inlineData']['data']
                if base64_str.startswith("data:image/"):

                    base64_str = base64_str.split(",", 1)[1]


                image_bytes = base64.b64decode(base64_str)

                image = Image.open(BytesIO(image_bytes))
                image.save(save_path)
            else:
                print('No Image')
    else:
        payload = json.dumps({
            "model": "doubao-seedream-3-0-t2i-250415",
            "prompt": text_prompt,
            "response_format": "url",
            "size": "1024x1024",
            "guidance_scale": 2.5,
            "watermark": False
        })
        headers = {
            'Authorization': '[API_KEY]',
            'Content-Type': 'application/json'
        }
        conn.request("POST", "/v1/images/generations", payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode("utf-8"))
        if "error" in data.keys():
            # print(data)
            return
        download_image(data['data'][0]['url'], save_path)




def edit_image(base64_image,text_prompt, save_path, model_name="doubao-seededit-3-0-i2i-250628"):
    conn = http.client.HTTPSConnection("yunwu.ai")
    text_prompt = text_prompt.replace("”","*")
    text_prompt = text_prompt.replace("“","*")
    text_prompt = text_prompt.replace("\"","*")
    if "gemini" in model_name:
        payload = json.dumps({
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {
                            "text": text_prompt,
                        },
                        {
                            "inline_data":{
                                "mime_type":"image/jpeg",
                                "data": base64_image[23:]
                            }
                        }
                    ]
                }
            ],
            "generationConfig": {
                "responseModalities": [
                    "TEXT",
                    "IMAGE"
                ]
            }
        })
        headers = {
            'Authorization': '[API_KEY]',
            'Content-Type': 'application/json'
        }
        conn.request("POST",
                     "/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=[API_KEY]",
                     payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode('utf-8'))
        for part in data['candidates'][0]['content']['parts']:
            if "inlineData" in part.keys():
                base64_str = part['inlineData']['data']
                if base64_str.startswith("data:image/"):
                    base64_str = base64_str.split(",", 1)[1]
                image_bytes = base64.b64decode(base64_str)

                image = Image.open(BytesIO(image_bytes))
                image.save(save_path)
            else:
                print('No Image')

    else:

        payload = json.dumps({
            "model": model_name,
            "prompt": text_prompt,
            "image": base64_image,
            "response_format": "url",
            "size": "1k",
            "guidance_scale": 5.5,
            "watermark": False

        })
        headers = {
            'Authorization': '[API_KEY]',
            'Content-Type': 'application/json'
        }
        conn.request("POST", "/v1/images/generations", payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode('utf-8'))

        download_image(data['data'][0]['url'], save_path)

def rearrange_dataset(file, start_index=1000):
    with open(file, 'r', encoding='utf-8') as f:
        dataset = json.load(f)
    for i in range(len(dataset)):
        dataset[i]['id'] = start_index + i
    with open(file, 'w', encoding='utf-8') as f:
        json.dump(dataset, f, ensure_ascii=False, indent=2)
