"""Host-side Materials Project tools.

MP_API_KEY stays on the host. The sandbox can use offline pymatgen on files
written to task_dir, but it must not receive the Materials Project key.
"""
from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
from typing import Any, Optional

from .base import Tool

try:  # patched in tests; missing dependency should produce a clean tool error.
    from mp_api.client import MPRester  # type: ignore
except Exception:  # pragma: no cover - exercised when mp-api is not installed
    MPRester = None  # type: ignore


_DEFAULT_SUMMARY_FIELDS = [
    "material_id",
    "formula_pretty",
    "symmetry",
    "energy_above_hull",
]


def _mp_key() -> str:
    key = os.environ.get("MP_API_KEY", "").strip()
    if not key:
        raise RuntimeError("MP_API_KEY env 未设置,无法查询 Materials Project")
    return key


def _to_plain(obj: Any) -> Any:
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, (list, tuple)):
        return [_to_plain(x) for x in obj]
    if isinstance(obj, dict):
        return {str(k): _to_plain(v) for k, v in obj.items()}
    if hasattr(obj, "model_dump"):
        return _to_plain(obj.model_dump())
    if hasattr(obj, "as_dict"):
        return _to_plain(obj.as_dict())
    # mp-api documents expose fields as attributes.
    out: dict[str, Any] = {}
    for name in _DEFAULT_SUMMARY_FIELDS:
        if hasattr(obj, name):
            out[name] = _to_plain(getattr(obj, name))
    return out or str(obj)


def _mpr():
    if MPRester is None:
        raise RuntimeError("mp-api 未安装,请在宿主环境安装 mp-api 后再启用 MP tool")
    return MPRester(_mp_key())


