from cloudgpt_utils import get_chat_completion, encode_image
from tenacity import retry, stop_after_attempt, wait_random_exponential, retry_if_exception_type
import time
import random
import os


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(60))
def openai_completion_text(prompt, engine="gpt-35-turbo-1106", max_tokens=1024, temperature=0):
    chat_message = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": prompt
                }
            ]
        }
    ]

    resp = get_chat_completion(
        engine=engine,
        messages=chat_message,
        max_tokens=max_tokens,
        timeout=10,
        # request_timeout=10,
        temperature=temperature,
        stop=["\n\n", "<|endoftext|>"]
    )
    # print(resp.choices[0].message)
    # print("\n%s: %s" % (engine, resp.choices[0].message.content))
    return resp.choices[0].message.content


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(60))
def openai_completion_vision(prompt, image, engine="gpt-4-turbo-20240409", max_tokens=700, temperature=0):
    attempt = 0
    max_attempts = 10
    # make multiple attempts to handle occasional failures
    while attempt < max_attempts:
        try:
            return attempt_openai_completion(prompt, image, engine, max_tokens, temperature)
        except Exception as e:
            attempt += 1
            if attempt < max_attempts:
                wait_time = random.randint(1, 60)
                print(f"Waiting for {wait_time} seconds before retrying...")
                time.sleep(wait_time)
            else:
                print(f"Error after multiple attempts with engine {engine}: {str(e)}")
                print("Switching to gpt-4-turbo-20240409 due to multiple failures.")
                return attempt_openai_completion(prompt, image, "gpt-4-turbo-20240409")


def attempt_openai_completion(prompt, image, engine="gpt-4-turbo-20240409", max_tokens=700, temperature=0):
    image_url = encode_image(image)
    chat_message = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": prompt
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ]
        }
    ]

    resp = get_chat_completion(
        engine=engine,
        messages=chat_message,
        max_tokens=max_tokens,
        timeout=10,
        # request_timeout=10,
        temperature=temperature,
        stop=["\n\n", "<|endoftext|>"]
    )
    # print(resp.choices[0].message)
    # print(resp)
    print("\n%s: %s" % (engine, resp.choices[0].message.content))
    return resp.choices[0].message.content


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(60))
def openai_completion_vision_CoT(sys_prompt, user_prompt, image, engine="gpt-4-turbo-20240409", api_key=None,
                                 base_url=None, max_tokens=12800, temperature=0):
    global_attempt, local_attempt = 0, 0
    global_max_attempts, local_max_attempts = 2, 3
    # make multiple attempts to handle occasional failures
    while global_attempt < global_max_attempts:
        try:
            try:
                return attempt_openai_completion_CoT(sys_prompt, user_prompt, image, engine, api_key, base_url,
                                                     max_tokens, temperature)
            except Exception as e:
                local_attempt += 1
                if local_attempt < local_max_attempts:
                    wait_time = random.randint(1, 60)
                    print(f"Waiting for {wait_time} seconds before retrying...")
                    time.sleep(wait_time)
                else:
                    print(f"Error after multiple attempts with engine {engine}: {str(e)} with sample {image}")
                    print("Switching to gpt-41-0414-global due to multiple failures.")
                    return attempt_openai_completion_CoT(sys_prompt, user_prompt, image, "gpt-41-0414-global")
        except Exception as e:
            global_attempt += 1
            if global_attempt == global_max_attempts:
                print(
                    f"Bad Error after multiple attempts with engine {engine}: {str(e)} with sample {image}, return None!")
                return ""


