"""
Chinese Legal Element Extractor for Labor Arbitration.
Uses a fine-tuned Chinese RoBERTa multi-task model (or falls back to rules).
"""

from __future__ import annotations

import re
from typing import Any, Optional

# NOTE: must be imported before torch on Windows
import transformers  # noqa: F401

import torch

import os as _os
DEFAULT_MODEL_PATH = _os.environ.get(
    "NLP_MODEL_PATH",
    _os.path.join(_os.path.dirname(__file__), "..", "..", "models", "chinese_roberta_labor_extractor_v2"),
)
# Fallback to absolute path if relative doesn't exist
if not _os.path.isdir(DEFAULT_MODEL_PATH):
    DEFAULT_MODEL_PATH = "D:/Graduation Project/second_type/nlp-service/models/chinese_roberta_labor_extractor_v2"

# Fallback rule-based extraction patterns (used when model is not available)
_CAUSE_PATTERNS = [
    ("生育保险待遇纠纷", ["生育津贴", "生育医疗", "产假工资", "生育保险待遇"]),
    ("赔偿金纠纷", ["违法解除劳动", "违法终止劳动", "违法辞退", "赔偿金", "2N"]),
    ("经济补偿金纠纷", ["经济补偿金", "代通知金", "N+1", "离职补偿"]),
    ("追索劳动报酬", ["追索劳动报酬", "拖欠工资", "工资差额", "加班费", "高温津贴"]),
    ("工伤保险待遇纠纷", ["工伤保险待遇", "一次性伤残", "停工留薪", "工伤认定"]),
    ("劳动关系纠纷类", ["确认劳动关系", "未订立书面劳动合同", "二倍工资", "双倍工资"]),
]


