from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from threading import RLock
from typing import Any, Dict, List, Optional
from uuid import uuid4

from utils.http_util import upload_file
from utils.doc_util import DocBase
from core.config import META_KEY, FILE_SUFFIX, use_lufa

logger = logging.getLogger(__name__)


_ALLOWED_RISK_LEVELS = {"H", "M", "L", "H,M", ""}
FINDING_KEY_REVIEW = "review"
FINDING_KEY_REFLECT = "reflect"
FINDING_KEY_MERGE = "merge"

_DEFAULT_REVIEW_KEY = FINDING_KEY_REVIEW
_FINDING_KEY_SHEET_NAMES = {
    FINDING_KEY_REVIEW: "审查结果",
    FINDING_KEY_REFLECT: "复核结果",
    FINDING_KEY_MERGE: "合并结果",
}


@dataclass
class Finding:
    rule_title: str
    segment_id: int
    original_text: str
    issue: str
    risk_level: str
    suggestion: str
    id: str = ""
    result: str = ""

    def __post_init__(self) -> None:
        level = (self.risk_level or "").upper()
        if level not in _ALLOWED_RISK_LEVELS:
            raise ValueError(
                f"risk_level must be one of {_ALLOWED_RISK_LEVELS}, got {self.risk_level}"
            )
        self.risk_level = level

    @classmethod
    def from_dict(cls, data: Dict) -> "Finding":
        data = data or {}
        return cls(
            id=str(data.get("id", "")),
            rule_title=str(data.get("rule_title", "")),
            segment_id=int(data.get("segment_id", 0) or 0),
            original_text=str(data.get("original_text", "")),
            issue=str(data.get("issue", "")),
            risk_level=str(data.get("risk_level", "")),
            suggestion=str(data.get("suggestion", "")),
            result=str(data.get("result", "")),
        )

    def __repr__(self):
        return (
            f"Finding(id={self.id!r}, rule_title={self.rule_title!r}, segment_id={self.segment_id}, "
            f"issue={self.issue!r}, risk_level={self.risk_level!r}, result={self.result!r})"
        )


