from __future__ import annotations
from datetime import UTC, datetime
from pydantic import model_validator, Field, BaseModel  # type: ignore
from typing import (
    Dict,
    Literal,
    ClassVar,
    cast,
    Any,
    AsyncGenerator,
    List,
    Optional
#    override
)  # NOTE: You must import Optional for late binding to work
import inspect
import random
import re
import json
import asyncio
import time
import asyncio
import time
import os
import hashlib

from .base import Agent, agent_registry, LLMMessage
from models import ActivityType, ApiActivityType, Candidate, ChatMessage, ChatMessageError, ChatMessageMetaData, ApiMessageType, ChatMessageStatus, ChatMessageStreaming, ChatMessageUser, ChatOptions, ChatSenderType, ApiStatusType, Tunables
import model_cast
from logger import logger
import defines
import backstory_traceback as traceback

from image_generator.image_model_cache import ImageModelCache
from image_generator.profile_image import generate_image, ImageRequest

seed = int(time.time())
random.seed(seed)

class ImageGenerator(Agent):
    agent_type: Literal["generate_image"] = "generate_image"  # type: ignore
    _agent_type: ClassVar[str] = agent_type  # Add this for registration
    agent_persist: bool = False

    system_prompt: str = "" # No system prompt is used

    async def generate(
        self, llm: Any, model: str, 
        session_id: str, prompt: str, 
        tunables: Optional[Tunables] = None, 
        temperature=0.7
    ) -> AsyncGenerator[ChatMessage | ChatMessageStatus | ChatMessageError | ChatMessageStreaming, None]:
        if not self.user:
            logger.error("User is not set for ImageGenerator agent.")
            raise ValueError("User must be set before generating images.")
        user = self.user

        file_path = os.path.join(defines.user_dir, user.username, "profile.png")
        try:
            #
            # Generate the profile picture
            #
            status_message = ChatMessageStatus(
                session_id=session_id,
                activity=ApiActivityType.GENERATING_IMAGE,
                content=f"Generating {prompt}...",
            )
            yield status_message

            logger.info(f"Image generation: {file_path} <- {prompt}")
            request = ImageRequest(filepath=file_path, session_id=session_id, prompt=prompt, iterations=4, height=256, width=256, guidance_scale=7.5)
            generated_message = None
            async for generated_message in generate_image(
                request=request 
            ):
                if generated_message.status == ApiStatusType.ERROR:
                    yield generated_message
                    return
                if generated_message.status != ApiStatusType.DONE:
                    yield generated_message
                    continue

            if generated_message is None:
                error_message = ChatMessageError(
                    session_id=session_id,
                    content="Image generation failed to produce a valid response."
                )
                logger.error(f"⚠️ {error_message.content}")
                yield error_message
                return

            logger.info("Image generation done...")

            user.profile_image = "profile.png"
            
            # Image generated
            generated_image = ChatMessage(
                session_id=session_id,
                status=ApiStatusType.DONE,
                content = f"{defines.api_prefix}/profile/{user.username}",
                metadata=generated_message.metadata
            )
            yield generated_image
            return
        
        except Exception as e:
            error_message = ChatMessageError(
                session_id=session_id,
                content=f"Error generating image: {str(e)}"
            )
            logger.error(traceback.format_exc())
            logger.error(f"⚠️ {error_message.content}")
            yield error_message
            return

# Register the base agent
agent_registry.register(ImageGenerator._agent_type, ImageGenerator)