class ChineseLegalElementExtractor:
    """
    Element extraction for labor arbitration case texts.
    Uses the trained Chinese RoBERTa multi-task model when available,
    falling back to rule-based extraction otherwise.
    """

    def __init__(self, model_path: Optional[str] = None):
        self._model = None
        self._tokenizer = None
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        try:
            from .bert_multi_task_model import ChineseRobertaMultiTask
            from transformers import AutoTokenizer

            path = model_path or DEFAULT_MODEL_PATH
            self._model = ChineseRobertaMultiTask.from_pretrained(path)
            self._model.to(self._device)
            self._model.eval()

            model_config_path = f"{path}/model_config.json"
            import json, os
            if os.path.exists(model_config_path):
                with open(model_config_path, encoding="utf-8") as f:
                    config = json.load(f)
                model_name = config.get("model_name", "hfl/chinese-roberta-wwm-ext")
            else:
                model_name = "hfl/chinese-roberta-wwm-ext"

            self._tokenizer = AutoTokenizer.from_pretrained(
                path if os.path.exists(f"{path}/vocab.txt") else model_name
            )
            self._use_model = True
        except Exception as e:
            import traceback
            print(f"[NLP] Model load failed: {e}")
            traceback.print_exc()
            self._use_model = False

    def extract(self, text: str) -> dict[str, Any]:
        if self._use_model and self._model is not None:
            return self._extract_with_model(text)
        return self._extract_with_rules(text)

    @torch.no_grad()
    def _extract_with_model(self, text: str) -> dict[str, Any]:
        # Use rules for entity spans (reliable) + model for classification
        rule_result = self._extract_with_rules(text)

        encoding = self._tokenizer(
            text,
            truncation=True,
            max_length=512,
            padding="max_length",
            return_tensors="pt",
        )
        input_ids = encoding["input_ids"].to(self._device)
        attention_mask = encoding["attention_mask"].to(self._device)

        preds = self._model.predict(input_ids, attention_mask, self._tokenizer)

        cause_classes = [
            "劳动关系纠纷类", "工伤保险待遇纠纷", "追索劳动报酬",
            "经济补偿金纠纷", "赔偿金纠纷", "生育保险待遇纠纷",
        ]
        contract_classes = [
            "无固定期限劳动合同", "固定期限劳动合同",
            "未订立书面劳动合同", "未知",
        ]

        cause_idx = preds["cause_predictions"][0] if preds["cause_predictions"] else 0
        contract_idx = preds["contract_predictions"][0] if preds["contract_predictions"] else 3

        # Decode NER spans as supplement to rules
        ner_spans = self._decode_ner_spans(
            input_ids[0].cpu().tolist(),
            preds["ner_predictions"][0] if preds["ner_predictions"] else [],
        )

        # Merge: model classification + rule-based entities
        result = {
            "parties": {
                "applicant_name": (
                    ner_spans.get("APPLICANT_NAME")
                    or rule_result.get("parties", {}).get("applicant_name")
                ),
                "respondent_name": (
                    ner_spans.get("RESPONDENT_NAME")
                    or rule_result.get("parties", {}).get("respondent_name")
                ),
                "worker_position": (
                    ner_spans.get("WORKER_POSITION")
                    or rule_result.get("parties", {}).get("worker_position")
                ),
                "arbitration_org": (
                    ner_spans.get("ARBITRATION_ORG")
                    or rule_result.get("parties", {}).get("arbitration_org")
                ),
            },
            "case_cause": {"type": cause_classes[cause_idx]},
            "tmpl_primary_cause": cause_classes[cause_idx],
            "facts": {
                "entry_date": (
                    ner_spans.get("ENTRY_DATE")
                    or rule_result.get("facts", {}).get("entry_date")
                ),
                "leave_date": (
                    ner_spans.get("LEAVE_DATE")
                    or rule_result.get("facts", {}).get("leave_date")
                ),
                "filing_date": (
                    ner_spans.get("FILING_DATE")
                    or rule_result.get("facts", {}).get("filing_date")
                ),
                "month_salary": (
                    ner_spans.get("MONTH_SALARY")
                    or rule_result.get("facts", {}).get("month_salary")
                ),
                "work_duration_text": (
                    ner_spans.get("WORK_DURATION")
                    or rule_result.get("facts", {}).get("work_duration_text")
                ),
                "overtime_desc": (
                    ner_spans.get("OVERTIME_DESC")
                    or rule_result.get("facts", {}).get("overtime_desc")
                ),
                "termination_reason": (
                    ner_spans.get("TERMINATION_REASON")
                    or rule_result.get("facts", {}).get("termination_reason")
                ),
            },
            "claims": {
                "amount_total": (
                    ner_spans.get("CLAIM_AMOUNT")
                    or rule_result.get("claims", {}).get("amount_total")
                ),
            },
            "contract_type": contract_classes[contract_idx],
            "law_refs": rule_result.get("law_refs", []),
            "evidence_materials": rule_result.get("evidence_materials", []),
        }
        return result

    def _decode_ner_spans(
        self,
        input_ids: list[int],
        ner_labels: list[int],
    ) -> dict[str, Any]:
        """Decode BIO NER predictions into span text using the tokenizer."""
        if self._tokenizer is None:
            return {}

        # Inline BIO decoding (avoids dependency on training/bio_schema.py)
        _ENTITY_TYPES = [
            "APPLICANT_NAME", "RESPONDENT_NAME", "ENTRY_DATE", "LEAVE_DATE",
            "FILING_DATE", "MONTH_SALARY", "CLAIM_AMOUNT", "WORKER_POSITION",
            "ARBITRATION_ORG", "LAW_REF", "CASE_NUMBER", "EVIDENCE",
            "TERMINATION_REASON", "OVERTIME_DESC", "WORK_DURATION",
        ]
        _id_to_label = {0: "O"}
        _idx = 1
        for et in _ENTITY_TYPES:
            _id_to_label[_idx] = f"B-{et}"
            _id_to_label[_idx + 1] = f"I-{et}"
            _idx += 2

        # Convert token IDs back to tokens
        tokens = self._tokenizer.convert_ids_to_tokens(input_ids)
        tokens = [t for t in tokens if t not in ("[PAD]", "[CLS]", "[SEP]")]

        # Skip [CLS] label
        n_skip = len(input_ids) - len(tokens)
        ner_labels = ner_labels[n_skip:n_skip + len(tokens)]
        ner_labels = ner_labels[:len(tokens)]
        while len(ner_labels) < len(tokens):
            ner_labels.append(0)

        # Decode BIO → entity spans (inline logic)
        spans: list[dict[str, Any]] = []
        current_entity = None
        current_tokens: list[str] = []
        for i, (token, label_id) in enumerate(zip(tokens, ner_labels)):
            label = _id_to_label.get(label_id, "O")
            if label.startswith("B-"):
                if current_entity is not None:
                    spans.append({
                        "entity": current_entity,
                        "text": "".join(current_tokens),
                    })
                current_entity = label[2:]
                current_tokens = [token]
            elif label.startswith("I-") and current_entity == label[2:]:
                current_tokens.append(token)
            else:
                if current_entity is not None:
                    spans.append({
                        "entity": current_entity,
                        "text": "".join(current_tokens),
                    })
                current_entity = None
                current_tokens = []
        if current_entity is not None:
            spans.append({"entity": current_entity, "text": "".join(current_tokens)})

        result: dict[str, list[str]] = {}
        for span in spans:
            text = span["text"].replace("##", "")
            if span["entity"] not in result:
                result[span["entity"]] = []
            result[span["entity"]].append(text)

        return {k: "；".join(v) for k, v in result.items()}

    @staticmethod
    def _extract_with_rules(text: str) -> dict[str, Any]:
        """Fallback: rule-based extraction for when model is not available."""
        if not text or not text.strip():
            return {
                "parties": {},
                "case_cause": {"type": "劳动争议"},
                "facts": {},
                "claims": {},
                "law_refs": [],
            }

        # Detect cause type
        cause_type = "劳动争议"
        for name, keywords in _CAUSE_PATTERNS:
            if any(kw in text for kw in keywords):
                cause_type = name
                break

        # Extract dates
        date_pattern = r"([12][0-9]{3})[年./-]([01]?[0-9])[月./-]([0-3]?[0-9])[日]?"
        entry_match = re.search(r"(入职|到岗).{0,10}" + date_pattern, text)
        leave_match = re.search(r"(离职|解除|终止).{0,10}" + date_pattern, text)

        # Extract amounts
        amount_match = re.search(r"([0-9]{3,8}(?:\.[0-9]{1,2})?)\s*元", text)
        salary_match = re.search(r"(月工资|月薪).{0,5}([0-9]{3,6})\s*元", text)

        # Extract legal references
        law_refs = re.findall(r"《[^》]{2,30}》", text)

        # Extract entity names
        applicant = re.search(r"申请人[:：]\s*([^\n，,。；]{2,20})", text)
        respondent = re.search(r"被申请人[:：]\s*([^\n，,。；]{2,40})", text)
        position = re.search(r"(岗位|职位|工种)[:：]?\s*([^\n，,。；]{2,20})", text)

        return {
            "parties": {
                "applicant_name": applicant.group(1).strip() if applicant else None,
                "respondent_name": respondent.group(1).strip() if respondent else None,
                "worker_position": position.group(2).strip() if position else None,
            },
            "case_cause": {"type": cause_type},
            "facts": {
                "entry_date": _format_date(entry_match) if entry_match else None,
                "leave_date": _format_date(leave_match) if leave_match else None,
                "month_salary": float(salary_match.group(2)) if salary_match else None,
                "overtime_desc": "含加班" if "加班" in text else None,
                "termination_reason": None,
            },
            "claims": {
                "amount_total": float(amount_match.group(1)) if amount_match else None,
            },
            "law_refs": list(set(law_refs)) if law_refs else [],
        }


def _format_date(match: re.Match) -> str | None:
    try:
        y, m, d = int(match.group(1)), int(match.group(2)), int(match.group(3))
        return f"{y:04d}-{m:02d}-{d:02d}"
    except (IndexError, ValueError):
        return None