from typing import Dict, Any
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from llm.mistral_client import get_mistral_llm


def jd_agent(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    JD Intelligence Agent
    ---------------------
    - Extracts required skills
    - Nice-to-have skills
    - Experience level
    - Hiring priorities
    """

    llm = get_mistral_llm()

    prompt = PromptTemplate(
        template="""
You are a senior hiring manager AI.

Analyze the following Job Description and extract information
in a STRICT JSON format with the following keys:

{
  "role_title": string,
  "required_skills": [string],
  "nice_to_have_skills": [string],
  "experience_level": string,
  "responsibilities": [string],
  "hiring_priorities": [string]
}

Rules:
- Do NOT add extra keys
- Do NOT return text outside JSON
- Be concise and accurate

Job Description:
{job_description}
""",
        input_variables=["job_description"]
    )

    parser = JsonOutputParser()

    chain = prompt | llm | parser

    try:
        result = chain.invoke(
            {"job_description": state["job_description"]}
        )
    except Exception as e:
        # Fallback in case LLM returns invalid JSON
        result = {
            "role_title": "Unknown",
            "required_skills": [],
            "nice_to_have_skills": [],
            "experience_level": "Unknown",
            "responsibilities": [],
            "hiring_priorities": [],
            "error": str(e)
        }

    state["jd_requirements"] = result
    return state