def attempt_openai_completion_CoT(sys_prompt, user_prompt, image, engine="gpt-41-0414-global", api_key=None,
                                  base_url=None, max_tokens=16384, temperature=0):
    image_url = encode_image(image)
    chat_message = [
        {
            "role": "system",
            "content": sys_prompt
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": user_prompt
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ]
        }
    ]

    # for reasoning model
    if engine[0] == "o":
        resp = get_chat_completion(
            api_key=api_key,
            base_url=base_url,
            engine=engine,
            messages=chat_message,
            max_tokens=max_tokens,
            timeout=256,
            # request_timeout=10,
            temperature=temperature,
            # stop=["\n\n", "<|endoftext|>"]
        )
    else:
        resp = get_chat_completion(
            api_key=api_key,
            base_url=base_url,
            engine=engine,
            messages=chat_message,
            max_tokens=max_tokens,
            timeout=128,
            # request_timeout=10,
            temperature=temperature,
            stop=["\n\n", "<|endoftext|>"]
        )

    # print(resp.choices[0].message)
    # print(resp)
    # print("\n%s: %s" % (engine, resp.choices[0].message.content))
    return resp.choices[0].message.content


def attempt_openai_completion_CoT_multi_images(
        sys_prompt,
        user_prompt,
        images,  # Can be a single image or list of images
        image_context=None,  # Additional context about images (search/edit results)
        engine="gpt-41-0414-global",
        api_key=None,
        base_url=None,
        max_tokens=16384,
        temperature=0
):
    """
    Handle multiple images including original, searched, and edited images

    Parameters:
    -----------
    sys_prompt : str
        System prompt
    user_prompt : str
        User prompt text
    images : str, list, or dict
        Can be:
        - Single image path/url
        - List of image paths/urls
        - Dict with 'original' and optionally 'search'/'edited' keys
    image_context : dict, optional
        Context from tools_pool operations (search results, edit info)
    engine : str
        Model engine to use
    api_key : str
        API key
    base_url : str
        API base URL
    max_tokens : int
        Maximum tokens for response
    temperature : float
        Temperature for generation

    Returns:
    --------
    str
        Model response content
    """

    # Build content array starting with text
    content = [
        {
            "type": "text",
            "text": user_prompt
        }
    ]

    # Process images based on type
    if isinstance(images, dict):
        # Structured format with original/search/edited

        # Add original image
        if 'original' in images:
            content.append({
                "type": "text",
                "text": "\n[Original Image]"
            })
            try:
                image_url = encode_image(images['original'])
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_url}
                })
            except Exception as e:
                print(f"Warning: Failed to encode original image: {str(e)[:100]}")

        # Add search results if present
        if 'search' in images and images['search']:
            search_items = images['search'] if isinstance(images['search'], list) else [images['search']]
            for i, search_item in enumerate(search_items):
                if isinstance(search_item, dict):
                    title = search_item.get('title', f'Search Result {i + 1}')
                    path = search_item.get('path', search_item.get('local_image_path'))
                else:
                    title = f'Search Result {i + 1}'
                    path = search_item

                if path and os.path.exists(path):
                    content.append({
                        "type": "text",
                        "text": f"\n[Search Reference {i + 1}: {title}]"
                    })
                    try:
                        image_url = encode_image(path)
                        content.append({
                            "type": "image_url",
                            "image_url": {"url": image_url}
                        })
                    except Exception as e:
                        print(f"Warning: Failed to encode search image {i + 1}: {str(e)[:100]}")

        # Add edited image if present
        if 'edited' in images and images['edited']:
            content.append({
                "type": "text",
                "text": "\n[Edited Image]"
            })
            try:
                image_url = encode_image(images['edited'])
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_url}
                })
            except Exception as e:
                print(f"Warning: Failed to encode edited image: {str(e)[:100]}")

    elif isinstance(images, list):
        # List format - process each image
        for i, image_item in enumerate(images):
            if isinstance(image_item, tuple) and len(image_item) == 2:
                image, caption = image_item
                if caption:
                    content.append({
                        "type": "text",
                        "text": f"\n[{caption}]"
                    })
            else:
                image = image_item

            try:
                image_url = encode_image(image)
                content.append({
                    "type": "image_url",
                    "image_url": {"url": image_url}
                })
            except Exception as e:
                print(f"Warning: Failed to encode image {i + 1}: {str(e)[:100]}")

    else:
        # Single image
        try:
            image_url = encode_image(images)
            content.append({
                "type": "image_url",
                "image_url": {"url": image_url}
            })
        except Exception as e:
            print(f"Warning: Failed to encode image: {str(e)[:100]}")

    # Add any additional context information
    if image_context:
        context_text = "\n\n[Additional Context]:\n"
        if 'search_query' in image_context:
            context_text += f"Search Query: {image_context['search_query']}\n"
        if 'edit_instruction' in image_context:
            context_text += f"Edit Instruction: {image_context['edit_instruction']}\n"
        if 'tool_results' in image_context:
            context_text += f"Tool Results: {image_context['tool_results']}\n"

        content.insert(1, {  # Insert after main prompt
            "type": "text",
            "text": context_text
        })

    # Build the chat message
    chat_message = [
        {
            "role": "system",
            "content": sys_prompt
        },
        {
            "role": "user",
            "content": content
        }
    ]

    # Call API with appropriate settings
    if engine[0] == "o":  # Reasoning models
        resp = get_chat_completion(
            api_key=api_key,
            base_url=base_url,
            engine=engine,
            messages=chat_message,
            max_tokens=max_tokens,
            timeout=256,
            temperature=temperature,
        )
    else:
        resp = get_chat_completion(
            api_key=api_key,
            base_url=base_url,
            engine=engine,
            messages=chat_message,
            max_tokens=max_tokens,
            timeout=128,
            temperature=temperature,
            stop=["\n\n", "<|endoftext|>"]
        )

    if resp and hasattr(resp, 'choices') and resp.choices:
        return resp.choices[0].message.content
    else:
        print(f"Failed to get valid response from {engine}")
        return ""


