import time
import uuid

def generate_id():
    return f"chatcmpl-{uuid.uuid4().hex}"

def _format_non_streaming_response(
    generated_text,
    model="custom-model",
    id=None,
    created=None,
    usage=None,
    finish_reason="stop",
):
    if id is None:
        id = generate_id()
    if created is None:
        created = int(time.time())
    
    response = {
        "id": id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": generated_text,
                },
                "finish_reason": finish_reason,
                "logprobs": None,
            }
        ],
    }
    
    if usage is not None:
        response["usage"] = usage
    
    return response

def _format_streaming_response(
    generated_chunks,
    model="custom-model",
    id=None,
    created=None,
    finish_reason="stop",
):
    if id is None:
        id = generate_id()
    if created is None:
        created = int(time.time())
    
    for chunk in generated_chunks:
        yield {
            "id": id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {
                        "content": chunk,
                    },
                    "finish_reason": None,
                    "logprobs": None,
                }
            ],
        }
    
    # Final chunk indicating completion
    yield {
        "id": id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {},
                "finish_reason": finish_reason,
                "logprobs": None,
            }
        ],
    }

def openai_response_format(
    generated_text,
    model="custom-model",
    id=None,
    created=None,
    usage=None,
    finish_reason="stop",
    stream=True,
):
    if stream:
        return _format_streaming_response(
            generated_text, model, id, created, finish_reason
        )
    else:
        return _format_non_streaming_response(
            generated_text, model, id, created, usage, finish_reason
)
from typing import List, Dict, Union, Optional
import base64
from PIL import Image
import io

def openai_to_hf_chat_messages(
    messages: List[Dict[str, str]],
    tokenizer: Optional[object] = None
) -> List[Dict[str, Union[str, Dict]]]:
    """
    Converts OpenAI-style chat messages into Hugging Face chat template format.
    
    Args:
        messages: List of OpenAI-style messages (e.g., [{"role": "user", "content": "Hello"}]).
        model_family: Optional model family (e.g., "llava", "llama") for special handling.
        tokenizer: Optional tokenizer to check for chat template support.
    
    Returns:
        List of messages in Hugging Face chat format.
    """
    hf_messages = []
    
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        
        # Handle multimodal content (e.g., images in OpenAI format)
        if isinstance(content, list):
            # OpenAI-style multimodal: [{"type": "text", "text": "..."}, {"type": "image_url", "image_url": "..."}]
            new_content = []
            for item in content:
                if item["type"] == "text":
                    new_content.append(item["text"])
                elif item["type"] == "image_url":
                    # Handle image (extract base64 or URL)
                    image_url = item["image_url"]["url"]
                    if image_url.startswith("data:image"):
                        # Base64-encoded image
                        image_data = image_url.split(",")[1]
                        image_bytes = base64.b64decode(image_data)
                        image = Image.open(io.BytesIO(image_bytes))
                        new_content.append({"image": image})
                    else:
                        # URL (model must handle downloads)
                        new_content.append({"url": image_url})
            content = " ".join(new_content) if all(isinstance(c, str) for c in new_content) else new_content
        elif not isinstance(content, str):
            raise ValueError(f"Unsupported content type: {type(content)}")
        
        # Add to HF messages
        hf_messages.append({"role": role, "content": content})
    
    # Apply model-specific adjustments
    if tokenizer is not None and hasattr(tokenizer, "apply_chat_template"):
        # Let Hugging Face tokenizer handle further formatting if needed
        try:
            return tokenizer.apply_chat_template(hf_messages, tokenize=False)
        except:
            pass  # Fall back to manual formatting
    
    return hf_messages

def convert_openai_to_hf_messages(openai_messages):
    """
    Converts OpenAI-style chat messages into a format compatible with Hugging Face's
    `tokenizer.apply_chat_template()` function, supporting all modalities (text, images, etc.).
    
    Args:
        openai_messages (list): List of OpenAI-style messages, where each message is a dict with
                                'role' (str) and 'content' (str or list of parts).
    
    Returns:
        list: Hugging Face-compatible messages. Each message is a dict with 'role' and 'content'.
              Content is a string (text-only) or a list of parts (multimodal).
    """
    hf_messages = []
    for msg in openai_messages:
        role = msg['role']
        content = msg['content']
        
        if isinstance(content, list):
            # Handle multimodal content (e.g., text + images)
            converted_content = []
            for part in content:
                if part['type'] == 'text':
                    converted_content.append({'type': 'text', 'text': part['text']})
                elif part['type'] == 'image_url':
                    # Handle image (extract base64 or URL)
                    image_url = part["image_url"]["url"]
                    if image_url.startswith("data:image"):
                        # Base64-encoded image
                        b64_img = image_url.split(",")[1]
                        converted_content.append({
                            'type': 'image',
                            'base64': b64_img
                        })
                    else:
                        # URL (model must handle downloads)
                        converted_content.append({
                            'type': 'image',
                            'url': image_url
                        })
                elif part['type'] == 'video_url':
                    video_url = part["video_url"]["url"]
                    if video_url.startswith("data:video"):
                        ValueError("Base64 video data is not supported in HF format.")
                    else:
                        # URL (model must handle downloads)
                        converted_content.append({
                            'type': 'video',
                            'url': video_url
                        })
                else:
                    raise ValueError(f"Unsupported content type: {part['type']} for conversion.")
            hf_content = converted_content
        else:
            # Text-only content (string)
            hf_content = content
        
        hf_messages.append({'role': role, 'content': hf_content})
    
    return hf_messages
