from __future__ import annotations
from pydantic import BaseModel, Field, model_validator  # type: ignore
from typing import (
    Literal,
    get_args,
    List,
    AsyncGenerator,
    TYPE_CHECKING,
    Optional,
    ClassVar,
    Any,
    TypeAlias,
    Dict,
    Tuple,
)
import json
import time
import inspect
import re
from abc import ABC
import asyncio
from datetime import datetime, UTC
from prometheus_client import Counter, Summary, CollectorRegistry  # type: ignore
import numpy as np  # type: ignore

from models import ( ApiActivityType, ChatMessageError, ChatMessageRagSearch, ChatMessageStatus, ChatMessageStreaming, LLMMessage, ChatQuery, ChatMessage, ChatOptions, ChatMessageUser, Tunables, ApiMessageType, ChatSenderType, ApiStatusType, ChatMessageMetaData, Candidate)
from logger import logger
import defines
from .registry import agent_registry
from metrics import Metrics
import model_cast
import backstory_traceback as traceback

from rag import ( ChromaDBGetResponse )

class Agent(BaseModel, ABC):
    """
    Base class for all agent types.
    This class defines the common attributes and methods for all agent types.
    """
    class Config:
        arbitrary_types_allowed = True  # Allow arbitrary types like RedisDatabase

    # Agent management with pydantic
    agent_type: Literal["base"] = "base"
    _agent_type: ClassVar[str] = agent_type  # Add this for registration

    user: Optional[Candidate] = None
    prometheus_collector: CollectorRegistry = Field(..., description="Prometheus collector for this agent, used to track metrics.", exclude=True)

    # Tunables (sets default for new Messages attached to this agent)
    tunables: Tunables = Field(default_factory=Tunables)
    metrics: Metrics = Field(
        None, description="Metrics collector for this agent, used to track performance and usage."
    )

    @model_validator(mode="after")
    def initialize_metrics(self) -> "Agent":
        if self.metrics is None:
            self.metrics = Metrics(prometheus_collector=self.prometheus_collector)
        return self
    
    # Agent properties
    system_prompt: str = ""
    context_tokens: int = 0

    # context_size is shared across all subclasses
    _context_size: ClassVar[int] = int(defines.max_context * 0.5)

    conversation: List[ChatMessage] = Field(
        default_factory=list,
        description="Conversation history for this agent, used to maintain context across messages."
    )

    @property
    def context_size(self) -> int:
        return Agent._context_size

    @context_size.setter
    def context_size(self, value: int):
        Agent._context_size = value

    async def get_last_item(self, generator):
        last_item = None
        async for item in generator:
            last_item = item
        return last_item

    def set_optimal_context_size(
        self, llm: Any, model: str, prompt: str, ctx_buffer=2048
    ) -> int:
        # Most models average 1.3-1.5 tokens per word
        word_count = len(prompt.split())
        tokens = int(word_count * 1.4)

        # Add buffer for safety
        total_ctx = tokens + ctx_buffer

        if total_ctx > self.context_size:
            logger.info(
                f"Increasing context size from {self.context_size} to {total_ctx}"
            )

        # Grow the context size if necessary
        self.context_size = max(self.context_size, total_ctx)
        # Use actual model maximum context size
        return self.context_size

    # Class and pydantic model management
    def __init_subclass__(cls, **kwargs) -> None:
        """Auto-register subclasses"""
        super().__init_subclass__(**kwargs)
        # Register this class if it has an agent_type
        if hasattr(cls, "agent_type") and cls.agent_type != Agent._agent_type:
            agent_registry.register(cls.agent_type, cls)

    def model_dump(self, *args, **kwargs) -> Any:
        # Ensure context is always excluded, even with exclude_unset=True
        kwargs.setdefault("exclude", set())
        if isinstance(kwargs["exclude"], set):
            kwargs["exclude"].add("context")
        elif isinstance(kwargs["exclude"], dict):
            kwargs["exclude"]["context"] = True
        return super().model_dump(*args, **kwargs)

    @classmethod
    def valid_agent_types(cls) -> set[str]:
        """Return the set of valid agent_type values."""
        return set(get_args(cls.__annotations__["agent_type"]))

    # Agent methods
    def get_agent_type(self):
        return self._agent_type

    # async def process_tool_calls(
    #     self,
    #     llm: Any,
    #     model: str,
    #     message: ChatMessage,
    #     tool_message: Any, # llama response
    #     messages: List[LLMMessage],
    # ) -> AsyncGenerator[ChatMessage, None]:
    #     logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")

    #     self.metrics.tool_count.labels(agent=self.agent_type).inc()
    #     with self.metrics.tool_duration.labels(agent=self.agent_type).time():

    #         if not self.context:
    #             raise ValueError("Context is not set for this agent.")
    #         if not message.metadata.tools:
    #             raise ValueError("tools field not initialized")

    #         tool_metadata = message.metadata.tools
    #         tool_metadata["tool_calls"] = []

    #         message.status = "tooling"

    #         for i, tool_call in enumerate(tool_message.tool_calls):
    #             arguments = tool_call.function.arguments
    #             tool = tool_call.function.name

    #             # Yield status update before processing each tool
    #             message.content = (
    #                 f"Processing tool {i+1}/{len(tool_message.tool_calls)}: {tool}..."
    #             )
    #             yield message
    #             logger.info(f"LLM - {message.content}")

    #             # Process the tool based on its type
    #             match tool:
    #                 case "TickerValue":
    #                     ticker = arguments.get("ticker")
    #                     if not ticker:
    #                         ret = None
    #                     else:
    #                         ret = TickerValue(ticker)

    #                 case "AnalyzeSite":
    #                     url = arguments.get("url")
    #                     question = arguments.get(
    #                         "question", "what is the summary of this content?"
    #                     )

    #                     # Additional status update for long-running operations
    #                     message.content = (
    #                         f"Retrieving and summarizing content from {url}..."
    #                     )
    #                     yield message
    #                     ret = await AnalyzeSite(
    #                         llm=llm, model=model, url=url, question=question
    #                     )

    #                 case "GenerateImage":
    #                     prompt = arguments.get("prompt", None)
    #                     if not prompt:
    #                         logger.info("No prompt supplied to GenerateImage")
    #                         ret = { "error": "No prompt supplied to GenerateImage" }

    #                     # Additional status update for long-running operations
    #                     message.content = (
    #                         f"Generating image for {prompt}..."
    #                     )
    #                     yield message
    #                     ret = await GenerateImage(
    #                         llm=llm, model=model, prompt=prompt
    #                     )
    #                     logger.info("GenerateImage returning", ret)

    #                 case "DateTime":
    #                     tz = arguments.get("timezone")
    #                     ret = DateTime(tz)

    #                 case "WeatherForecast":
    #                     city = arguments.get("city")
    #                     state = arguments.get("state")

    #                     message.content = (
    #                         f"Fetching weather data for {city}, {state}..."
    #                     )
    #                     yield message
    #                     ret = WeatherForecast(city, state)

    #                 case _:
    #                     logger.error(f"Requested tool {tool} does not exist")
    #                     ret = None

    #             # Build response for this tool
    #             tool_response = {
    #                 "role": "tool",
    #                 "content": json.dumps(ret),
    #                 "name": tool_call.function.name,
    #             }

    #             tool_metadata["tool_calls"].append(tool_response)

    #             if len(tool_metadata["tool_calls"]) == 0:
    #                 message.status = "done"
    #                 yield message
    #                 return

    #             message_dict = LLMMessage(
    #                 role=tool_message.get("role", "assistant"),
    #                 content=tool_message.get("content", ""),
    #                 tool_calls=[
    #                     {
    #                         "function": {
    #                             "name": tc["function"]["name"],
    #                             "arguments": tc["function"]["arguments"],
    #                         }
    #                     }
    #                     for tc in tool_message.tool_calls
    #                 ],
    #             )

    #             messages.append(message_dict)
    #             messages.extend(tool_metadata["tool_calls"])

    #         message.status = "thinking"
    #         message.content = "Incorporating tool results into response..."
    #         yield message

    #         # Decrease creativity when processing tool call requests
    #         message.content = ""
    #         start_time = time.perf_counter()
    #         for response in llm.chat(
    #             model=model,
    #             messages=messages,
    #             options={
    #                 **message.metadata.options,
    #             },
    #             stream=True,
    #         ):
    #             # logger.info(f"LLM::Tools: {'done' if response.finish_reason else 'processing'} - {response}")
    #             message.status = "streaming"
    #             message.chunk = response.content
    #             message.content += message.chunk

    #             if not response.finish_reason:
    #                 yield message

    #             if response.finish_reason:
    #                 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.status = "done"
    #                 yield message

    #         end_time = time.perf_counter()
    #         message.metadata.timers["llm_with_tools"] = end_time - start_time
    #         return

    def collect_metrics(self, response):
        self.metrics.tokens_prompt.labels(agent=self.agent_type).inc(
            response.usage.prompt_eval_count
        )
        self.metrics.tokens_eval.labels(agent=self.agent_type).inc(response.usage.eval_count)

    def get_rag_context(self, rag_message: ChatMessageRagSearch) -> str:
        """
        Extracts the RAG context from the rag_message.
        """
        if not rag_message.content:
            return ""
        
        context = []
        for chroma_results in rag_message.content:
            for index, metadata in enumerate(chroma_results.metadatas):
                content = "\n".join([
                    line.strip()
                    for line in chroma_results.documents[index].split("\n")
                    if line
                ]).strip()
                context.append(f"""
Source: {metadata.get("doc_type", "unknown")}: {metadata.get("path", "")}
Document reference: {chroma_results.ids[index]}
Content: {content}
""")
        return "\n".join(context)

    async def generate_rag_results(
        self,
        session_id: str,
        prompt: str,
        top_k: int=defines.default_rag_top_k, 
        threshold: float=defines.default_rag_threshold,
    ) -> AsyncGenerator[ChatMessageRagSearch | ChatMessageError | ChatMessageStatus, None]:
        """
        Generate RAG results for the given query.

        Args:
            query: The query string to generate RAG results for.

        Returns:
            A list of dictionaries containing the RAG results.
        """
        if not self.user:
            error_message = ChatMessageError(
                session_id=session_id,
                content="No user set for RAG generation."
            )
            yield error_message
            return
        
        results : List[ChromaDBGetResponse] = []
        entries: int = 0
        user: Candidate = self.user
        for rag in user.rags:
            if not rag.enabled:
                continue

            status_message = ChatMessageStatus(
                session_id=session_id,
                activity=ApiActivityType.SEARCHING,
                content = f"Searching RAG context {rag.name}..."
            )
            yield status_message

            try:
                chroma_results = await user.file_watcher.find_similar(
                    query=prompt, top_k=top_k, threshold=threshold
                )
                if not chroma_results:
                    continue
                query_embedding = np.array(chroma_results["query_embedding"]).flatten()

                umap_2d = user.file_watcher.umap_model_2d.transform([query_embedding])[0]
                umap_3d = user.file_watcher.umap_model_3d.transform([query_embedding])[0]

                rag_metadata = ChromaDBGetResponse( 
                    name=rag.name,
                    query_embedding=query_embedding.tolist(),
                    ids=chroma_results.get("ids", []),
                    embeddings=chroma_results.get("embeddings", []),
                    documents=chroma_results.get("documents", []),
                    metadatas=chroma_results.get("metadatas", []),
                    umap_embedding_2d=umap_2d.tolist(),
                    umap_embedding_3d=umap_3d.tolist(),
                )
                results.append(rag_metadata)
            except Exception as e:
                continue_message = ChatMessageStatus(
                    session_id=session_id,
                    activity=ApiActivityType.SEARCHING,
                    content=f"Error searching RAG context {rag.name}: {str(e)}"
                )
                yield continue_message

        final_message = ChatMessageRagSearch(
            session_id=session_id,
            content=results,
            status=ApiStatusType.DONE,
        )
        yield final_message
        return
        
    async def llm_one_shot(
            self, 
            llm: Any, model: str, 
            session_id: str, prompt: str, system_prompt: str, 
            tunables: Optional[Tunables] = None,
             temperature=0.7) -> AsyncGenerator[ChatMessageStatus | ChatMessageError | ChatMessageStreaming | ChatMessage, None]:

        self.set_optimal_context_size(
            llm=llm, model=model, prompt=prompt
        )

        options = ChatOptions(
            seed=8911,
            num_ctx=self.context_size,
            temperature=temperature,
        )

        messages: List[LLMMessage] = [
            LLMMessage(role="system", content=system_prompt),
            LLMMessage(role="user", content=prompt),
        ]

        status_message = ChatMessageStatus(
            session_id=session_id,
            activity=ApiActivityType.GENERATING,
            content=f"Generating response..."
        )
        yield status_message

        logger.info(f"Message options: {options.model_dump(exclude_unset=True)}")
        response = None
        content = ""
        async for response in llm.chat_stream(
            model=model,
            messages=messages,
            options={
                **options.model_dump(exclude_unset=True),
            },
            stream=True,
        ):
            if not response:
                error_message = ChatMessageError(
                    session_id=session_id,
                    content="No response from LLM."
                )
                yield error_message
                return

            content += response.content

            if not response.finish_reason:
                streaming_message = ChatMessageStreaming(
                    session_id=session_id,
                    content=response.content,
                    status=ApiStatusType.STREAMING,
                )
                yield streaming_message

        if not response:
            error_message = ChatMessageError(
                session_id=session_id,
                content="No response from LLM."
            )
            yield error_message
            return
        
        self.collect_metrics(response)
        self.context_tokens = (
            response.usage.prompt_eval_count + response.usage.eval_count
        )

        chat_message = ChatMessage(
            session_id=session_id,
            tunables=tunables,
            status=ApiStatusType.DONE,
            content=content,
            metadata = ChatMessageMetaData(
                options=options,
                eval_count=response.usage.eval_count,
                eval_duration=response.usage.eval_duration,
                prompt_eval_count=response.usage.prompt_eval_count,
                prompt_eval_duration=response.usage.prompt_eval_duration,

            )
        )
        yield chat_message
        return

    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:
            error_message = ChatMessageError(
                session_id=session_id,
                content="No user set for chat generation."
            )
            yield error_message
            return
        
        user_message = ChatMessageUser(
            session_id=session_id,
            content=prompt,
        )
        user = self.user

        self.metrics.generate_count.labels(agent=self.agent_type).inc()
        with self.metrics.generate_duration.labels(agent=self.agent_type).time():
            context = None
            if self.user:
                rag_message = None
                async for rag_message in self.generate_rag_results(session_id=session_id, prompt=prompt):
                    if rag_message.status == ApiStatusType.ERROR:
                        yield rag_message
                        return
                    # Only yield messages that are in a streaming state
                    if rag_message.status == ApiStatusType.STATUS:
                        yield rag_message

                if not isinstance(rag_message, ChatMessageRagSearch):
                    raise ValueError(
                        f"Expected ChatMessageRagSearch, got {type(rag_message)}"
                    )
                
                context = self.get_rag_context(rag_message)

            # Create a pruned down message list based purely on the prompt and responses,
            # discarding the full preamble generated by prepare_message
            messages: List[LLMMessage] = [
                LLMMessage(role="system", content=self.system_prompt)
            ]
            # Add the conversation history to the messages
            messages.extend([
                LLMMessage(role="user" if isinstance(m, ChatMessageUser) else "assistant", content=m.content)
                for m in self.conversation
            ])
            # Add the RAG context to the messages if available
            if context:
                messages.append(
                    LLMMessage(
                        role="user",
                        content=f"<|context|>\nThe following is context information about {self.user.full_name}:\n{context}\n</|context|>\n\nPrompt to respond to:\n{prompt}\n"
                    )
                )
            else:
                # Only the actual user query is provided with the full context message
                messages.append(
                    LLMMessage(role="user", content=prompt)
                )

            llm_history = messages

            # use_tools = message.tunables.enable_tools and len(self.context.tools) > 0
            # message.metadata.tools = {
            #     "available": llm_tools(self.context.tools),
            #     "used": False,
            # }
            # tool_metadata = message.metadata.tools

            # if use_tools:
            #     message.status = "thinking"
            #     message.content = f"Performing tool analysis step 1/2..."
            #     yield message

            #     logger.info("Checking for LLM tool usage")
            #     start_time = time.perf_counter()
            #     # Tools are enabled and available, so query the LLM with a short context of messages
            #     # in case the LLM did something like ask "Do you want me to run the tool?" and the
            #     # user said "Yes" -- need to keep the context in the thread.
            #     tool_metadata["messages"] = (
            #         [{"role": "system", "content": self.system_prompt}] + messages[-6:]
            #         if len(messages) >= 7
            #         else messages
            #     )

            #     response = llm.chat(
            #         model=model,
            #         messages=tool_metadata["messages"],
            #         tools=tool_metadata["available"],
            #         options={
            #             **message.metadata.options,
            #         },
            #         stream=False,  # No need to stream the probe
            #     )
            #     self.collect_metrics(response)

            #     end_time = time.perf_counter()
            #     message.metadata.timers["tool_check"] = end_time - start_time
            #     if not response.tool_calls:
            #         logger.info("LLM indicates tools will not be used")
            #         # The LLM will not use tools, so disable use_tools so we can stream the full response
            #         use_tools = False
            #     else:
            #         tool_metadata["attempted"] = response.tool_calls

            # if use_tools:
            #     logger.info("LLM indicates tools will be used")

            #     # Tools are enabled and available and the LLM indicated it will use them
            #     message.content = (
            #         f"Performing tool analysis step 2/2 (tool use suspected)..."
            #     )
            #     yield message

            #     logger.info(f"Performing LLM call with tools")
            #     start_time = time.perf_counter()
            #     response = llm.chat(
            #         model=model,
            #         messages=tool_metadata["messages"],  # messages,
            #         tools=tool_metadata["available"],
            #         options={
            #             **message.metadata.options,
            #         },
            #         stream=False,
            #     )
            #     self.collect_metrics(response)

            #     end_time = time.perf_counter()
            #     message.metadata.timers["non_streaming"] = end_time - start_time

            #     if not response:
            #         message.status = "error"
            #         message.content = "No response from LLM."
            #         yield message
            #         return

            #     if response.tool_calls:
            #         tool_metadata["used"] = response.tool_calls
            #         # Process all yielded items from the handler
            #         start_time = time.perf_counter()
            #         async for message in self.process_tool_calls(
            #             llm=llm,
            #             model=model,
            #             message=message,
            #             tool_message=response,
            #             messages=messages,
            #         ):
            #             if message.status == "error":
            #                 yield message
            #                 return
            #             yield message
            #         end_time = time.perf_counter()
            #         message.metadata.timers["process_tool_calls"] = end_time - start_time
            #         message.status = "done"
            #         return

            #     logger.info("LLM indicated tools will be used, and then they weren't")
            #     message.content = response.content
            #     message.status = "done"
            #     yield message
            #     return

            # not use_tools
            status_message = ChatMessageStatus(
                session_id=session_id,
                activity=ApiActivityType.GENERATING,
                content=f"Generating response..."
            )
            yield status_message

            # Set the response for streaming
            self.set_optimal_context_size(
                llm, model, prompt=prompt
            )

            options = ChatOptions(
                seed=8911,
                num_ctx=self.context_size,
                temperature=temperature,
            )
            logger.info(f"Message options: {options.model_dump(exclude_unset=True)}")
            content = ""
            start_time = time.perf_counter()
            response = None
            async for response in llm.chat_stream(
                model=model,
                messages=messages,
                options={
                    **options.model_dump(exclude_unset=True),
                },
                stream=True,
            ):
                if not response:
                    error_message = ChatMessageError(
                        session_id=session_id,
                        content="No response from LLM."
                    )
                    yield error_message
                    return

                content += response.content

                if not response.finish_reason:
                    streaming_message = ChatMessageStreaming(
                        session_id=session_id,
                        content=response.content,
                    )
                    yield streaming_message

            if not response:
                error_message = ChatMessageError(
                    session_id=session_id,
                    content="No response from LLM."
                )
                yield error_message
                return
            
            self.collect_metrics(response)
            self.context_tokens = (
                response.usage.prompt_eval_count + response.usage.eval_count
            )
            end_time = time.perf_counter()

            chat_message = ChatMessage(
                session_id=session_id,
                tunables=tunables,
                status=ApiStatusType.DONE,
                content=content,
                metadata = ChatMessageMetaData(
                    options=options,
                    eval_count=response.usage.eval_count,
                    eval_duration=response.usage.eval_duration,
                    prompt_eval_count=response.usage.prompt_eval_count,
                    prompt_eval_duration=response.usage.prompt_eval_duration,
                    timers={
                        "llm_streamed": end_time - start_time,
                        "llm_with_tools": 0,  # Placeholder for tool processing time
                    },
                )
            )


        # Add the user and chat messages to the conversation
        self.conversation.append(user_message)
        self.conversation.append(chat_message)
        yield chat_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}")

    #     self.metrics.process_count.labels(agent=self.agent_type).inc()
    #     with self.metrics.process_duration.labels(agent=self.agent_type).time():

    #         if not self.context:
    #             raise ValueError("Context is not set for this agent.")

    #         logger.info(
    #             "TODO: Implement delay queing; busy for same agent, otherwise return queue size and estimated wait time"
    #         )
    #         spinner: List[str] = ["\\", "|", "/", "-"]
    #         tick: int = 0
    #         while self.context.processing:
    #             message.status = "waiting"
    #             message.content = (
    #                 f"Busy processing another request. Please wait. {spinner[tick]}"
    #             )
    #             tick = (tick + 1) % len(spinner)
    #             yield message
    #             await asyncio.sleep(1)  # Allow the event loop to process the write

    #         self.context.processing = True

    #         message.system_prompt = (
    #             f"<|system|>\n{self.system_prompt.strip()}\n</|system|>"
    #         )
    #         message.context_prompt = ""
    #         for p in message.preamble.keys():
    #             message.context_prompt += (
    #                 f"\n<|{p}|>\n{message.preamble[p].strip()}\n</|{p}>\n\n"
    #             )
    #         message.context_prompt += f"{message.prompt}"

    #         # Estimate token length of new messages
    #         message.content = f"Optimizing context..."
    #         message.status = "thinking"
    #         yield message

    #         message.context_size = self.set_optimal_context_size(
    #             llm, model, prompt=message.context_prompt
    #         )

    #         message.content = f"Processing {'RAG augmented ' if message.metadata.rag else ''}query..."
    #         message.status = "thinking"
    #         yield message

    #         async for message in self.generate_llm_response(
    #             llm=llm, model=model, message=message
    #         ):
    #             # logger.info(f"LLM: {message.status} - {f'...{message.content[-20:]}' if len(message.content) > 20 else message.content}")
    #             if message.status == "error":
    #                 yield message
    #                 self.context.processing = False
    #                 return
    #             yield message

    #         # Done processing, add message to conversation
    #         message.status = "done"
    #         self.conversation.add(message)
    #         self.context.processing = False

    #     return

    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")

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