class MaterialsProjectSearchSummaryTool(Tool):
    name = "mp_search_summary"
    description = (
        "Search Materials Project summary data using the host MP_API_KEY. "
        "Single query: pass formula / material_ids / elements → returns trimmed JSON. "
        "Batch: pass `formulas` (a list) to look up MANY formulae in ONE tool call — "
        "full results are written to task_dir/materials/*.json and only a compact "
        "per-formula digest is returned. Always prefer batch over calling this tool "
        "once per formula. Use mp_get_structure to save a CIF for offline pymatgen analysis."
    )
    parameters = {
        "type": "object",
        "properties": {
            "formula": {"type": "string", "description": "Optional single formula such as Ca3SiO5."},
            "formulas": {"type": "array", "items": {"type": "string"}, "description": "Optional list of formulae for a BATCH lookup, e.g. ['Ca3SiO5','Ca2SiO4',...]. One tool call instead of one per formula; full results saved to a JSON file, only a digest returned. Max 300 per call — split larger lists. When you have >1 formula to look up, use THIS, not repeated formula= calls."},
            "material_ids": {"type": "array", "items": {"type": "string"}, "description": "Optional Materials Project ids such as mp-123."},
            "elements": {"type": "array", "items": {"type": "string"}, "description": "Optional element symbols for a chemical system search."},
            "fields": {"type": "array", "items": {"type": "string"}, "description": "Fields to return; defaults to material_id/formula_pretty/symmetry/energy_above_hull. Do NOT put the search formula here — use formula_pretty (not 'formula'). Common fields: formula_pretty, symmetry, energy_above_hull, band_gap, density, volume, is_stable, structure."},
            "limit": {"type": "integer", "default": 10, "description": "Maximum records returned, 1-50. In batch mode this is PER formula (keep small — 1-3 gives the most-stable phases). Server-side bounded to stay within MP fair-use."},
        },
    }

    _MAX_BATCH = 300

    def __init__(
        self,
        *,
        working_dir: Optional[Path] = None,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
    ) -> None:
        super().__init__(base_dir=base_dir, user_root=user_root)
        self.working_dir = Path(working_dir) if working_dir is not None else None

    def execute(
        self,
        formula: str = "",
        formulas: Optional[list[str]] = None,
        material_ids: Optional[list[str]] = None,
        elements: Optional[list[str]] = None,
        fields: Optional[list[str]] = None,
        limit: int = 10,
    ) -> str:
        limit = min(max(int(limit), 1), 50)
        chosen_fields = fields or _DEFAULT_SUMMARY_FIELDS
        batch = [str(f).strip() for f in (formulas or []) if str(f).strip()]
        if batch:
            return self._batch(batch, chosen_fields, limit)
        kwargs: dict[str, Any] = {"fields": chosen_fields}
        if formula:
            kwargs["formula"] = formula
        if material_ids:
            kwargs["material_ids"] = material_ids
        if elements:
            kwargs["elements"] = elements
        if len(kwargs) == 1:
            return "[Error] formula / formulas / material_ids / elements 至少传一个"
        try:
            with _mpr() as mpr:
                # num_chunks=1 + chunk_size=limit 让服务端单次只回 limit 条、不翻页。
                # 否则 mp-api 默认 chunk_size=1000 且自动拉完所有页 —— 整库级下载,
                # 会被 MP 判定为 abusive traffic 并封 IP/ASN。
                docs = mpr.materials.summary.search(
                    num_chunks=1, chunk_size=limit, **kwargs
                )
        except Exception as e:
            return f"[Error] mp_search_summary failed: {type(e).__name__}: {e}"
        plain = [_to_plain(d) for d in list(docs)[:limit]]
        return json.dumps(plain, ensure_ascii=False, indent=2)

    @staticmethod
    def _hull(rec: Any) -> float:
        v = rec.get("energy_above_hull") if isinstance(rec, dict) else None
        return v if isinstance(v, (int, float)) else float("inf")

    def _batch(self, formulas: list[str], fields: list[str], limit: int) -> str:
        if len(formulas) > self._MAX_BATCH:
            return (
                f"[Error] 一次最多 {self._MAX_BATCH} 个化学式(收到 {len(formulas)});请分批调用。"
            )
        seen: set[str] = set()
        uniq = [f for f in formulas if not (f in seen or seen.add(f))]
        try:
            session = _mpr()
        except Exception as e:
            return f"[Error] mp_search_summary batch failed: {type(e).__name__}: {e}"
        agg: list[dict[str, Any]] = []
        n_ok = 0
        # 单个 formula 出错不连坐整批;复用一个 MPRester 会话避免逐条重认证。
        with session as mpr:
            for f in uniq:
                try:
                    docs = mpr.materials.summary.search(
                        formula=f, fields=fields, num_chunks=1, chunk_size=limit
                    )
                    results = [_to_plain(d) for d in list(docs)[:limit]]
                    results.sort(key=self._hull)  # 稳定相优先
                    if results:
                        n_ok += 1
                    agg.append({"formula": f, "n": len(results), "results": results})
                except Exception as e:
                    agg.append({"formula": f, "error": f"{type(e).__name__}: {e}"})
        return self._render_batch(agg, fields, n_ok)

    def _render_batch(self, agg: list[dict[str, Any]], fields: list[str], n_ok: int) -> str:
        empties = [a["formula"] for a in agg if "error" in a or a.get("n", 0) == 0]
        header = (
            f"mp_search_summary batch: {len(agg)} 个化学式 → {n_ok} 个有结果,"
            f"{len(empties)} 个空/错误。\n"
            f"每条记录字段: {', '.join(fields)}\n"
        )
        if self.working_dir is not None:
            h = hashlib.sha1(
                ",".join(a["formula"] for a in agg).encode("utf-8")
            ).hexdigest()[:8]
            dest = self.working_dir / "materials" / f"mp_search_batch_{h}.json"
            try:
                dest.parent.mkdir(parents=True, exist_ok=True)
                dest.write_text(
                    json.dumps(agg, ensure_ascii=False, indent=2), encoding="utf-8"
                )
                body = (
                    f"完整结果已存: {self._display(dest)}\n"
                    "  结构 [{\"formula\",\"n\",\"results\":[...]}...];"
                    "用 read 或 run_python/pandas 直接读该文件处理,勿逐条 dump 进对话。\n"
                )
            except Exception as e:
                body = (
                    f"[warn] 写文件失败({type(e).__name__}: {e}),改为内联返回完整 JSON:\n"
                    + json.dumps(agg, ensure_ascii=False, indent=2)
                    + "\n"
                )
        else:
            body = json.dumps(agg, ensure_ascii=False, indent=2) + "\n"
        tail = ""
        if empties:
            shown = ", ".join(empties[:20])
            more = f" …(+{len(empties) - 20})" if len(empties) > 20 else ""
            tail = f"空/错误化学式: {shown}{more}\n"
        return header + body + tail


