from __future__ import annotations
from pydantic import model_validator, Field  # type: ignore
from typing import (
    Dict,
    Literal,
    ClassVar,
    Any,
    AsyncGenerator,
    List,
    Optional
#    override
)  # NOTE: You must import Optional for late binding to work
import inspect
import re
import json
import traceback
import asyncio
import time
import asyncio
import numpy as np  # type: ignore

from . base import Agent, agent_registry, LLMMessage
from .. message import Message
from .. rag import ChromaDBGetResponse
from .. setup_logging import setup_logging

logger = setup_logging()


system_generate_resume = """
""".strip()

system_user_qualifications = """
Answer questions about the job description.
"""


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

    system_prompt: str = system_generate_resume
    job_description: str

    llm: Any = Field(default=None, exclude=True)
    model: str = Field(default=None, exclude=True)

    @model_validator(mode="after")
    def validate_job_description(self):
        if not self.job_description.strip():
            raise ValueError("Job description cannot be empty")
        return self

    #@override
    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
        # Always add the job description, user resume, and question
        message.preamble["job_description"] = self.job_description
        message.preamble["resume"] = self.context.user_resume
        yield message
        return

    #@override
    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.")

        # Generating a resume should not use any tools
        message.tunables.enable_tools = False

        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.response = (
                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
        job_description = message.preamble["job_description"]
        resume = message.preamble["resume"]

        original_message = message.model_copy()
        original_message.prompt = job_description
        original_message.response = ""
        self.conversation.add(original_message)

        self.llm = llm
        self.model = model
        self.metrics.generate_count.labels(agent=self.agent_type).inc()
        with self.metrics.generate_duration.labels(agent=self.agent_type).time():

            try:
                async for message in self.generate_resume(
                    message=message, job_description=job_description, resume=resume
                ):
                    if message.status != "done":
                        yield message
                message.status = "done"
                yield message

            except Exception as e:
                message.status = "error"
                logger.error(message.response)
                message.response = f"Error in resume generation process: {str(e)}"
                logger.error(message.response)
                logger.error(traceback.format_exc())
                yield message
                return

        # Done processing, add message to conversation
        self.context.processing = False

        resume_generation = message.metadata.resume_generation
        if not resume_generation:
            message.response = (
                "Generation did not generate metadata necessary for processing."
            )
            message.status = "error"
            yield message
            return

        generated_resume = resume_generation.get(
            "resume", "Resume generation resulted in an empty resume."
        )

        facts = resume_generation.get("match_stats", {"results": "No skills matched"})

        # Add the "Job requirements" message

        self.system_prompt = system_user_qualifications

        resume_agent = self.context.get_agent(agent_type="resume")
        if not resume_agent:
            resume_agent = self.context.get_or_create_agent(
                agent_type="resume", resume=generated_resume
            )
        else:
            resume_agent.conversation.reset()
        resume_message = message.model_copy()
        resume_message.prompt = ""
        resume_message.response = generated_resume
        resume_message.status = "done"
        resume_agent.conversation.add(resume_message)

        message.actions.append("resume_generated")

        fact_check_agent = self.context.get_agent(agent_type="fact_check")
        if not fact_check_agent:
            fact_check_agent = self.context.get_or_create_agent(
                agent_type="fact_check", facts=json.dumps(facts, indent=2)
            )
        else:
            fact_check_agent.conversation.reset()

        fact_check_message = message.model_copy()
        fact_check_message.prompt = ""
        clean_json = re.sub(r"```", "'''", json.dumps(facts, indent=2))
        fact_check_message.response = f"```json\n\n{clean_json}\n```\n"
        fact_check_message.status = "done"
        fact_check_agent.conversation.add(fact_check_message)

        message.actions.append("facts_checked")

        # Return the final message
        message.response = "Resume generated."
        yield message
        return

    # BEGIN

    # Helper functions
    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 validate_job_requirements(self, job_requirements: Dict) -> None:
        """Validate the structure of job requirements."""
        required_keys = ["job_requirements"]

        if not all(key in job_requirements for key in required_keys):
            missing = [key for key in required_keys if key not in job_requirements]
            raise ValueError(f"Missing required keys in job requirements: {missing}")

        # Additional validation can be added here

    def generate_resume_from_skill_assessments(
        self,
        candidate_info,
        skill_assessment_results,
        original_resume,
    ):
        """
        Generate a professional resume based on skill assessment results

        Parameters:
        - candidate_name (str): The candidate's full name
        - candidate_contact_info (dict): Contact details like email, phone, location
        - skill_assessment_results (list): List of individual skill assessment results from LLM queries
        - original_resume (str): Original resume text for reference

        Returns:
        - str: System prompt for generating a professional resume
        """

        # Extract and organize skill assessment data
        skills_by_strength = {"STRONG": [], "MODERATE": [], "WEAK": [], "NONE": []}

        experience_evidence = {}

        # Process each skill assessment
        for result in skill_assessment_results:
            # Parse the result if it's a string
            if isinstance(result, str):
                try:
                    result = json.loads(result)
                    assessment = result.get("skill_assessment", {})
                except:
                    continue
            else:
                assessment = result.get("skill_assessment", {})

            skill = assessment.get("skill", "")
            strength = assessment.get("evidence_strength", "NONE")

            # Add to appropriate strength category
            if skill and strength in skills_by_strength:
                skills_by_strength[strength].append(skill)

            # Collect experience evidence
            for evidence in assessment.get("evidence_details", []):
                source = evidence.get("source", "")
                if source:
                    if source not in experience_evidence:
                        experience_evidence[source] = []

                    experience_evidence[source].append(
                        {
                            "skill": skill,
                            "quote": evidence.get("quote", ""),
                            "context": evidence.get("context", ""),
                        }
                    )

        # Format contact info
        contact_info_str = ""
        contact_items = []
        for key, value in candidate_info.get("contact_info", {}).items():
            if value:
                contact_items.append(f"{key}: {value}")
        contact_info_str = "\n".join(contact_items)

        # Build the system prompt
        system_prompt = f"""You are a professional resume writer with expertise in highlighting candidate strengths and experiences. 
Create a polished, concise, and ATS-friendly resume for the candidate based on the assessment data provided.

## CANDIDATE INFORMATION:
Name: {candidate_info["full_name"]}
{contact_info_str}

## SKILL ASSESSMENT RESULTS:
"""

        if len(skills_by_strength["STRONG"]):
            system_prompt += f"""\

### Strong Skills (prominent in resume):
{", ".join(skills_by_strength["STRONG"])}
"""

        if len(skills_by_strength["MODERATE"]):
            system_prompt += f"""\

### Moderate Skills (demonstrated in resume):
{", ".join(skills_by_strength["MODERATE"])}
"""

        if len(skills_by_strength["WEAK"]):
            system_prompt += f"""\

### Weaker Skills (mentioned or implied):
{", ".join(skills_by_strength["WEAK"])}
"""
            system_prompt += """\
       
## EXPERIENCE EVIDENCE:
"""

        # Add experience evidence by source/position
        for source, evidences in experience_evidence.items():
            system_prompt += f"\n### {source}:\n"
            for evidence in evidences:
                system_prompt += f"- {evidence['skill']}: {evidence['context']}\n"

        # Add instructions for the resume creation
        system_prompt += """\
    
## INSTRUCTIONS:

1. Create a professional resume that emphasizes the candidate's strongest skills and most relevant experiences.
2. Format the resume in a clean, concise, and modern style that will pass ATS systems.
3. Include these sections:
  - Professional Summary (highlight strongest skills and experience level)
  - Skills (organized by strength)
  - Professional Experience (focus on achievements and evidence of the skill)
  - Education
  - Additional sections as appropriate
4. Use action verbs and quantifiable achievements where possible.
5. Maintain a professional tone throughout.
6. Be concise and impactful - the resume should be 1-2 pages MAXIMUM.
7. Ensure all information is accurate to the original resume - do not embellish or fabricate experiences.

## OUTPUT FORMAT:
Provide the resume in clean markdown format, ready for the candidate to use.

"""
# ## REFERENCE (Original Resume):
# """

#         # Add a truncated version of the original resume for reference if it's too long
#         max_resume_length = 25000  # Characters
#         if len(original_resume) > max_resume_length:
#             system_prompt += (
#                 original_resume[:max_resume_length]
#                 + "...\n[Original resume truncated due to length]"
#             )
#         else:
#             system_prompt += original_resume

        prompt = "Create a tailored professional resume that highlights candidate's skills and experience most relevant to the job requirements. Format it in clean, ATS-friendly markdown. Provide ONLY the resume with no commentary before or after."
        return system_prompt, prompt

    async def process_and_generate_resume(
        self,
        message: Message,
        job_requirements,
        skill_assessment_results,
        candidate_info,
        metadata: Dict[str, Any],
    ) -> AsyncGenerator[Message, None]:
        """
        Process skill assessment results and generate an optimized resume

        Parameters:
        - job_requirements (dict): The original job requirements
        - skill_assessment_results (dict): Results from skill assessments
        - candidate_info (dict): Basic candidate information

        Returns:
        - dict: Contains the generated resume and match analysis
        """

        # Extract candidate information
        original_resume = candidate_info.get("original_resume", "")

        # Generate resume prompt
        system_prompt, prompt = self.generate_resume_from_skill_assessments(
            candidate_info,
            skill_assessment_results.values(),
            original_resume,
        )
        metadata["system_prompt"] = system_prompt
        metadata["prompt"] = prompt

        message.status = "thinking"
        message.response = "Generating tailored resume..."
        yield message

        # This would be replaced with actual LLM query
        async for message in self.call_llm(
            message=message, system_prompt=system_prompt, prompt=prompt
        ):
            if message.status != "done":
                yield message
        if message.status == "error":
            return

        metadata["resume"] = message.response
        metadata["match_stats"] = self.calculate_match_statistics(
            job_requirements, skill_assessment_results
        )
        # Calculate skill match statistics
        yield message
        return

    def calculate_match_statistics(self, job_requirements, skill_assessment_results) -> dict[str, dict[str, Any]]:
        """
        Calculate statistics about how well the candidate matches job requirements

        Parameters:
        - job_requirements (dict): The job requirements dictionary
        - skill_assessment_results (dict): Results from skill assessments

        Returns:
        - dict: Match statistics
        """

        # Initialize counters
        required_skills_total = 0
        required_skills_matched = 0
        preferred_skills_total = 0
        preferred_skills_matched = 0

        # Count required technical skills
        tech_required = job_requirements.get("technical_skills", {}).get("required", [])
        required_skills_total += len(tech_required)

        # Count preferred technical skills
        tech_preferred = job_requirements.get("technical_skills", {}).get(
            "preferred", []
        )
        preferred_skills_total += len(tech_preferred)

        # Count required experience
        exp_required = job_requirements.get("experience_requirements", {}).get(
            "required", []
        )
        required_skills_total += len(exp_required)

        # Count preferred experience
        exp_preferred = job_requirements.get("experience_requirements", {}).get(
            "preferred", []
        )
        preferred_skills_total += len(exp_preferred)

        # Education requirements count toward required
        edu_required = job_requirements.get("education_requirements", [])
        required_skills_total += len(edu_required)

        # Soft skills count toward preferred
        soft_skills = job_requirements.get("soft_skills", [])
        preferred_skills_total += len(soft_skills)

        # Industry knowledge counts toward preferred
        industry_knowledge = job_requirements.get("industry_knowledge", [])
        preferred_skills_total += len(industry_knowledge)

        # Check matches in assessment results
        for key, result in skill_assessment_results.items():
            # Parse if string
            if isinstance(result, str):
                try:
                    result = json.loads(result)
                except:
                    continue

            assessment = result.get("skill_assessment", {})
            evidence_found = assessment.get("evidence_found", False)
            evidence_strength = assessment.get("evidence_strength", "NONE")

            # Consider STRONG and MODERATE evidence as matches
            is_match = evidence_found and evidence_strength in ["STRONG", "MODERATE"]

            if not is_match:
                continue

            # Determine which category this skill belongs to
            if key.startswith("technical_required_"):
                required_skills_matched += 1
            elif key.startswith("technical_preferred_"):
                preferred_skills_matched += 1
            elif key.startswith("experience_required_"):
                required_skills_matched += 1
            elif key.startswith("experience_preferred_"):
                preferred_skills_matched += 1
            elif key.startswith("education_"):
                required_skills_matched += 1
            elif key.startswith("soft_skill_"):
                preferred_skills_matched += 1
            elif key.startswith("industry_knowledge_"):
                preferred_skills_matched += 1

        # Calculate percentages
        required_match_percent = (
            (required_skills_matched / required_skills_total * 100)
            if required_skills_total > 0
            else 0
        )
        preferred_match_percent = (
            (preferred_skills_matched / preferred_skills_total * 100)
            if preferred_skills_total > 0
            else 0
        )
        overall_total = required_skills_total + preferred_skills_total
        overall_matched = required_skills_matched + preferred_skills_matched
        overall_match_percent = (
            (overall_matched / overall_total * 100) if overall_total > 0 else 0
        )

        return {
            "required_skills": {
                "total": required_skills_total,
                "matched": required_skills_matched,
                "percentage": round(required_match_percent, 1),
            },
            "preferred_skills": {
                "total": preferred_skills_total,
                "matched": preferred_skills_matched,
                "percentage": round(preferred_match_percent, 1),
            },
            "overall_match": {
                "total": overall_total,
                "matched": overall_matched,
                "percentage": round(overall_match_percent, 1),
            },
        }

    def generate_skill_assessment_prompt(self, skill, resume, rag_content):
        """
        Generate a system prompt to query the LLM for evidence of a specific skill

        Parameters:
        - skill (str): The specific skill to assess from job requirements
        - resume (str): The full text of the candidate's resume
        - rag_content (str, optional): Additional RAG content queried from candidate documents

        Returns:
        - str: A system prompt tailored to assess the specific skill
        """

        system_prompt = f"""You are an objective skill assessor. Your task is to determine if a candidate possesses 
a SPECIFIC skill based solely on their resume and supporting evidence.

## SKILL TO ASSESS:
"{skill}"

## INSTRUCTIONS:
1. Focus exclusively on assessing the candidate's proficiency with the skill: "{skill}".
2. Examine the resume and supporting documents for both explicit mentions and clearly demonstrated applications of the skill.
3. Do NOT infer the skill unless it is either:
   - Directly mentioned in association with experience, or
   - Clearly evidenced through relevant tools, technologies, responsibilities, or outcomes.
   - Referenced as being used or listed as part of a broader initiative.
4. Evaluate each piece of evidence and assign a confidence rating:
   - STRONG: Explicit use with clear context or repeated/prolonged experience
   - MODERATE: Inferred through tools, environments, or outcomes (e.g., Python used in a listed project/tool)
   - WEAK: Mentioned in a list or indirectly implied without context
   - NONE: No relevant evidence
5. Provide up to 10 evidence_details. Each should include:
   - source: where the evidence appears
   - quote: an exact snippet from the document(s)
   - context: a short rationale explaining how this supports the skill claim
6. When no supporting evidence exists, output a "NONE" rating with an empty evidence details list.
7. IMPORTANT: Even if the skill is only listed in a technologies or languages section, consider this valid evidence.

## OUTPUT FORMAT:
```json
{{
  "skill_assessment": {{
    "skill": "{skill}",
    "evidence_found": true/false,
    "evidence_strength": "STRONG/MODERATE/WEAK/NONE",
    "description": "short (two to three sentence) description of what {skill} means with a concise example of what you're looking for",
    "evidence_details": [
      {{
        "source": "resume section/position/project",
        "quote": "exact text from resume showing evidence",
        "context": "brief explanation of how this demonstrates the skill"
      }}
    ]
  }}
}}
```

"""
# ## CANDIDATE RESUME:
# {resume}
# """

        # Add RAG content if provided
        if rag_content:
            system_prompt += f"""

## ADDITIONAL SUPPORTING EVIDENCE:
{rag_content}
"""

        # Add final instruction to prevent hallucination
        system_prompt += f"""

IMPORTANT: Be factual and precise. If you cannot find strong evidence for this specific skill, it's better to indicate "evidence_strength": WEAK than to stretch for connections. Focus only on "{skill}" and not similar skills unless they directly demonstrate the required skill. Remember that a skill listed in a "Languages" or "Technologies" section should be considered valid evidence.
"""

        prompt = f'Analyze candidate resume and supporting evidence to determine if they have experience with "{skill}". Adhere strictly to the output format requested.'

        return system_prompt, prompt

    def process_job_requirements(self, job_requirements, resume, rag_function):
        """
        Process all job requirements and build individual skill assessment prompts

        Parameters:
        - job_requirements (dict): The dict returned from job analysis
        - resume (str): The candidate's full resume
        - rag_function (callable): Function to retrieve relevant RAG content for a skill

        Returns:
        - dict: A dictionary of skill assessment prompts
        """
        skill_assessment_prompts = {}

        # Process technical skills - required
        for skill in job_requirements.get("technical_skills", {}).get("required", []):
            rag_content, metadata = rag_function(skill)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                skill, resume, rag_content
            )
            skill_assessment_prompts[f"technical_required_{skill}"] = (
                system_prompt,
                prompt,
                "technical_required",
                skill,
                metadata,
            )

        # Process technical skills - preferred
        for skill in job_requirements.get("technical_skills", {}).get("preferred", []):
            rag_content, metadata = rag_function(skill)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                skill, resume, rag_content
            )
            skill_assessment_prompts[f"technical_preferred_{skill}"] = (
                system_prompt,
                prompt,
                "technical_preferred",
                skill,
                metadata,
            )

        # Process experience requirements - required
        for exp in job_requirements.get("experience_requirements", {}).get(
            "required", []
        ):
            rag_content, metadata = rag_function(exp)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                exp, resume, rag_content
            )
            skill_assessment_prompts[f"experience_required_{exp}"] = (
                system_prompt,
                prompt,
                "experience_resuired",
                exp,
                metadata,
            )

        # Process experience requirements - preferred
        for exp in job_requirements.get("experience_requirements", {}).get(
            "preferred", []
        ):
            rag_content, metadata = rag_function(exp)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                exp, resume, rag_content
            )
            skill_assessment_prompts[f"experience_preferred_{exp}"] = (
                system_prompt,
                prompt,
                "experience_preferred",
                exp,
                metadata,
            )

        # Process education requirements
        for edu in job_requirements.get("education_requirements", []):
            rag_content, metadata = rag_function(edu)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                edu, resume, rag_content
            )
            skill_assessment_prompts[f"education_{edu}"] = (
                system_prompt,
                prompt,
                "education",
                edu,
                metadata,
            )

        # Process soft skills
        for skill in job_requirements.get("soft_skills", []):
            rag_content, metadata = rag_function(skill)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                skill, resume, rag_content
            )
            skill_assessment_prompts[f"soft_skill_{skill}"] = (
                system_prompt,
                prompt,
                "soft_skills",
                skill,
                metadata,
            )

        # Process industry knowledge
        for knowledge in job_requirements.get("industry_knowledge", []):
            rag_content, metadata = rag_function(knowledge)
            system_prompt, prompt = self.generate_skill_assessment_prompt(
                knowledge, resume, rag_content
            )
            skill_assessment_prompts[f"industry_knowledge_{knowledge}"] = (
                system_prompt,
                prompt,
                "industry_knowledge",
                knowledge,
                metadata,
            )

        return skill_assessment_prompts

    # Stage 1A: Job Analysis Implementation
    def create_job_analysis_prompt(self, job_description: str) -> tuple[str, str]:
        """Create the prompt for job requirements analysis."""
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
        system_prompt = """
      You are an objective job requirements analyzer. Your task is to extract and categorize the specific skills,
      experiences, and qualifications required in a job description WITHOUT any reference to any candidate.
      
      ## INSTRUCTIONS:
      
      1. Analyze ONLY the job description provided.
      2. Extract and categorize all requirements and preferences.
      3. DO NOT consider any candidate information - this is a pure job analysis task.
      
      ## OUTPUT FORMAT:
      
      ```json
      {
        "job_requirements": {
          "technical_skills": {
            "required": ["skill1", "skill2"],
            "preferred": ["skill1", "skill2"]
          },
          "experience_requirements": {
            "required": ["exp1", "exp2"],
            "preferred": ["exp1", "exp2"]
          },
          "education_requirements": ["req1", "req2"],
          "soft_skills": ["skill1", "skill2"],
          "industry_knowledge": ["knowledge1", "knowledge2"],
          "responsibilities": ["resp1", "resp2"],
          "company_values": ["value1", "value2"]
        }
      }
      ```
      
      Be specific and detailed in your extraction. Break down compound requirements into individual components.
      For example, "5+ years experience with React, Node.js and MongoDB" should be separated into:
      - Experience: "5+ years software development"
      - Technical skills: "React", "Node.js", "MongoDB"
      
      Avoid vague categorizations and be precise about whether skills are explicitly required or just preferred.
      """

        prompt = f"Job Description:\n{job_description}"
        return system_prompt, prompt

    async def analyze_job_requirements(
        self, message, job_description: str, metadata: Dict[str, Any]
    ) -> AsyncGenerator[Message, None]:
        """Analyze job requirements from job description."""
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
        try:
            system_prompt, prompt = self.create_job_analysis_prompt(job_description)
            metadata["system_prompt"] = system_prompt
            metadata["prompt"] = prompt
            async for message in self.call_llm(
                message=message, system_prompt=system_prompt, prompt=prompt
            ):
                if message.status != "done":
                    yield message
            if message.status == "error":
                return

            # Extract JSON from response
            json_str = self.extract_json_from_text(message.response)
            job_requirements = json.loads(json_str)

            self.validate_job_requirements(job_requirements)

            metadata["results"] = job_requirements["job_requirements"]

            message.status = "done"
            yield message
            return

        except Exception as e:
            metadata["error"] = message.response
            raise

    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 retrieve_rag_content(self, skill: str) -> tuple[str, ChromaDBGetResponse]:
        if self.context is None or self.context.user is None or self.context.user.file_watcher is None:
            raise ValueError("self.context or self.context.user.file_watcher is None")

        try:
            rag_results = ""
            rag_metadata = ChromaDBGetResponse()
            chroma_results = self.context.user.file_watcher.find_similar(query=skill, top_k=10, threshold=0.5)
            if chroma_results:
                query_embedding = np.array(chroma_results["query_embedding"]).flatten()

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

                rag_metadata = ChromaDBGetResponse( 
                    query=skill,
                    query_embedding=query_embedding.tolist(),
                    name="JPK",
                    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(),
                    size=self.context.user.file_watcher.collection.count()
                )

                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()
                    rag_results += f"""
Source: {metadata.get("doc_type", "unknown")}: {metadata.get("path", "")}
Document reference: {chroma_results["ids"][index]}
Content: { content }

"""
            return rag_results, rag_metadata

        except Exception as e:
            logger.error(traceback.format_exc())
            logger.error(e)
            exit(0)

    async def generate_resume(
        self, message: Message, job_description: str, resume: str
    ) -> AsyncGenerator[Message, None]:
        """
        Main function to generate a factually accurate tailored resume.

        Args:
            job_description: The job description text
            resume: The candidate's original resume text
            additional_context: Any additional context about the candidate (optional)

        Returns:
            Dict containing the generated resume and supporting analysis
        """
        if self.context is None:
            raise ValueError(f"context is None in {self.agent_type}")

        message.metadata.resume_generation = {}
        metadata = message.metadata.resume_generation
        # Stage 1A: Analyze job requirements
        streaming_message = Message(prompt="Analyze job requirements")
        streaming_message.status = "thinking"
        streaming_message.response = f"Analyzing job requirements\n\n"
        yield streaming_message

        metadata["requirements"] = {}
        async for message in self.analyze_job_requirements(
            message, job_description, metadata["requirements"]
        ):
            if message.status != "done":
                streaming_message.response += "."
                yield streaming_message
        if message.status == "error":
            return

        job_requirements = metadata.get("requirements", {}).get("results", {})
        metadata["candidate_skills"] = {}
        candidate_skills = metadata["candidate_skills"]
        prompts = self.process_job_requirements(
            job_requirements=job_requirements,
            resume=resume,
            rag_function=self.retrieve_rag_content,
        )

        # UI should persist this state of the message
        partial = message.model_copy()
        partial.status = "partial"
        partial.title = f"Found {len(prompts)} job requirements"
        yield partial
        self.conversation.add(partial)

        # Stage 1B: Analyze candidate qualifications
        message.status = "thinking"
        message.response = "Performing skill-based assessment"
        logger.info(message.response)
        yield message

        # Use each prompt to query an LLM
        total_prompts = len(prompts.items())
        index = 0
        skill_assessment_results: Dict[str, Any] = {}
        for skill_name, (
            system_prompt,
            prompt,
            section,
            description,
            rag,
        ) in prompts.items():
            index += 1
            # Query LLM with prompt
            streaming_message = Message(prompt=prompt)
            streaming_message.status = "thinking"
            streaming_message.response = (
                f"Analyzing skill {index}/{total_prompts}:\n\n* {description}\n\n"
            )
            yield streaming_message

            async for message in self.call_llm(
                message=message, system_prompt=system_prompt, prompt=prompt
            ):
                if message.status != "done":
                    streaming_message.response += "."
                    yield streaming_message
            if message.status == "error":
                return

            name = re.sub(rf"^{section}_", "", skill_name)
            json_str = self.extract_json_from_text(message.response)
            try:
                skill_assessment_results[skill_name] = json.loads(json_str)
                match_level = (
                    skill_assessment_results[skill_name]
                    .get("skill_assessment", {})
                    .get("evidence_strength", "UNKNOWN")
                )
                skill_description = (
                    skill_assessment_results[skill_name]
                    .get("skill_assessment", {})
                    .get("description", "")
                )
                if skill_description:  # Do not duplicate this item in the JSON data
                    del skill_assessment_results[skill_name]["skill_assessment"][
                        "description"
                    ]
            except Exception as e:
                skill_assessment_results[skill_name] = json_str
                match_level = "UNKNOWN"
                skill_description = ""

            # UI should persist this state of the message
            partial = message.model_copy()
            partial.status = "partial"
            partial.prompt = prompt
            partial.system_prompt = system_prompt
            partial.title = (
                f"Skill {index}/{total_prompts}: {description} [{match_level}]"
            )
            partial.metadata.rag = [rag] # Front-end expects a list of RAG retrievals
            if skill_description:
                partial.response = f"""
```json
{json.dumps(skill_assessment_results[skill_name]["skill_assessment"])}
```

{skill_description}
"""
            yield partial
            self.conversation.add(partial)

            if section not in candidate_skills:
                candidate_skills[section] = {}

            candidate_skills[section][name] = {
                "result": skill_assessment_results[skill_name],
            }

        candidate_info = {
            "full_name": self.context.user.full_name,
            "contact_info": self.context.user.contact_info
        }

        # Stage 2: Generate tailored resume
        message.status = "thinking"
        message.response = "Generating tailored resume"
        logger.info(message.response)
        yield message

        async for message in self.process_and_generate_resume(
            message=message,
            job_requirements=job_requirements,
            skill_assessment_results=skill_assessment_results,
            candidate_info=candidate_info,
            metadata=metadata,
        ):
            if message.status != "done":
                yield message
        if message.status == "error":
            return

        yield message
        return

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