"""LLM 投资事件抽取的数据模型 (M4)。

EventExtraction 是 LLM 严格输出 schema(JSON mode 解析后用 Pydantic 校验)。
ExtractedEvent 把 EventExtraction 与原文章元数据合并,作为 M4 最终落盘格式。
"""

from __future__ import annotations

import re
from datetime import datetime
from enum import StrEnum
from typing import Self

from pydantic import BaseModel, Field, field_validator, model_validator


class Sentiment(StrEnum):
    """事件情绪倾向。"""

    POSITIVE = "positive"  # 利好
    NEUTRAL = "neutral"    # 中性
    NEGATIVE = "negative"  # 利空


# 事件类型枚举(Prompt 中也会展示给 LLM)
EVENT_TYPES: tuple[str, ...] = (
    "业绩预告",
    "业绩快报",
    "财报披露",
    "合作签约",
    "投资并购",
    "重大合同",
    "产品发布",
    "技术突破",
    "监管处罚",
    "诉讼仲裁",
    "股东减持",
    "股东增持",
    "回购",
    "分红",
    "高管变动",
    "资产重组",
    "停牌复牌",
    "ST警示",
    "退市风险",
    "宏观政策",
    "行业政策",
    "国际局势",
    "其他",
)

# A 股股票代码:6 位数字(000xxx/300xxx/600xxx 等),也可带 .SH/.SZ/.BJ 后缀
_STOCK_CODE_RE = re.compile(r"^\d{6}(\.(SH|SZ|BJ))?$")

# 重要程度合理区间(LLM 偶尔会给 0/6/10,这里夹紧)
MIN_IMPORTANCE = 1
MAX_IMPORTANCE = 5


class EventExtraction(BaseModel):
    """LLM 输出的 JSON 直接映射到此模型。"""

    stock_codes: list[str] = Field(
        default_factory=list,
        description="A 股 6 位代码,允许带 .SH/.SZ/.BJ 后缀;无相关股票时为空",
    )
    company_names: list[str] = Field(
        default_factory=list, description="涉及公司中文简称,无关时为空"
    )
    industries: list[str] = Field(
        default_factory=list, description="所属行业(申万二级粒度优先);无关时为空"
    )
    sentiment: Sentiment = Field(..., description="positive/neutral/negative")
    importance: int = Field(
        ..., ge=MIN_IMPORTANCE, le=MAX_IMPORTANCE, description="1-5 重要程度"
    )
    event_type: str = Field(..., description="事件类型,见 EVENT_TYPES")
    summary: str = Field(
        default="",
        max_length=200,
        description="一句话事件摘要(≤ 100 字),便于人工浏览",
    )

    @field_validator("stock_codes")
    @classmethod
    def _strip_and_validate_stock_codes(cls, v: list[str]) -> list[str]:
        """剔除空字符串、统一大写、过滤明显非法格式。"""
        cleaned: list[str] = []
        for code in v:
            s = (code or "").strip().upper().replace(" ", "")
            if not s:
                continue
            if _STOCK_CODE_RE.match(s):
                cleaned.append(s)
        # 去重保持顺序
        seen: set[str] = set()
        out: list[str] = []
        for c in cleaned:
            if c not in seen:
                seen.add(c)
                out.append(c)
        return out

    @field_validator("company_names", "industries")
    @classmethod
    def _strip_text_lists(cls, v: list[str]) -> list[str]:
        cleaned = [(s or "").strip() for s in v]
        cleaned = [s for s in cleaned if s]
        seen: set[str] = set()
        out: list[str] = []
        for c in cleaned:
            if c not in seen:
                seen.add(c)
                out.append(c)
        return out

    @field_validator("event_type")
    @classmethod
    def _normalize_event_type(cls, v: str) -> str:
        s = (v or "").strip()
        if not s:
            return "其他"
        return s

    @model_validator(mode="after")
    def _post_check(self) -> Self:
        """中性情绪时 importance 应较低(1-3),纠正常见误判。"""
        # 不强制纠正,留作后续校准。占位,便于以后扩展。
        return self


class ExtractedEvent(BaseModel):
    """落盘格式:文章元数据 + LLM 抽取结果 + 调用元信息。"""

    # ---- 来源标识 ----
    source_id: str
    url: str
    url_hash: str
    title: str
    publish_time: datetime | None = None

    # ---- 抽取结果 ----
    event: EventExtraction

    # ---- 调用元信息 ----
    provider: str = Field(..., description="deepseek / qwen 等")
    model: str
    extracted_at: datetime = Field(default_factory=datetime.now)
    attempts: int = Field(default=1, ge=1, description="LLM 实际调用次数(含重试)")
    prompt_tokens: int | None = None
    completion_tokens: int | None = None

    def short_summary(self) -> str:
        ev = self.event
        codes = ",".join(ev.stock_codes) or "-"
        return (
            f"[{self.source_id}] {self.title[:30]} "
            f"-> {ev.sentiment.value}/{ev.importance}/{ev.event_type} "
            f"({codes})"
        )


class LLMCallError(Exception):
    """LLM 调用失败(网络 / 解析 / 校验)。"""

    def __init__(self, reason: str, *, attempts: int = 0) -> None:
        super().__init__(reason)
        self.reason = reason
        self.attempts = attempts