src/llm_client.py

import json
import re
import time
from openai import OpenAI
from config import NVIDIA_API_KEY, GLM_BASE_URL, GLM_MODEL


class LLMClient:
    # Hard cap per API request — a hung call must fail fast, not stall the pipeline
    REQUEST_TIMEOUT = 90.0

    def __init__(self):
        self.client = OpenAI(
            base_url=GLM_BASE_URL,
            api_key=NVIDIA_API_KEY,
            timeout=self.REQUEST_TIMEOUT,
            max_retries=0,  # we do our own retries with backoff
        )
        self.model = GLM_MODEL

    def _call(self, system: str, user: str, max_tokens: int = 512, retries: int = 3) -> str:
        for attempt in range(retries):
            try:
                completion = self.client.chat.completions.create(
                    model=self.model,
                    messages=[
                        {"role": "system", "content": system},
                        {"role": "user", "content": user},
                    ],
                    temperature=0.2,
                    top_p=0.9,
                    max_tokens=max_tokens,
                    stream=False,
                )
                return completion.choices[0].message.content or ""
            except Exception as e:
                if attempt < retries - 1:
                    time.sleep(2 ** attempt)
                else:
                    raise

    def _extract_json(self, text: str) -> dict | list:
        try:
            return json.loads(text)
        except Exception:
            pass
        match = re.search(r"```(?:json)?\s*([\s\S]+?)```", text)
        if match:
            try:
                return json.loads(match.group(1))
            except Exception:
                pass
        for start_char, end_char in [('{', '}'), ('[', ']')]:
            idx = text.find(start_char)
            if idx >= 0:
                depth = 0
                for i, ch in enumerate(text[idx:], idx):
                    if ch == start_char:
                        depth += 1
                    elif ch == end_char:
                        depth -= 1
                    if depth == 0:
                        try:
                            return json.loads(text[idx:i+1])
                        except Exception:
                            break
        raise ValueError(f"Cannot parse JSON: {text[:200]}")

    # ──────────────────────────────────────────────────────────
    # BATCH ASSESSMENT — sends 8 jobs per API call (8× faster)
    # ──────────────────────────────────────────────────────────
    def assess_jobs_batch(self, jobs_batch: list[dict], compact_profile: str) -> list[dict]:
        """
        Assess a batch of up to 8 jobs in one API call.
        Returns list of assessment dicts in the same order as jobs_batch.
        """
        system = (
            "You are a PM recruiter. Rate job-candidate fit. "
            "Return ONLY a JSON array — no markdown, no text."
        )

        jobs_text = ""
        for i, job in enumerate(jobs_batch, 1):
            desc = (job.get("description") or "")[:300].replace("\n", " ")
            jobs_text += (
                f"\nJOB {i}: {job.get('title','')} at {job.get('company','')} | {job.get('location','')}\n"
                f"DESC: {desc}\n"
            )

        user = f"""CANDIDATE: {compact_profile}

{jobs_text}
Return a JSON array with one object per job (in order):
[
  {{
    "job_index": 1,
    "score": <1-10>,
    "match_pct": <0-100>,
    "exp_match": "Good fit|Under-qualified|Over-qualified",
    "matching": ["skill1","skill2"],
    "missing": ["skill1"],
    "strengths": ["point1","point2"],
    "note": "<1 sentence>",
    "keywords": ["kw1","kw2","kw3"],
    "priority": "High|Medium|Low"
  }}
]"""

        response = self._call(system, user, max_tokens=150 * len(jobs_batch))
        try:
            result = self._extract_json(response)
            if isinstance(result, list):
                return result
            # Sometimes model wraps in object
            if isinstance(result, dict):
                for v in result.values():
                    if isinstance(v, list):
                        return v
        except Exception:
            pass
        # Fallback: return neutral scores
        return [self._neutral_assessment(i + 1) for i in range(len(jobs_batch))]

    def _neutral_assessment(self, idx: int) -> dict:
        return {
            "job_index": idx, "score": 5, "match_pct": 50,
            "exp_match": "Unknown", "matching": [], "missing": [],
            "strengths": [], "note": "Auto-assessment failed.",
            "keywords": [], "priority": "Medium",
        }

    def _call_with_cfg(self, cfg: dict, system: str, user: str, max_tokens: int = 2000) -> str:
        """Call any NVIDIA model using the provided model config dict."""
        from openai import OpenAI
        client = OpenAI(base_url=cfg.get("base_url"), api_key=cfg["api_key"],
                        timeout=self.REQUEST_TIMEOUT, max_retries=0)
        extra_body = cfg.get("extra_body") or None
        for attempt in range(3):
            try:
                kwargs = dict(
                    model=cfg["model"],
                    messages=[
                        {"role": "system", "content": system},
                        {"role": "user",   "content": user},
                    ],
                    temperature=0.2,
                    top_p=0.9,
                    max_tokens=max_tokens,
                    stream=False,
                )
                if extra_body:
                    kwargs["extra_body"] = extra_body
                completion = client.chat.completions.create(**kwargs)
                return completion.choices[0].message.content or ""
            except Exception:
                if attempt < 2:
                    time.sleep(2 ** attempt)
                else:
                    raise

    def customize_resume_fast(self, cfg: dict, resume_text: str, job_description: str,
                               job_title: str, company: str, assessment: dict) -> dict:
        """Customize resume using a fast model (Kimi/Step/Qwen) instead of GLM."""
        system = (
            "You are an ATS resume writer for PM roles. "
            "Return ONLY valid JSON, no markdown."
        )
        matching_skills = ", ".join(assessment.get("matching_skills", assessment.get("matching", [])))
        ats_keywords    = ", ".join(assessment.get("ats_keywords", assessment.get("keywords", [])))

        user = self._resume_customize_prompt(
            resume_text, job_description, job_title, company, ats_keywords, matching_skills
        )
        for attempt in range(2):
            try:
                response = self._call_with_cfg(cfg, system, user, max_tokens=4000)
                data = self._extract_json(response)
                # Some models wrap the object in an array
                if isinstance(data, list):
                    data = next((d for d in data if isinstance(d, dict)), {})
                if self._customization_valid(data):
                    return data
            except Exception:
                pass
        return self._empty_customization()

    @staticmethod
    def _customization_valid(data) -> bool:
        """A usable customization must have a real summary and skills list."""
        return (
            isinstance(data, dict)
            and len(data.get("professional_summary", "") or "") > 50
            and len(data.get("core_competencies", []) or []) >= 5
        )

    # ──────────────────────────────────────────────────────────
    # RESUME CUSTOMIZATION
    # ──────────────────────────────────────────────────────────
    def customize_resume(self, resume_text: str, job_description: str, job_title: str, company: str, assessment: dict) -> dict:
        system = (
            "You are an ATS resume writer for PM roles. "
            "Return ONLY valid JSON, no markdown."
        )
        matching_skills = ", ".join(assessment.get("matching_skills", assessment.get("matching", [])))
        ats_keywords    = ", ".join(assessment.get("ats_keywords", assessment.get("keywords", [])))
        user = self._resume_customize_prompt(
            resume_text, job_description, job_title, company, ats_keywords, matching_skills
        )
        try:
            response = self._call(system, user, max_tokens=4000)
            data = self._extract_json(response)
            if isinstance(data, list):
                data = next((d for d in data if isinstance(d, dict)), {})
            if self._customization_valid(data):
                return data
        except Exception:
            pass
        return self._empty_customization()

    def _resume_customize_prompt(self, resume_text, job_description, job_title, company,
                                  ats_keywords, matching_skills) -> str:
        kw_list = ats_keywords if ats_keywords else "product manager, agile, roadmap, stakeholder, KPI, user research"
        return f"""You are an expert ATS resume writer. Rewrite this resume to score 95%+ on ATS for the role below.

TARGET ROLE: {job_title} at {company}

JOB DESCRIPTION:
{job_description[:2000]}

CANDIDATE'S ORIGINAL RESUME:
{resume_text[:2500]}

MANDATORY ATS KEYWORDS (you MUST include ALL of these naturally in the resume):
{kw_list}

RULES FOR 95%+ ATS SCORE:
1. Mirror the exact language from the JD — use the same phrases, not synonyms
2. Every bullet point MUST start with a strong action verb (Led, Built, Drove, Scaled, Launched, Reduced, Increased, Delivered)
3. Every bullet MUST include a quantified metric (%, numbers, $ impact, time saved, users impacted)
4. Professional summary must open with the exact job title from the JD and include 3+ keywords from the list
5. Core competencies must include ALL mandatory keywords above plus 6+ tools/frameworks from the JD
6. Include PM-specific terms: product roadmap, go-to-market, sprint, backlog, user story, A/B testing, funnel, retention
7. Do NOT add skills the candidate doesn't have — rephrase existing experience to match JD language

Return ONLY valid JSON (no markdown):
{{
  "professional_summary": "<4-5 sentences. Open with exact job title. Include 5+ keywords. Quantify impact.>",
  "core_competencies": ["skill1","skill2","skill3","skill4","skill5","skill6","skill7","skill8","skill9","skill10","skill11","skill12","skill13","skill14","skill15"],
  "experience_bullets": {{
    "role_name_1": ["• Led X resulting in Y% improvement", "• Built Z used by N users", "• Drove A increasing B by C%"],
    "role_name_2": ["• Launched X achieving Y", "• Reduced X by N%"]
  }},
  "key_achievements": ["Achieved X resulting in Y", "Built Z growing metric by N%", "Led team of N to deliver X on time"],
  "cover_letter_intro": "<2 sentences tailored to {company} specifically>",
  "tailoring_notes": "<which keywords were added and where>"
}}"""

    def _empty_customization(self) -> dict:
        return {
            "professional_summary": "",
            "core_competencies": [],
            "experience_bullets": {},
            "key_achievements": [],
            "cover_letter_intro": "",
            "tailoring_notes": "Auto-customization failed.",
        }

    # ──────────────────────────────────────────────────────────
    # PROFILE EXTRACTION
    # ──────────────────────────────────────────────────────────
    def extract_profile_summary(self, resume_text: str) -> str:
        system = "You are a resume parser. Return ONLY valid JSON."
        user = f"""Parse this resume for job matching.

RESUME:
{resume_text[:3000]}

Return ONLY:
{{
  "name": "<name>",
  "current_role": "<role>",
  "total_experience_years": <N>,
  "summary": "<2 sentence summary>",
  "core_skills": ["s1","s2","s3","s4","s5","s6","s7","s8"],
  "domain_expertise": ["d1","d2"],
  "industries": ["i1","i2"],
  "education": "<degree + institution>",
  "certifications": ["c1"],
  "notable_achievements": ["a1","a2","a3"]
}}"""

        response = self._call(system, user, max_tokens=800)
        try:
            data = self._extract_json(response)
            return json.dumps(data, indent=2)
        except Exception:
            return resume_text[:1500]

    def extract_profile_summary_fast(self, cfg: dict, resume_text: str) -> str:
        """Like extract_profile_summary but uses a fast model (Kimi/Step) instead of GLM."""
        system = "You are a resume parser. Return ONLY valid JSON."
        user = f"""Parse this resume for job matching.

RESUME:
{resume_text[:3000]}

Return ONLY:
{{
  "name": "<name>",
  "current_role": "<role>",
  "total_experience_years": <N>,
  "summary": "<2 sentence summary>",
  "core_skills": ["s1","s2","s3","s4","s5","s6","s7","s8"],
  "domain_expertise": ["d1","d2"],
  "industries": ["i1","i2"],
  "education": "<degree + institution>",
  "certifications": ["c1"],
  "notable_achievements": ["a1","a2","a3"]
}}"""
        try:
            response = self._call_with_cfg(cfg, system, user, max_tokens=800)
            data = self._extract_json(response)
            return json.dumps(data, indent=2)
        except Exception:
            return self.extract_profile_summary(resume_text)  # fallback to GLM

    def build_compact_profile(self, profile_json: str) -> str:
        """Build a short ~200-char profile string for batch assessments."""
        try:
            d = json.loads(profile_json)
            skills = ", ".join(d.get("core_skills", [])[:8])
            return (
                f"{d.get('name','')} | {d.get('current_role','')} | "
                f"{d.get('total_experience_years','')} yrs exp | "
                f"Skills: {skills} | "
                f"Domain: {', '.join(d.get('domain_expertise',[])[:3])}"
            )
        except Exception:
            return profile_json[:300]