src/agents/planner.py

from __future__ import annotations

import asyncio
import json
import logging
import re
from typing import Any

from src.agents.prompts import PLANNER_SYSTEM_PROMPT
from src.core.config import get_llm_client, get_settings
from src.core.constants import INDIAN_CITIES, INDIAN_COMPANIES
from src.core.models import (
    ExperienceRequirements,
    LocationRequirements,
    ParsedQuery,
    PreferredSkill,
    RequiredSkill,
    SkillImportance,
)
from src.language.code_mixed import CodeMixedProcessor
from src.matching.skill_matcher import SKILL_ALIASES

logger = logging.getLogger(__name__)


def _strip_json_fences(content: str) -> str:
    """Strip markdown JSON code fences (```json ... ```) from LLM output."""
    content = content.strip()
    if content.startswith("```"):
        content = re.sub(r"^```(?:json)?\s*", "", content)
        content = re.sub(r"\s*```$", "", content)
    return content.strip()


class PlannerAgent:
    def __init__(self) -> None:
        self._client = None
        self._code_mixed_processor: CodeMixedProcessor | None = None
        settings = get_settings()
        self.model = settings.openai_model

    @property
    def client(self) -> Any:
        if self._client is None:
            try:
                self._client = get_llm_client()
            except Exception:
                logger.warning("LLM client unavailable for planner")
                self._client = None
        return self._client

    async def plan(self, raw_query: str) -> ParsedQuery:  # noqa: C901
        # mypy: return type is inferred from the return statements below
        try:
            if self._code_mixed_processor is None:
                self._code_mixed_processor = CodeMixedProcessor()
            processor = self._code_mixed_processor
            if processor.detect_code_mixed(raw_query):
                logger.info("Code-mixed query detected, applying TinT prompting")
                tint_query = (
                    "[Translate-in-Thought] The following query contains Hinglish "
                    "(Hindi-English code-mixed text). Internally translate it to English "
                    "before parsing, then output the JSON result.\n\n"
                    "Query: " + raw_query
                )
            else:
                tint_query = raw_query

            from langchain_core.messages import HumanMessage, SystemMessage

            messages = [
                SystemMessage(content=PLANNER_SYSTEM_PROMPT),
                HumanMessage(content=tint_query),
            ]
            response = await asyncio.wait_for(
                self.client.ainvoke(messages), timeout=30.0,
            )
            content = response.content if hasattr(response, "content") else str(response)
            if not content or not content.strip():
                logger.warning("Planner LLM returned empty content, using fallback")
                return self._fallback_parse(raw_query)
            content = _strip_json_fences(content)
            try:
                parsed = json.loads(content)
                return ParsedQuery(**parsed)
            except (json.JSONDecodeError, Exception) as e:
                logger.warning(f"Planner LLM parse failed: {e}. Raw: {content[:200]}")
                return self._fallback_parse(raw_query)
        except TimeoutError:
            logger.warning("Planner LLM timed out after 30s, using fallback")
            return self._fallback_parse(raw_query)
        except Exception as e:
            logger.warning(f"Planner LLM failed: {type(e).__name__}: {e}, using fallback")
            return self._fallback_parse(raw_query)

    async def replan(
        self, original_query: str, previous_params: dict[str, Any], feedback: str,
    ) -> ParsedQuery:
        try:
            prompt = (
                "Original query: " + original_query + "\n"
                "Previous params: " + json.dumps(previous_params) + "\n"
                "Feedback: " + feedback + "\n"
                "Revise the search parameters. Output valid JSON only."
            )
            from langchain_core.messages import HumanMessage, SystemMessage

            messages = [
                SystemMessage(content=PLANNER_SYSTEM_PROMPT),
                HumanMessage(content=prompt),
            ]
            response = await asyncio.wait_for(
                self.client.ainvoke(messages), timeout=30.0,
            )
            content = response.content if hasattr(response, "content") else str(response)
            content = _strip_json_fences(content)
            if not content or not content.strip():
                logger.warning("Replan LLM returned empty content, using fallback")
                relaxed = self._relax_params(previous_params)
                return ParsedQuery(**relaxed)
            parsed = json.loads(content)
            return ParsedQuery(**parsed)
        except TimeoutError:
            logger.warning("Replan LLM timed out after 30s, using fallback")
            relaxed = self._relax_params(previous_params)
            return ParsedQuery(**relaxed)
        except Exception as e:
            logger.warning(f"Replan LLM failed, using fallback: {e}")
            relaxed = self._relax_params(previous_params)
            return ParsedQuery(**relaxed)

    def _fallback_parse(self, query: str) -> ParsedQuery:
        required: list[RequiredSkill] = []
        preferred: list[PreferredSkill] = []
        min_years: float | None = None
        max_years: float | None = None
        industry: str | None = None
        city: str | None = None

        for alias, aliases in SKILL_ALIASES.items():
            candidates = [alias] + aliases
            if any(qs in query.lower() for qs in candidates):
                required.append(
                    RequiredSkill(name=alias.title(), importance=SkillImportance.REQUIRED)
                )

        year_match = re.search(r"(\d+)\+?\s*(?:years?|yrs?)", query.lower())
        if year_match:
            min_years = float(year_match.group(1))

        year_range = re.search(r"(\d+)\s*[-to]+\s*(\d+)\s*(?:years?|yrs?)", query.lower())
        if year_range:
            min_years = float(year_range.group(1))
            max_years = float(year_range.group(2))

        for c in INDIAN_CITIES:
            if c.lower() in query.lower():
                city = c
                break

        for comp in INDIAN_COMPANIES:
            if comp.lower() in query.lower():
                industry = "technology"
                break

        return ParsedQuery(
            required_skills=required,
            preferred_skills=preferred,
            experience=ExperienceRequirements(  # noqa: E501
                min_years=min_years, max_years=max_years, industry=industry,
            ),
            location=LocationRequirements(city=city, remote_ok="remote" in query.lower()),
        )

    def _relax_params(self, params: dict) -> dict:
        params = dict(params)
        exp = dict(params.get("experience", {}))
        if exp.get("min_years") is not None:
            exp["min_years"] = max(0, exp["min_years"] - 2)
        if exp.get("max_years") is not None:
            exp["max_years"] = (exp["max_years"] or 0) + 3
        params["experience"] = exp

        loc = dict(params.get("location", {}))
        loc["city"] = None
        loc["remote_ok"] = True
        params["location"] = loc

        params["required_skills"] = []
        return params