class MaterialsProjectGetStructureTool(Tool):
    name = "mp_get_structure"
    description = (
        "Download a Materials Project structure by material_id and save it as CIF in task_dir/materials/."
    )
    parameters = {
        "type": "object",
        "properties": {
            "material_id": {"type": "string", "description": "Materials Project id, e.g. mp-123."},
            "filename": {"type": "string", "description": "Optional CIF filename. Defaults to <material_id>.cif."},
        },
        "required": ["material_id"],
    }

    def __init__(
        self,
        *,
        working_dir: Path,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
    ) -> None:
        super().__init__(base_dir=base_dir, user_root=user_root)
        self.working_dir = Path(working_dir)

    def execute(self, material_id: str, filename: str = "") -> str:
        material_id = (material_id or "").strip()
        if not material_id:
            return "[Error] material_id 不可为空"
        safe_name = (filename or f"{material_id}.cif").replace("/", "_").replace("\\", "_").replace("..", "_")
        if not safe_name.lower().endswith(".cif"):
            safe_name += ".cif"
        dest = self.working_dir / "materials" / safe_name
        try:
            with _mpr() as mpr:
                struct = mpr.get_structure_by_material_id(material_id)
            dest.parent.mkdir(parents=True, exist_ok=True)
            struct.to(filename=str(dest))
        except Exception as e:
            return f"[Error] mp_get_structure failed: {type(e).__name__}: {e}"
        return f"saved: {self._display(dest)}"


class MaterialsProjectGetEntriesTool(Tool):
    name = "mp_get_entries"
    description = (
        "Fetch Materials Project computed entries for a chemical system and save trimmed JSON to task_dir/materials/. "
        "Downloads the FULL chemical system (all sub-systems) — volume grows fast with element count; call sparingly and reuse the saved file rather than re-querying."
    )
    parameters = {
        "type": "object",
        "properties": {
            "elements": {"type": "array", "items": {"type": "string"}, "description": "Chemical system elements, e.g. ['Ca','Si','O','H']. More elements = much heavier download."},
            "filename": {"type": "string", "description": "Optional JSON filename. Defaults to mp_entries_<chemsys>.json."},
            "limit": {"type": "integer", "default": 200, "description": "Max entries SAVED to disk, 1-1000. NOTE: this only trims the saved JSON — the full chemsys is still fetched from MP. It does not reduce network traffic."},
        },
        "required": ["elements"],
    }

    def __init__(
        self,
        *,
        working_dir: Path,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
    ) -> None:
        super().__init__(base_dir=base_dir, user_root=user_root)
        self.working_dir = Path(working_dir)

    def execute(
        self,
        elements: list[str],
        filename: str = "",
        limit: int = 200,
    ) -> str:
        elems = [e.strip() for e in (elements or []) if str(e).strip()]
        if not elems:
            return "[Error] elements 不可为空"
        limit = min(max(int(limit), 1), 1000)
        chemsys = "-".join(elems)
        safe_name = filename or f"mp_entries_{chemsys}.json"
        safe_name = safe_name.replace("/", "_").replace("\\", "_").replace("..", "_")
        if not safe_name.lower().endswith(".json"):
            safe_name += ".json"
        dest = self.working_dir / "materials" / safe_name
        try:
            with _mpr() as mpr:
                entries = mpr.get_entries_in_chemsys(elems)
            payload = [_to_plain(e) for e in list(entries)[:limit]]
            dest.parent.mkdir(parents=True, exist_ok=True)
            dest.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        except Exception as e:
            return f"[Error] mp_get_entries failed: {type(e).__name__}: {e}"
        return f"saved: {self._display(dest)}"