# Specialized function for tool integration
def attempt_openai_completion_with_tools(
        sys_prompt,
        user_prompt,
        original_image,
        tool_results=None,  # Results from tools_pool operations
        engine="gpt-41-0414-global",
        api_key=None,
        base_url=None,
        max_tokens=16384,
        temperature=0
):
    """
    Integrate with tools_pool search and edit results

    Parameters:
    -----------
    original_image : str
        Path to the original image
    tool_results : dict
        Results from tools_pool containing 'search' and/or 'edit' results
        Example: {
            'search': {'images': [...], 'text': '...'},
            'edited_image': 'path/to/edited.png',
            'tool_type': 'searching' or 'image_editing',
            'query': 'search query or edit instruction'
        }
    """

    # Prepare images dictionary
    images_dict = {'original': original_image}
    image_context = {}

    if tool_results:
        # Handle search results
        if 'search' in tool_results and tool_results['search']:
            search_data = tool_results['search']
            if isinstance(search_data, dict) and 'images' in search_data:
                images_dict['search'] = search_data['images']
                # Add search text to prompt if available
                if 'text' in search_data:
                    user_prompt = f"{user_prompt}\n\n{search_data['text']}"
            elif isinstance(search_data, list):
                images_dict['search'] = search_data

        # Handle edited image
        if 'edited_image' in tool_results and tool_results['edited_image']:
            images_dict['edited'] = tool_results['edited_image']

        # Add context
        if 'tool_type' in tool_results:
            image_context['tool_type'] = tool_results['tool_type']
        if 'query' in tool_results:
            if tool_results.get('tool_type') == 'searching':
                image_context['search_query'] = tool_results['query']
            elif tool_results.get('tool_type') == 'image_editing':
                image_context['edit_instruction'] = tool_results['query']

    return attempt_openai_completion_CoT_multi_images(
        sys_prompt=sys_prompt,
        user_prompt=user_prompt,
        images=images_dict,
        image_context=image_context,
        engine=engine,
        api_key=api_key,
        base_url=base_url,
        max_tokens=max_tokens,
        temperature=temperature
    )