import argparse
import base64
import csv as _csv
import json
import os
import re
import requests
from datetime import datetime, timezone
from typing import Optional, List, Dict, Any
from zoneinfo import ZoneInfo

TAIPEI_TZ = ZoneInfo("Asia/Taipei")

DEFAULT_URL = "http://192.168.42.124:13000/function"
DEFAULT_TOKEN = "6R0W53R135510"

CSV_FIELDS = [
    "id",
    "name",
    "gender",
    "remaining",
    "expire_time",
    "expire_at",
    "latitude",
    "longitude",
    "iv_pct",
    "iv_atk",
    "iv_def",
    "iv_sta",
    "cp",
    "level",
    "is_perfect",
    "fast_move",
    "charge_move",
]


class PokemonRadarClient:
    """
    調用 twpk_radar.mjs 的封裝客戶端。
    - 注入座標/截圖旗標
    - 注入 state（cookies/localStorage/UA/langs）
    - 呼叫 Function Server
    - 套用篩選、輸出 JSON/CSV/截圖
    """

    def __init__(
        self,
        url: str = DEFAULT_URL,
        token: str = DEFAULT_TOKEN,
        script: str = "twpk_radar.mjs",
    ):
        self.url = url
        self.token = token
        self.script = script

    # ---------- 準備 JS 代碼 ----------
    def _load_script(self) -> str:
        path = self.script
        if not os.path.isfile(path):
            fallback = "fn.mjs"
            if os.path.isfile(fallback):
                print("ℹ️  {} 不存在，改用 {}".format(path, fallback))
                path = fallback
            else:
                raise FileNotFoundError(
                    "找不到 JS 腳本：{}（也找不到 fn.mjs）".format(self.script)
                )
        with open(path, "r", encoding="utf-8") as f:
            return f.read()

    def _inject_target_and_shot(
        self, code: str, lat: float, lng: float, zoom: int, shot: bool
    ) -> str:
        code = re.sub(
            r"const TARGET\s*=\s*\{[^}]+\};",
            "const TARGET = {{ lat: {}, lng: {}, zoom: {} }};".format(lat, lng, zoom),
            code,
        )
        code = re.sub(
            r"const DO_SCREENSHOT\s*=\s*(true|false);",
            "const DO_SCREENSHOT = {};".format("true" if shot else "false"),
            code,
        )
        return code

    def _inject_state(
        self, code: str, state_path: Optional[str], ua_from_state: bool
    ) -> str:
        payload: Dict[str, Any] = {
            "cookies": [],
            "localStorage": {},
            "ua": None,
            "langs": None,
        }
        if state_path:
            try:
                with open(state_path, "r", encoding="utf-8") as sf:
                    raw = json.load(sf)
                payload["cookies"] = raw.get("cookies", raw.get("Cookies", []))
                payload["localStorage"] = raw.get(
                    "localStorage", raw.get("LocalStorage", {})
                )
                if ua_from_state:
                    payload["ua"] = raw.get("ua")
                    payload["langs"] = raw.get("langs")
                print(
                    "[{}] 🧩 Loaded state: cookies={}, ls={}, ua={}".format(
                        datetime.now().astimezone().isoformat(timespec="seconds"),
                        len(payload["cookies"]),
                        len(payload["localStorage"]),
                        "Y" if payload["ua"] else "N",
                    )
                )
            except Exception as e:
                print("⚠️ load state error: {}".format(e))
        inject = (
            "globalThis.__POKE_STATE__ = "
            + json.dumps(payload, ensure_ascii=False)
            + ";\n"
        )
        return inject + code

    # ---------- 呼叫 function server ----------
    def _call_server(self, code: str) -> Dict[str, Any]:
        resp = requests.post(
            self.url,
            params={"token": self.token, "timeout": "150000"}, # 告知 Function Server 150秒超時
            headers={"Content-Type": "application/json"},
            json={"code": code},
            timeout=160, # Python 客戶端 160秒就斷線，不要等到 200
        )
        resp.raise_for_status()
        return json.loads(resp.json()["data"])

    # ---------- 對外：抓取 ----------
    def fetch(
        self,
        lat: float,
        lng: float,
        zoom: int = 16,
        shot: bool = False,
        state: Optional[str] = None,
        ua_from_state: bool = False,
    ) -> Dict[str, Any]:
        code = self._load_script()
        code = self._inject_target_and_shot(code, lat, lng, zoom, shot)
        code = self._inject_state(code, state, ua_from_state)
        parsed = self._call_server(code)
        return {
            "items": parsed.get("results", []),
            "debug": parsed.get("debug", {}),
            "screenshot": parsed.get("screenshot"),
        }

    # ---------- 工具：IV% 回推 ----------
    @staticmethod
    def _iv_percent_from_triplet(atk: Any, df: Any, sta: Any) -> Optional[int]:
        try:
            if atk is None or df is None or sta is None:
                return None
            return int(round((int(atk) + int(df) + int(sta)) * 100.0 / 45.0))
        except Exception:
            return None

    # ---------- 後處理 + 篩選 ----------
    def postprocess(
        self,
        items: List[Dict[str, Any]],
        min_remaining: Optional[int] = None,
        min_iv: Optional[int] = None,
        perfect_only: bool = False,
        species: str = "",
    ) -> List[Dict[str, Any]]:
        # 標準化 + 回填 iv_pct + 算 expire_at
        norm: List[Dict[str, Any]] = []
        for r in items:
            row: Dict[str, Any] = {
                "id": r.get("id"),
                "name": r.get("name"),
                "gender": r.get("gender"),
                "remaining": r.get("remaining"),
                "expire_time": r.get("expire_time"),
                "latitude": r.get("latitude"),
                "longitude": r.get("longitude"),
                "iv_pct": r.get("iv_pct"),
                "iv_atk": r.get("iv_atk"),
                "iv_def": r.get("iv_def"),
                "iv_sta": r.get("iv_sta"),
                "cp": r.get("cp"),
                "level": r.get("level"),
                "is_perfect": r.get("is_perfect", False),
                "fast_move": r.get("fast_move"),
                "charge_move": r.get("charge_move"),
            }

            if row["iv_pct"] is None:
                alt = self._iv_percent_from_triplet(
                    row["iv_atk"], row["iv_def"], row["iv_sta"]
                )
                if alt is not None:
                    row["iv_pct"] = alt

            expire_ts = row.get("expire_time")
            if expire_ts is not None:
                try:
                    row["expire_at"] = datetime.fromtimestamp(
                        int(expire_ts), tz=TAIPEI_TZ
                    ).isoformat()

                except Exception:
                    row["expire_at"] = None
            else:
                row["expire_at"] = None

            norm.append(row)

        # 物種條件
        species_raw = [s.strip() for s in (species or "").split(",") if s.strip()]
        species_ids, species_kw = set(), []
        for s in species_raw:
            if re.fullmatch(r"\d+", s):
                species_ids.add(int(s))
            else:
                species_kw.append(s.lower())

        def species_match(row: Dict[str, Any]) -> bool:
            if not species_ids and not species_kw:
                return True
            ok_id = (row.get("id") in species_ids) if species_ids else False
            nm = row.get("name") or ""
            ok_kw = (
                any(k in nm.lower() for k in species_kw) if species_kw and nm else False
            )
            return ok_id or ok_kw

        def remaining_match(row: Dict[str, Any]) -> bool:
            if min_remaining is None:
                return True
            rem = row.get("remaining")
            return (rem is not None) and (int(rem) >= int(min_remaining))

        def iv_match(row: Dict[str, Any]) -> bool:
            if min_iv is None:
                return True
            ivp = row.get("iv_pct")
            return (ivp is not None) and (int(ivp) >= int(min_iv))

        out = [
            r for r in norm if species_match(r) and remaining_match(r) and iv_match(r)
        ]
        if perfect_only:
            out = [r for r in out if r.get("is_perfect")]

        out.sort(
            key=lambda x: (
                x["remaining"] if x["remaining"] is not None else 9_999_999,
                x.get("id") or 0,
            )
        )
        return out

    # ---------- 輸出 ----------
    def save(
        self,
        items: List[Dict[str, Any]],
        debug: Dict[str, Any],
        screenshot_b64: Optional[str],
        json_name: str = "pokemon_data.json",
        csv_name: str = "pokemon_data.csv",
        shot_name: str = "screenshot.png",
    ) -> None:
        with open(json_name, "w", encoding="utf-8") as f:
            json.dump({"items": items, "debug": debug}, f, ensure_ascii=False, indent=2)
        print("💾 Saved JSON to {}".format(json_name))

        with open(csv_name, "w", newline="", encoding="utf-8") as f:
            w = _csv.DictWriter(f, fieldnames=CSV_FIELDS)
            w.writeheader()
            for row in items:
                w.writerow({k: row.get(k) for k in CSV_FIELDS})
        print("🧾 Saved CSV to {}".format(csv_name))

        if screenshot_b64:
            try:
                b64 = (
                    screenshot_b64.split(",", 1)[1]
                    if "," in screenshot_b64
                    else screenshot_b64
                )
                with open(shot_name, "wb") as f:
                    f.write(base64.b64decode(b64))
                print("📸 Screenshot saved to {}".format(shot_name))
            except Exception as e:
                print("⚠️ screenshot save error: {}".format(e))


