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

    preamble_types = [f"<|{p}|>" for p in message.preamble.keys()]
    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
registry.register(Resume._agent_type, Resume)