@dataclass
class MemoryStore:
    """简化的记忆存储：合同事实 facts 与问题 findings。线程安全并支持 JSON 持久化。"""

    storage_name: Optional[Path] = "default.json"

    def __init__(self, storage_name: str = "default.json") -> None:
        self._storage_path = Path(__file__).resolve().parent.parent / "tmp" / storage_name  # type: ignore[arg-type]
        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = RLock()
        self.facts: List[Dict[str, Any]] = []
        self.merge_facts: List[Dict[str, Any]] = []
        self.findings: Dict[str, List[Finding]] = {}
        self._load()

    # ---------------------- facts ----------------------
    def set_facts(self, facts: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        with self._lock:
            self.facts = facts or []
            self._persist()
            return self.facts

    def add_facts(self, partial: Dict[str, Any]) -> List[Dict[str, Any]]:
        with self._lock:
            self.facts.append(partial)
            self._persist()
            return self.facts

    def get_facts(self) -> List[Dict[str, Any]]:
        with self._lock:
            return self.facts  # deep copy

    def set_merge_facts(
        self, merge_facts: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        with self._lock:
            self.merge_facts = merge_facts or []
            self._persist()
            return self.merge_facts

    def add_merge_facts(self, partial: Dict[str, Any]) -> List[Dict[str, Any]]:
        with self._lock:
            self.merge_facts.append(partial)
            self._persist()
            return self.merge_facts

    def get_merge_facts(self) -> List[Dict[str, Any]]:
        with self._lock:
            return self.merge_facts

    def search_facts(self, keywords: List[str]) -> List[Any]:
        keys = [str(k).strip().lower() for k in (keywords or []) if str(k).strip()]
        if not keys:
            return []

        def _key_match(name: Any) -> bool:
            key_name = str(name).strip().lower()
            return bool(key_name) and any(k in key_name or key_name in k for k in keys)

        matched_values: List[Any] = []
        with self._lock:
            all_facts = list(self.facts)

        for item in all_facts:
            if not isinstance(item, dict):
                continue

            for top_key, top_value in item.items():
                if _key_match(top_key):
                    matched_values.append(
                        {
                            top_key: top_value,
                            META_KEY: item.get(
                                META_KEY, {}
                            ),  # include metadata if exists
                        }
                    )

        return matched_values

    # -------------------- findings ---------------------
    def add_finding(self, key: str, finding: Finding) -> Finding:
        return self._add_finding(key, finding)

    def list_findings(self, key: str) -> List[Finding]:
        return self._list_findings(self._get_findings_bucket(key))

    def get_findings_by_segment(self, key: str, segment_id: int) -> List[Finding]:
        return self._get_findings_by_segment(self._get_findings_bucket(key), segment_id)

    def delete_findings_by_segment(self, key: str, segment_id: int) -> int:
        return self._delete_findings_by_segment(key, segment_id)

    def search_findings(
        self,
        key: str,
        keyword: str,
        rule_title: Optional[str] = None,
        risk_level: Optional[str] = None,
    ) -> List[Finding]:
        return self._search_findings(
            self._get_findings_bucket(key), keyword, rule_title, risk_level
        )

    def list_findings_grouped(self) -> Dict[str, List[Finding]]:
        with self._lock:
            return {k: list(v) for k, v in self.findings.items()}

    def _add_finding(self, key: str, finding: Finding) -> Finding:
        with self._lock:
            finding_key = self._normalize_finding_key(key)
            if not finding.id:
                finding.id = uuid4().hex
            bucket = self.findings.setdefault(finding_key, [])
            bucket.append(finding)
            self._persist()
            return finding

    def _get_findings_bucket(self, key: str) -> List[Finding]:
        finding_key = self._normalize_finding_key(key)
        return self.findings.setdefault(finding_key, [])

    def _list_findings(self, target: List[Finding]) -> List[Finding]:
        with self._lock:
            return list(target)

    def _get_findings_by_segment(
        self, target: List[Finding], segment_id: int
    ) -> List[Finding]:
        with self._lock:
            return [f for f in target if f.segment_id == segment_id]

    def _delete_findings_by_segment(self, key: str, segment_id: int) -> int:
        with self._lock:
            current = self._get_findings_bucket(key)
            before = len(current)
            updated = [f for f in current if f.segment_id != segment_id]
            self.findings[key] = updated
            removed = before - len(updated)
            if removed:
                self._persist()
            return removed

    def _search_findings(
        self,
        target: List[Finding],
        keyword: str,
        rule_title: Optional[str] = None,
        risk_level: Optional[str] = None,
    ) -> List[Finding]:
        key = (keyword or "").strip().lower()
        with self._lock:
            candidates = list(target)
            if rule_title:
                candidates = [
                    f
                    for f in candidates
                    if (f.rule_title or "").lower() == rule_title.strip().lower()
                ]
            if risk_level:
                lvl = risk_level.strip().upper()
                candidates = [f for f in candidates if f.risk_level == lvl]
            if not key:
                return candidates

            def _matches(f: Finding) -> bool:
                hay = " ".join(
                    [
                        f.rule_title,
                        f.original_text,
                        f.issue,
                        f.suggestion,
                        f.result,
                    ]
                ).lower()
                return key in hay

            return [f for f in candidates if _matches(f)]

    # ------------------- housekeeping ------------------
    def clear(self) -> None:
        with self._lock:
            self.facts.clear()
            self.merge_facts.clear()
            self.findings.clear()
            self._persist()

    def _persist(self) -> None:
        payload = {
            "facts": self.facts,
            "merge_facts": self.merge_facts,
            "findings": {
                key: [asdict(f) for f in values]
                for key, values in self.findings.items()
            },
        }
        try:
            self._storage_path.write_text(
                json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
            )
        except Exception as exc:
            logger.error("Failed to persist memory store: %s", exc)

    def _load(self) -> None:
        try:
            if not self._storage_path.exists():
                return
            raw = self._storage_path.read_text(encoding="utf-8")
            data = json.loads(raw or "{}")
            if isinstance(data, dict):
                self.facts = data.get("facts") or []
                self.merge_facts = data.get("merge_facts") or []
                loaded_findings = data.get("findings", {})

                findings_map: Dict[str, List[Finding]] = {}
                if isinstance(loaded_findings, dict):
                    for key, items in loaded_findings.items():
                        normalized_key = self._normalize_finding_key(str(key))
                        findings_map[normalized_key] = [
                            Finding.from_dict(item) for item in (items or [])
                        ]

                self.findings = findings_map
                needs_persist = False
                for bucket in self.findings.values():
                    for finding in bucket:
                        if not finding.id:
                            finding.id = uuid4().hex
                            needs_persist = True
                if needs_persist:
                    self._persist()
        except Exception as exc:
            logger.error("Failed to load memory store: %s", exc)

    def export_to_excel(self, file_name: Optional[str] = None) -> Dict[str, Any]:
        """Export findings and facts to Excel, upload, then delete the local file."""
        try:
            from openpyxl import Workbook  # type: ignore
        except ImportError as exc:
            raise ImportError(
                "openpyxl is required for export_to_excel; install via 'pip install openpyxl'"
            ) from exc
        file_suffix = FILE_SUFFIX + datetime.now().strftime("%Y%m%d_%H%M%S")
        name = file_name or f"memory_export.xlsx"
        name = Path(name).stem + file_suffix + ".xlsx"
        # print(f"Exporting to Excel with file name: {name}")
        output_path = Path(__file__).resolve().parent.parent / "tmp" / name

        with self._lock:
            wb = Workbook()
            finding_headers = [
                ("id", "ID"),
                ("rule_title", "规则标题"),
                ("segment_id", "分段ID"),
                ("original_text", "原文"),
                ("issue", "问题描述"),
                ("risk_level", "风险等级"),
                ("result", "合格性"),
                ("suggestion", "建议"),
            ]

            grouped_items = list(self.findings.items())
            if grouped_items:
                first_key, first_values = grouped_items[0]
                ws_first = wb.active
                first_sheet_name = _FINDING_KEY_SHEET_NAMES.get(
                    self._normalize_finding_key(first_key), first_key
                )
                ws_first.title = self._safe_sheet_name(first_sheet_name)
                ws_first.append([label for _, label in finding_headers])
                for f in first_values:
                    ws_first.append([getattr(f, key, "") for key, _ in finding_headers])

                for key, values in grouped_items[1:]:
                    sheet_name = _FINDING_KEY_SHEET_NAMES.get(
                        self._normalize_finding_key(key), key
                    )
                    ws = wb.create_sheet(self._safe_sheet_name(sheet_name))
                    ws.append([label for _, label in finding_headers])
                    for f in values:
                        ws.append(
                            [
                                getattr(f, item_key, "")
                                for item_key, _ in finding_headers
                            ]
                        )
            else:
                ws_empty = wb.active
                ws_empty.title = self._safe_sheet_name(
                    _FINDING_KEY_SHEET_NAMES.get(
                        _DEFAULT_REVIEW_KEY, _DEFAULT_REVIEW_KEY
                    )
                )
                ws_empty.append([label for _, label in finding_headers])

            ws_facts = wb.create_sheet("合同事实")
            if self.facts:
                ws_facts.append(["元信息", "事实内容"])
                for item in self.facts:
                    if not isinstance(item, dict):
                        ws_facts.append(["事实", self._format_summary_for_export(item)])
                        continue
                    meta_info = item.get(META_KEY, None)
                    ws_facts.append(
                        [
                            json.dumps(meta_info, ensure_ascii=False),
                            self._format_summary_for_export(item),
                        ]
                    )
            else:
                ws_facts.append(["元信息", "事实内容"])

            wb.save(output_path)

        try:
            res = upload_file(str(output_path))
        finally:
            try:
                output_path.unlink()
            except Exception:
                logger.warning("Failed to delete temp excel: %s", output_path)

        return res

    def export_facts_to_excel(self, file_name: Optional[str] = None) -> Dict[str, Any]:
        """Export raw facts and merged facts to Excel, upload, then delete the local file."""
        try:
            from openpyxl import Workbook  # type: ignore
        except ImportError as exc:
            raise ImportError(
                "openpyxl is required for export_facts_to_excel; install via 'pip install openpyxl'"
            ) from exc

        file_suffix = FILE_SUFFIX + datetime.now().strftime("%Y%m%d_%H%M%S")
        name = file_name or "facts_export.xlsx"
        name = Path(name).stem + file_suffix + ".xlsx"
        output_path = Path(__file__).resolve().parent.parent / "tmp" / name

        with self._lock:
            wb = Workbook()
            ws_facts = wb.active
            ws_facts.title = self._safe_sheet_name("合同事实")
            self._append_facts_sheet(ws_facts, self.facts)

            ws_merge_facts = wb.create_sheet(self._safe_sheet_name("合并事实"))
            self._append_merge_facts_sheet(ws_merge_facts, self.merge_facts)

            wb.save(output_path)

        try:
            res = upload_file(str(output_path))
        finally:
            try:
                output_path.unlink()
            except Exception:
                logger.warning("Failed to delete temp facts excel: %s", output_path)

        return res

    def export_findings_to_doc_comments(
        self,
        doc_obj: DocBase,
        file_name: Optional[str] = None,
        remove_prefix: bool = False,
        finding_key: str = _DEFAULT_REVIEW_KEY,
    ) -> Dict[str, Any]:
        """Add all findings as comments to a document, upload, then delete the local file."""
        if doc_obj is None:
            raise ValueError("doc_obj is required")
        # build suffix
        file_suffix = FILE_SUFFIX + datetime.now().strftime("%Y%m%d_%H%M%S")
        # derive file name
        name = file_name or getattr(doc_obj, "_doc_name", "") or "memory_export.docx"
        name = Path(name).stem + file_suffix + (Path(name).suffix or ".docx")
        output_path = Path(__file__).resolve().parent.parent / "tmp" / name

        target_key = self._normalize_finding_key(finding_key)

        with self._lock:
            target_findings = list(self._get_findings_bucket(target_key))
            comments: List[Dict[str, Any]] = []
            for idx, f in enumerate(target_findings, start=1):
                segment_id = int(f.segment_id or 0)
                chunk_id = max(segment_id, 0)
                if use_lufa:
                    suggest_parts = []
                    if f.risk_level:
                        suggest_parts.append(f"风险等级：{f.risk_level}")
                    if f.issue:
                        suggest_parts.append(f"问题：{f.issue}")
                    if f.suggestion:
                        suggest_parts.append(f"建议：{f.suggestion}")
                    suggest_text = "\n".join(suggest_parts).strip()
                else:
                    suggest_text = f"建议：{f.suggestion}".strip()
                comments.append(
                    {
                        "id": str(idx),
                        "key_points": f.rule_title or "风险提示",
                        "original_text": f.original_text or "",
                        "details": f.issue or "",
                        "chunk_id": chunk_id,
                        "result": f.result or "不合格",
                        "suggest": suggest_text,
                    }
                )

            if comments:
                doc_obj.add_chunk_comment(0, comments)

            doc_obj.to_file(str(output_path), remove_prefix=remove_prefix)

        try:
            res = upload_file(str(output_path))
        finally:
            try:
                output_path.unlink()
            except Exception:
                logger.warning("Failed to delete temp doc: %s", output_path)

        return res

    @staticmethod
    def _safe_sheet_name(name: str) -> str:
        # Excel sheet names cannot exceed 31 chars or include certain symbols.
        safe = (name or _DEFAULT_REVIEW_KEY).strip() or _DEFAULT_REVIEW_KEY
        for ch in [":", "\\", "/", "?", "*", "[", "]"]:
            safe = safe.replace(ch, "_")
        return safe[:31]

    def _append_facts_sheet(self, ws: Any, facts: List[Dict[str, Any]]) -> None:
        ws.append(["元信息", "事实内容"])
        for item in facts or []:
            if not isinstance(item, dict):
                ws.append(["事实", self._format_summary_for_export(item)])
                continue
            meta_info = item.get(META_KEY, None)
            ws.append(
                [
                    json.dumps(meta_info, ensure_ascii=False),
                    self._format_summary_for_export(item),
                ]
            )

    def _append_merge_facts_sheet(
        self, ws: Any, merge_facts: List[Dict[str, Any]]
    ) -> None:
        ws.append(["提取项", "提取内容"])
        rows: Dict[str, Any] = {}
        for item in merge_facts or []:
            if not isinstance(item, dict):
                continue
            for summary_name, value in item.items():
                if summary_name == META_KEY:
                    continue
                summary_text = str(summary_name or "").strip()
                if not summary_text:
                    continue
                rows[summary_text] = value

        for summary_name, value in rows.items():
            ws.append([summary_name, self._format_summary_for_export(value)])

    @staticmethod
    def _format_summary_for_export(value: Any, level: int = 0) -> str:
        indent = "  " * level
        if isinstance(value, dict):
            lines: List[str] = []
            for key, child in value.items():
                if key == META_KEY:
                    continue
                key_text = str(key)
                if isinstance(child, (dict, list)):
                    lines.append(f"{indent}{key_text}：")
                    child_text = MemoryStore._format_summary_for_export(child, level + 1)
                    if child_text:
                        lines.append(child_text)
                else:
                    lines.append(f"{indent}{key_text}：{MemoryStore._format_scalar(child)}")
            return "\n".join(lines)

        if isinstance(value, list):
            lines = []
            for item in value:
                if isinstance(item, (dict, list)):
                    item_text = MemoryStore._format_summary_for_export(item, level)
                    if item_text:
                        lines.append(item_text)
                else:
                    lines.append(f"{indent}{MemoryStore._format_scalar(item)}")
            return "\n".join(lines)

        return f"{indent}{MemoryStore._format_scalar(value)}"

    @staticmethod
    def _format_scalar(value: Any) -> str:
        if value is None:
            return ""
        return str(value)

    @staticmethod
    def _normalize_finding_key(key: str) -> str:
        normalized = (key or "").strip().lower()
        if not normalized:
            return _DEFAULT_REVIEW_KEY
        return normalized


def test_export_findings_to_doc_comments(doc_path: str) -> None:

    store = MemoryStore()
    finding = Finding(
        rule_title="违约责任",
        segment_id=1,
        original_text="湖南麓谷发展集团有限公司",
        issue="未约定违约金上限，可能导致赔偿范围过大",
        risk_level="H",
        suggestion="建议增加‘赔偿金额不超过合同总额的30%’",
        result="不合格",
    )
    store.add_finding(FINDING_KEY_REFLECT, finding)
    """测试：将 findings 作为批注写入文档并上传。"""
    if not doc_path:
        print("doc_path 为空，跳过批注导出测试")
        return
    if not Path(doc_path).exists():
        print(f"文件不存在，跳过批注导出测试: {doc_path}")
        return
    try:
        from utils.spire_word_util import SpireWordDoc
    except Exception as exc:
        print(f"加载 SpireWordDoc 失败，跳过批注导出测试: {exc}")
        return
    doc = SpireWordDoc()
    doc.load(doc_path)
    res = store.export_findings_to_doc_comments(doc)
    print("Export doc comments:")
    print(json.dumps(res, ensure_ascii=False, indent=2))


def test_memory_and_export_excel():
    # 简单示例：设置事实 -> 写入问题 -> 读取/搜索
    store = MemoryStore()
    store.add_facts(
        {
            "公司": {"甲方": "A 公司", "乙方": "B 公司"},
            "支付": {"方式": "银行转账", "期限": "验收后30日内"},
            META_KEY: {"segment_id": 1},
        }
    )
    # print( store.search_facts(['支付']))
    finding1 = Finding(
        rule_title="违约责任",
        segment_id=1,
        original_text="违约方应赔偿全部损失",
        issue="未约定违约金上限，可能导致赔偿范围过大",
        risk_level="H",
        suggestion="建议增加‘赔偿金额不超过合同总额的30%’",
    )
    finding2 = Finding(
        rule_title="违约责任",
        segment_id=2,
        original_text="违约方应赔偿全部损失",
        issue="未约定违约金上限，可能导致赔偿范围过大",
        risk_level="H",
        suggestion="建议增加‘赔偿金额不超过合同总额的30%’",
    )
    store.add_finding(FINDING_KEY_REVIEW, finding1)
    store.add_finding(FINDING_KEY_REFLECT, finding2)
    print(store.get_findings_by_segment(FINDING_KEY_REVIEW, 1))

    # print("Facts:\n" + json.dumps(store.get_facts(), ensure_ascii=False, indent=2))
    # hits = store.search_findings("赔偿", rule_title="违约责任")
    # print("Findings search:")
    # for f in hits:
    #   print(json.dumps(asdict(f), ensure_ascii=False, indent=2))
    print(store.export_to_excel("测试"))


if __name__ == "__main__":
    test_export_findings_to_doc_comments(
        "/home/ccran/lufa-contract/tmp/1_金盘箱变采购合同.docx"
    )
    # test_memory_and_export_excel()
