from __future__ import annotations
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 traceback
import asyncio
import time
import asyncio
import time
import os
import hashlib

from . base import Agent, agent_registry, LLMMessage
from .. message import Message
from .. rag import ChromaDBGetResponse
from .. setup_logging import setup_logging
from .. profile_image import generate_image, ImageRequest
from .. import defines
from .. user import User

logger = setup_logging()

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

class ImageGenerator(Agent):
    agent_type: Literal["image"] = "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
    username: str
    filename: str = ""

    llm: Any = Field(default=None, exclude=True)
    model: str = Field(default=None, exclude=True)
    user: Dict[str, Any] = Field(default={}, exclude=True)

    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
        message.status = "done"
        yield message
        return

    async def process_message(
        self, llm: Any, model: str, message: Message
    ) -> AsyncGenerator[Message, None]:
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
        if not self.context:
            raise ValueError("Context is not set for this agent.")

        self.llm = llm
        self.model = model

        prompt = message.prompt
        user_info = os.path.join(defines.user_dir, self.username, defines.user_info_file)
        with open(user_info, "r") as f:
            self.user = json.loads(f.read())
        images = self.user.get("images", None)
        if not images:
            self.user["images"] = {}
            images = self.user["images"]
            with open(user_info, "w") as f:
                f.write(json.dumps(self.user))

        if self.filename != "profile.png":
            # Convert the prompt to a hash using MD5
            hash_object = hashlib.md5(prompt.encode('utf-8'))
            # Get the hexadecimal representation of the hash
            hash_string = hash_object.hexdigest()
            # Return the filename with the specified extension
            self.filename = f"{hash_string}.png"
            file_path = os.path.join(defines.user_dir, self.username, "images")
            os.makedirs(file_path, exist_ok=True)
            file_path=os.path.join(file_path, self.filename)
        else:
            file_path = os.path.join(defines.user_dir, self.username, self.filename)
        
        images[self.filename] = { "status": "thinking", "response": "Generating image" }

        if self.context.processing:
            tick: int = 0
            while self.context.processing:
                with open(user_info, "r") as f:
                    self.user = json.loads(f.read())
                image = self.user["images"].get(self.filename, { "status": "waiting", "response": f"Waiting for image generation slot.{'.' * tick}"})
                message.status = image["status"]
                message.response = image["response"]
                tick = (tick + 1) % 5
                yield message
                await asyncio.sleep(1)  # Allow the event loop to process the write

            # Processing of active image is complete. Check if it was this image, and return it if so
            with open(user_info, "r") as f:
                self.user = json.loads(f.read())
            image = self.user["images"].get(self.filename, None)
            if image:
                message.status = image["status"]
                message.response = image["response"]
                yield message
                return

        self.context.processing = True

        try:
            #
            # Generate the profile picture
            #
            message.status = "thinking"
            message.response = f"Generating: {prompt}"
            yield message

            logger.info("Beginning image generation...", self.user)
            if os.path.exists(file_path) and self.filename != "profile.png":               
                logger.info(f"Image already exists: {file_path}")
                message.status = "done"
                message.response = f"/api/u/{self.username}/images/{self.filename}"
                self.user["images"][self.filename] = { "status": message.status, "response": message.response}
                with open(user_info, "w") as f:
                    f.write(json.dumps(self.user))
                yield message
                return
            
            logger.info(f"Image generation: {file_path} <- {prompt}")
            request = ImageRequest(filepath=file_path, prompt=prompt)
            async for message in generate_image(
                message=message,
                request=request 
            ):
                self.user["images"][self.filename] = { "status": message.status, "response": message.response}
                with open(user_info, "w") as f:
                    f.write(json.dumps(self.user))
                if message.status != "done":
                    yield message
            logger.info("Image generation done...")
            if not os.path.exists(file_path):
                logger.info(f"Generated image does not exist: {file_path}")
                logger.error(f"{message.status} {message.response}")
            else:
                if self.filename == "profile.png":
                    self.user["has_profile"] = True

            self.user["images"][self.filename] = { "status": message.status, "response": message.response}
            with open(user_info, "w") as f:
                f.write(json.dumps(self.user))

            #
            # Write out the completed user information
            #
            with open(user_info, "w") as f:
                f.write(json.dumps(self.user))

            # Image generated
            message.status = "done"
            if self.filename != "profile.png":
                message.response = f"/api/u/{self.username}/images/{self.filename}"
            else:
                message.response = f"/api/u/{self.username}/profile"

        except Exception as e:
            message.status = "error"
            logger.error(traceback.format_exc())
            logger.error(message.response)
            message.response = f"Error in image generation: {str(e)}"
            logger.error(message.response)
            self.user["images"][self.filename] = { "status": message.status, "response": message.response}
            with open(user_info, "w") as f:
                f.write(json.dumps(self.user))
            yield message
            return

        # Done processing, add message to conversation
        self.user["images"][self.filename] = { "status": message.status, "response": message.response}
        with open(user_info, "w") as f:
            f.write(json.dumps(self.user))
        self.context.processing = False
        # Return the final message
        yield message
        return

    async def call_llm(self, message: Message, system_prompt, prompt, temperature=0.7):
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")

        messages: List[LLMMessage] = [
            LLMMessage(role="system", content=system_prompt),
            LLMMessage(role="user", content=prompt),
        ]
        message.metadata.options = {
            "seed": 8911,
            "num_ctx": self.context_size,
            "temperature": temperature,  # Higher temperature to encourage tool usage
        }

        message.status = "streaming"
        yield message

        last_chunk_time = 0
        message.chunk = ""
        message.response = ""
        for response in self.llm.chat(
            model=self.model,
            messages=messages,
            options={
                **message.metadata.options,
            },
            stream=True,
        ):
            if not response:
                message.status = "error"
                message.response = "No response from LLM."
                yield message
                return

            message.status = "streaming"
            message.chunk += response.message.content
            message.response += response.message.content

            if not response.done:
                now = time.perf_counter()
                if now - last_chunk_time > 0.25:
                    yield message
                    last_chunk_time = now
                    message.chunk = ""

            if response.done:
                self.collect_metrics(response)
                message.metadata.eval_count += response.eval_count
                message.metadata.eval_duration += response.eval_duration
                message.metadata.prompt_eval_count += response.prompt_eval_count
                message.metadata.prompt_eval_duration += response.prompt_eval_duration
                self.context_tokens = response.prompt_eval_count + response.eval_count
                message.chunk = ""
                message.status = "done"
                yield message

    def extract_json_from_text(self, text: str) -> str:
        """Extract JSON string from text that may contain other content."""
        json_pattern = r"```json\s*([\s\S]*?)\s*```"
        match = re.search(json_pattern, text)
        if match:
            return match.group(1).strip()

        # Try to find JSON without the markdown code block
        json_pattern = r"({[\s\S]*})"
        match = re.search(json_pattern, text)
        if match:
            return match.group(1).strip()

        raise ValueError("No JSON found in the response")

    def extract_markdown_from_text(self, text: str) -> str:
        """Extract Markdown string from text that may contain other content."""
        markdown_pattern = r"```(md|markdown)\s*([\s\S]*?)\s*```"
        match = re.search(markdown_pattern, text)
        if match:
            return match.group(2).strip()

        raise ValueError("No Markdown found in the response")
    
# Register the base agent
agent_registry.register(ImageGenerator._agent_type, ImageGenerator)
