from __future__ import annotations
from pydantic import model_validator  
from typing import (
    Literal,
    ClassVar,
    Optional,
    Any,
    AsyncGenerator,
    List,
)  # NOTE: You must import Optional for late binding to work
from datetime import datetime
import inspect

from .base import Agent, agent_registry
from ..conversation import Conversation
from ..message import Message
from ..setup_logging import setup_logging

logger = setup_logging()

system_fact_check = f"""
Launched on {datetime.now().isoformat()}.

You are a professional resume fact checker. Your task is answer any questions about items identified in the <|discrepancies|>.
The <|discrepancies|> indicate inaccuracies or unsupported claims in the <|generated-resume|> based on content from the <|resume|> and <|context|>.

When answering queries, follow these steps:
- If there is information in the <|context|> or <|resume|> sections to enhance the answer, incorporate it seamlessly and refer to it using natural language instead of mentioning '<|context|>' (etc.) or quoting it directly.
- Avoid phrases like 'According to the <|context|>' or similar references to the <|context|>, <|generated-resume|>, or <|resume|> tags.
""".strip()


class FactCheck(Agent):
    agent_type: Literal["fact_check"] = "fact_check"  
    _agent_type: ClassVar[str] = agent_type  # Add this for registration

    system_prompt: str = system_fact_check
    facts: str

    @model_validator(mode="after")
    def validate_facts(self):
        if not self.facts.strip():
            raise ValueError("Facts cannot be empty")
        return self

    async def prepare_message(self, 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.")

        resume_agent = self.context.get_agent("resume")
        if not resume_agent:
            raise ValueError("resume agent does not exist")

        message.tunables.enable_tools = False

        async for message in super().prepare_message(message):
            if message.status != "done":
                yield message

        message.preamble["generated-resume"] = resume_agent.resume
        message.preamble["discrepancies"] = self.facts

        excluded = {"job_description"}
        preamble_types = [
            f"<|{p}|>" for p in message.preamble.keys() if p not in excluded
        ]
        preamble_types_AND = " and ".join(preamble_types)
        preamble_types_OR = " or ".join(preamble_types)
        message.preamble[
            "rules"
        ] = f"""\
- Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
- If there is no information in these sections, answer based on your knowledge, or use any available tools.
- Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
"""
        message.preamble["question"] = "Respond to:"

        yield message
        return


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