from __future__ import annotations
from pydantic import model_validator  # type: ignore
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 ..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 to identify any inaccuracies in the <|generated-resume|> based on the individual's <|context|> and <|resume|>.

If there are inaccuracies, list them in a bullet point format.

When answering queries, follow these steps:
- Analyze the <|generated-resume|> to identify any discrepancies or inaccuracies which are not supported by the <|context|> and <|resume|>.
- 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.

Do not generate a revised resume.
""".strip()

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

You are a hiring and job placing specialist. Your task is to answers about a resume and work history as it relates to a potential job.

When answering queries, follow these steps:
- Analyze the <|job_description|> and <|generated-resume|> to provide insights for the asked question.
- If any financial information is requested, be sure to account for inflation.
- If there is information in the <|context|>, <|job_description|>, <|generated-resume|>, or <|resume|> sections to enhance the answer, incorporate it seamlessly and refer to it using natural language instead of mentioning '<|job_description|>' (etc.) or quoting it directly.
- Avoid phrases like 'According to the <|context|>' or similar references to the <|context|>, <|job_description|>, <|resume|>, or <|context|> tags.
""".strip()


class Resume(Agent):
    agent_type: Literal["resume"] = "resume"  # type: ignore
    _agent_type: ClassVar[str] = agent_type  # Add this for registration

    system_prompt: str = system_fact_check
    resume: str

    @model_validator(mode="after")
    def validate_resume(self):
        if not self.resume.strip():
            raise ValueError("Resume content 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.")

        # Generating fact check or resume should not use any tools
        message.tunables.enable_tools = False

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

        message.preamble["generated-resume"] = self.resume
        job_description_agent = self.context.get_agent("job_description")
        if not job_description_agent:
            raise ValueError("job_description agent does not exist")

        message.preamble["job_description"] = job_description_agent.job_description

        excluded = {}
        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}.
"""
        fact_check_agent = self.context.get_agent(agent_type="fact_check")
        if fact_check_agent:
            message.preamble["question"] = "Respond to:"
        else:
            message.preamble["question"] = (
                f"Fact check the <|generated-resume|> based on the <|resume|>{' and <|context|>' if 'context' in message.preamble else ''}."
            )

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

        async for message in super().process_message(llm, model, message):
            if message.status != "done":
                yield message

        fact_check_agent = self.context.get_agent(agent_type="fact_check")
        if not fact_check_agent:
            # Switch agent from "Fact Check Generated Resume" mode
            # to "Answer Questions about Generated Resume"
            self.system_prompt = system_resume

            # Instantiate the "resume" agent, and seed (or reset) its conversation
            # with this message.
            fact_check_agent = self.context.get_or_create_agent(
                agent_type="fact_check", facts=message.response
            )
            first_fact_check_message = message.copy()
            first_fact_check_message.prompt = "Fact check the generated resume."
            fact_check_agent.conversation.add(first_fact_check_message)
            message.response = "Resume fact checked."

        # Return the final message
        yield message
        return


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