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 asyncio
import time
import asyncio
import numpy as np  # type: ignore

from .base import Agent, agent_registry, LLMMessage
from models import ApiActivityType, Candidate, ChatMessage, ChatMessageError, ChatMessageMetaData, ApiMessageType, ChatMessageStatus, ChatMessageUser, ChatOptions, ChatSenderType, ApiStatusType, JobRequirements, JobRequirementsMessage, Tunables
import model_cast
from logger import logger
import defines
import backstory_traceback as traceback

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

    # 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 company information, job title, and all requirements.
      3. Extract and categorize all requirements and preferences.
      4. DO NOT consider any candidate information - this is a pure job analysis task.
      
      ## OUTPUT FORMAT:
      
      ```json
      { 
        "company_name": "Company Name",
        "job_title": "Job Title",
        "job_summary": "Brief summary of the job",
        "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, llm: Any, model: str, session_id: str, prompt: str
    ) -> AsyncGenerator[ChatMessage | ChatMessageError, None]:
        """Analyze job requirements from job description."""
        system_prompt, prompt = self.create_job_analysis_prompt(prompt)
        generated_message = None
        async for generated_message in self.llm_one_shot(llm, model, session_id=session_id, prompt=prompt, system_prompt=system_prompt):
            if generated_message.status == ApiStatusType.ERROR:
                yield generated_message
                return
            if generated_message.status != ApiStatusType.DONE:
                yield generated_message

        if not generated_message:
            error_message = ChatMessageError(
                session_id=session_id,
                content="Job requirements analysis failed to generate a response.")
            logger.error(f"⚠️ {error_message.content}")
            yield error_message
            return
        
        yield generated_message
        return

    async def generate(
        self, llm: Any, model: str, session_id: str, prompt: str, tunables: Optional[Tunables] = None, temperature=0.7
    ) -> AsyncGenerator[ChatMessage, None]:
        # Stage 1A: Analyze job requirements
        status_message = ChatMessageStatus(
            session_id=session_id,
            content = f"Analyzing job requirements",
            activity=ApiActivityType.THINKING
        )
        yield status_message

        generated_message = None
        async for generated_message in self.analyze_job_requirements(llm, model, session_id, prompt):
            if generated_message.status == ApiStatusType.ERROR:
                yield generated_message
                return
            if generated_message.status != ApiStatusType.DONE:
                yield generated_message

        if not generated_message:
            error_message = ChatMessageError(
                session_id=session_id,
                content="Job requirements analysis failed to generate a response."
            )
            logger.error(f"⚠️ {error_message.content}")
            yield error_message
            return
        
        requirements = None
        job_requirements_data = ""
        company = ""
        summary = ""
        title = ""
        try:
            json_str = self.extract_json_from_text(generated_message.content)
            requirements_json = json.loads(json_str)

            company = requirements_json.get("company_name", "")
            title = requirements_json.get("job_title", "")
            summary = requirements_json.get("job_summary", "")
            job_requirements_data = requirements_json.get("job_requirements", None)
            requirements = JobRequirements.model_validate(job_requirements_data)
        except json.JSONDecodeError as e:
            status_message.status = ApiStatusType.ERROR
            status_message.content = f"Failed to parse job requirements JSON: {str(e)}\n\n{job_requirements_data}"
            logger.error(f"⚠️ {status_message.content}")
            yield status_message
            return
        except ValueError as e:
            status_message.status = ApiStatusType.ERROR
            status_message.content = f"Job requirements validation error: {str(e)}\n\n{job_requirements_data}"
            logger.error(f"⚠️ {status_message.content}")
            logger.error(f"Content: {prompt}")
            yield status_message
            return
        except Exception as e:
            status_message.status = ApiStatusType.ERROR
            status_message.content = f"Unexpected error processing job requirements: {str(e)}\n\n{job_requirements_data}"
            logger.error(traceback.format_exc())
            logger.error(f"⚠️ {status_message.content}")
            yield status_message
            return
        job_requirements_message = JobRequirementsMessage(
            session_id=session_id,
            status=ApiStatusType.DONE,
            requirements=requirements,
            company=company,
            title=title,
            summary=summary,
            description=prompt,
        )
        yield job_requirements_message
        logger.info(f"✅ Job requirements analysis completed successfully.")
        return

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