src/ai_classifier.py

"""
AI 分类模块 - 调用 LLM 对邮件进行分类、提取公司名、生成摘要

所有配置值均通过 src.config 模块读取（支持运行时动态更新）。
"""
import json
import logging
import asyncio
from dataclasses import dataclass
from typing import Optional

from openai import AsyncOpenAI
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

from . import config as cfg
from .email_parser import ParsedEmail

logger = logging.getLogger(__name__)


@dataclass
class ClassificationResult:
    """AI 分类结果"""
    category: str = ""
    company_name: str = ""
    summary: str = ""
    confidence: str = ""
    raw_response: str = ""
    error: str = ""

    @property
    def is_valid(self) -> bool:
        return not self.error and bool(self.category) and bool(self.company_name)


# 系统提示词
SYSTEM_PROMPT = """You are an expert email classifier for job application tracking.
Your task is to analyze recruitment reply emails and extract structured information.

You MUST respond with valid JSON only. No markdown, no explanation.

Classification rules (strictly follow this mapping):
- "Auto reply": Automated/out-of-office reply from a company system
- "career page": Reply suggests visiting their career/jobs page
- "connect more": They want more information about the applicant (e.g., "Tell me more about yourself")
- "considered": They'll keep your profile for future opportunities
- "not hiring": Not currently hiring for this role
- "do not exist": Delivery failure / bounce / email address doesn't exist
- "not fit": Applicant doesn't match the requirements
- "Forward": They forwarded your CV to HR or relevant team

Company name extraction rules:
- Extract the CORE company name only
- Remove suffixes: Inc., LLC, Ltd., Corp., Co., Limited, GmbH, etc.
- Example: "HomeLight Inc." -> "HomeLight"
- If unclear from body, use the sender email domain (e.g., @stripe.com -> "Stripe")
- If truly unknown, return "Unknown"

Summary rules:
- Extract the single most important sentence from the email body
- Should capture the key intent or action item
- Max 100 characters
- In English

Respond with this exact JSON structure:
{
  "category": "<one of the 8 categories>",
  "company_name": "<core company name>",
  "summary": "<key sentence from email>",
  "confidence": "<high|medium|low>"
}"""


def _build_user_prompt(email: ParsedEmail) -> str:
    body_truncated = email.body_text[:3000] if email.body_text else ""
    return f"""Analyze this recruitment reply email:

From: {email.sender}
Subject: {email.subject}
Date: {email.date_str}

Email Body:
{body_truncated}

Extract the category, company name, and key summary."""


def _parse_llm_response(content: str, email: ParsedEmail) -> ClassificationResult:
    """解析 LLM 返回的 JSON 内容"""
    content = content.strip()
    if content.startswith("```"):
        lines = content.split("\n")
        content = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
    content = content.strip()

    try:
        data = json.loads(content)
    except json.JSONDecodeError as e:
        logger.error(f"LLM 响应 JSON 解析失败: {e}\n原始内容: {content}")
        import re
        match = re.search(r'\{[^{}]+\}', content, re.DOTALL)
        if match:
            try:
                data = json.loads(match.group())
            except Exception:
                return ClassificationResult(
                    category="", company_name="", summary="",
                    confidence="low", raw_response=content,
                    error=f"JSON 解析失败: {e}"
                )
        else:
            return ClassificationResult(
                category="", company_name="", summary="",
                confidence="low", raw_response=content,
                error=f"JSON 解析失败: {e}"
            )

    category = data.get("category", "").strip()
    all_cats = cfg.ALL_CATEGORIES
    if category not in all_cats:
        for valid_cat in all_cats:
            if valid_cat.lower() == category.lower():
                category = valid_cat
                break
        else:
            logger.warning(f"未知分类标签: {category}，将使用 'Auto reply'")
            category = "Auto reply"

    company_name = data.get("company_name", "").strip()
    if not company_name or company_name.lower() == "unknown":
        company_name = email.sender_domain.title() if email.sender_domain else "Unknown"

    summary = data.get("summary", "").strip()[:200]
    confidence = data.get("confidence", "medium").lower()
    if confidence not in ("high", "medium", "low"):
        confidence = "medium"

    return ClassificationResult(
        category=category,
        company_name=company_name,
        summary=summary,
        confidence=confidence,
        raw_response=content,
    )


@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
    retry=retry_if_exception_type((Exception,)),
    reraise=True,
)
async def classify_email_async(email: ParsedEmail) -> ClassificationResult:
    """异步分类单封邮件"""
    if not email.is_valid:
        return ClassificationResult(
            category="", company_name="", summary="",
            confidence="low",
            error=f"邮件解析失败: {email.parse_error}"
        )

    try:
        api_key = cfg.get_llm_api_key()
        if not api_key:
            return ClassificationResult(
                category="", company_name=email.sender_domain.title() or "Unknown",
                summary="", confidence="low",
                error="❌ 未配置 LLM API Key，请在设置面板中填写"
            )

        model = cfg.get_llm_model()
        base_url = cfg.get_llm_base_url()

        client_kwargs = {"api_key": api_key}
        if base_url:
            client_kwargs["base_url"] = base_url

        client = AsyncOpenAI(**client_kwargs)
        user_prompt = _build_user_prompt(email)

        logger.debug(f"正在分类邮件: {email.filename} (模型: {model})")

        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.1,
            max_tokens=256,
            response_format={"type": "json_object"},
        )

        content = response.choices[0].message.content or ""
        result = _parse_llm_response(content, email)
        logger.info(
            f"分类完成: {email.filename} -> [{result.category}] {result.company_name} "
            f"(置信度: {result.confidence})"
        )
        return result

    except Exception as e:
        logger.error(f"LLM API 调用失败 ({email.filename}): {e}", exc_info=True)
        raise


def classify_email_sync(email: ParsedEmail) -> ClassificationResult:
    """同步版本"""
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                future = pool.submit(asyncio.run, classify_email_async(email))
                return future.result(timeout=30)
        else:
            return loop.run_until_complete(classify_email_async(email))
    except RuntimeError:
        return asyncio.run(classify_email_async(email))


async def batch_classify_emails(
    emails: list[ParsedEmail],
    progress_callback=None,
) -> list[ClassificationResult]:
    """批量异步分类邮件（带并发控制）"""
    max_concurrent = cfg.get_max_concurrent()
    semaphore = asyncio.Semaphore(max_concurrent)
    results = [None] * len(emails)
    completed = [0]

    async def classify_with_semaphore(idx: int, email: ParsedEmail):
        async with semaphore:
            try:
                result = await classify_email_async(email)
            except Exception as e:
                result = ClassificationResult(
                    category="", company_name=email.sender_domain.title() or "Unknown",
                    summary="", confidence="low",
                    error=f"分类失败: {str(e)}"
                )
            results[idx] = result
            completed[0] += 1
            if progress_callback:
                await progress_callback(completed[0], len(emails), result, email)
            return result

    tasks = [classify_with_semaphore(i, email) for i, email in enumerate(emails)]
    await asyncio.gather(*tasks)
    return results