# ---------------- CLI ----------------
def main() -> None:
    ap = argparse.ArgumentParser(
        description="Fetch Pokémon locations (state/UA + filters)"
    )
    ap.add_argument("--lat", type=float, default=25.0478)
    ap.add_argument("--lng", type=float, default=121.5170)
    ap.add_argument("--zoom", type=int, default=16)
    ap.add_argument("--shot", action="store_true")
    ap.add_argument("--perfect-only", action="store_true")
    ap.add_argument("--state", type=str, help="Path to state.json or state_full.json")
    ap.add_argument(
        "--ua-from-state",
        action="store_true",
        help="Use UA/langs from state if available",
    )
    ap.add_argument(
        "--min-remaining",
        type=int,
        default=None,
        help="Keep entries with remaining >= SECONDS",
    )
    ap.add_argument("--min-iv", type=int, default=None, help="Keep IV% >= N")
    ap.add_argument(
        "--species",
        type=str,
        default="",
        help="Comma list of dex ids or name keywords, e.g. '133,669,伊布'",
    )
    ap.add_argument(
        "--script", type=str, default="twpk_radar.mjs", help="JS worker filename"
    )
    ap.add_argument("--url", type=str, default=DEFAULT_URL)
    ap.add_argument("--token", type=str, default=DEFAULT_TOKEN)
    args = ap.parse_args()

    client = PokemonRadarClient(url=args.url, token=args.token, script=args.script)
    print(
        "🚀 Fetching at {},{} zoom={} (script={})".format(
            args.lat, args.lng, args.zoom, os.path.basename(args.script)
        )
    )
    payload = client.fetch(
        args.lat, args.lng, args.zoom, args.shot, args.state, args.ua_from_state
    )

    print("----- debug notes -----")
    for n in payload.get("debug", {}).get("notes", []):
        print(n)
    print("-----------------------")
    print("🎯 Got {} results (before filters)".format(len(payload.get("items", []))))

    filtered = client.postprocess(
        payload["items"],
        min_remaining=args.min_remaining,
        min_iv=args.min_iv,
        perfect_only=args.perfect_only,
        species=args.species,
    )
    print("🔎 After filters: {} results".format(len(filtered)))

    client.save(filtered, payload.get("debug", {}), payload.get("screenshot"))


if __name__ == "__main__":
    main()