"""
Entity Extractor — LangExtract + DeepSeek entity extraction.
Independent implementation for the GraphRAG Studio backend.
"""
from __future__ import annotations

import os
from pathlib import Path

from dotenv import load_dotenv

import langextract as lx
from langextract.providers.openai import OpenAILanguageModel

load_dotenv(Path(__file__).parent.parent / ".env", override=True)

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "")
DEEPSEEK_BASE_URL = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
MODEL_ID = "deepseek-chat"

PROMPT_DESCRIPTION = (
    "从建筑规范文本中按出现顺序提取命名实体。"
    "实体类型："
    "构件 — 建筑构件和部位（防火墙、防火门、楼梯间、电梯井、管道井、变形缝、外墙、屋面、楼板、梁、柱、承重墙、隔墙、门窗、幕墙、雨篷、阳台、走廊、前室、竖井等）；"
    "设备 — 消防设备和系统（消火栓、灭火器、喷淋系统、报警系统、防烟排烟系统、应急照明、疏散指示、防火卷帘、消防电梯、消防水箱、消防水泵、气体灭火系统、火灾自动报警系统、电气火灾监控系统等）；"
    "设施 — 建筑设施和配套（停车场、锅炉房、变配电室、发电机房、储油间、空调机房、通风机房、电梯机房、消防控制室、水泵房等）；"
    "指标 — 技术参数和量化要求（耐火等级、防火间距、疏散宽度、最大面积、高度限值、距离要求、时间要求、容量要求、温度限值、压力要求等）；"
    "场所 — 建筑分类和空间类型（高层建筑、地下建筑、公共建筑、住宅建筑、工业建筑、商业建筑、医疗建筑、教育建筑、娱乐场所、仓库、厂房、中庭、避难层、避难走道、敞开楼梯间等）；"
    "材料 — 建筑材料和产品（不燃材料、难燃材料、可燃材料、防火涂料、防火玻璃、防火密封件、保温材料、装饰材料、钢结构防火保护材料等）；"
    "措施 — 防火措施和策略（防火分隔、防火封堵、自然排烟、机械加压送风、机械排烟、安全疏散、防火保护、消防供电、消防水源、灭火救援等）；"
    "条款 — 规范条文引用（强制性条文、推荐性条文、术语定义、一般规定、基本要求等）。"
)

EXAMPLES = [
    lx.data.ExampleData(
        text="公共建筑应设置防火墙。甲级防火门耐火极限为2.0h。消防控制室采用不燃材料。",
        extractions=[
            lx.data.Extraction(extraction_class="场所", extraction_text="公共建筑"),
            lx.data.Extraction(extraction_class="构件", extraction_text="防火墙"),
            lx.data.Extraction(extraction_class="构件", extraction_text="甲级防火门"),
            lx.data.Extraction(extraction_class="指标", extraction_text="2.0h"),
            lx.data.Extraction(extraction_class="设施", extraction_text="消防控制室"),
            lx.data.Extraction(extraction_class="材料", extraction_text="不燃材料"),
        ],
    ),
    lx.data.ExampleData(
        text="住宅建筑应设室内消火栓。高位消防水箱容积为18m3。消防电梯前室设安全出口。",
        extractions=[
            lx.data.Extraction(extraction_class="场所", extraction_text="住宅建筑"),
            lx.data.Extraction(extraction_class="设备", extraction_text="室内消火栓"),
            lx.data.Extraction(extraction_class="设备", extraction_text="高位消防水箱"),
            lx.data.Extraction(extraction_class="指标", extraction_text="18m3"),
            lx.data.Extraction(extraction_class="构件", extraction_text="消防电梯前室"),
            lx.data.Extraction(extraction_class="构件", extraction_text="安全出口"),
        ],
    ),
]


def create_model() -> OpenAILanguageModel:
    if not DEEPSEEK_API_KEY:
        raise ValueError("DEEPSEEK_API_KEY not set in backend/.env")
    return OpenAILanguageModel(
        model_id=MODEL_ID,
        api_key=DEEPSEEK_API_KEY,
        base_url=DEEPSEEK_BASE_URL,
    )


def extract_entities(page_text: str, model: OpenAILanguageModel) -> lx.data.AnnotatedDocument:
    return lx.extract(
        text_or_documents=page_text,
        prompt_description=PROMPT_DESCRIPTION,
        examples=EXAMPLES,
        model=model,
        show_progress=False,
    )