from __future__ import annotations

import json
import logging
import os
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from niannian.config import load_config, load_providers_config, resolve_config_path, resolve_data_home
from niannian.storage.baidu_backup import run_baidu_backup
from niannian.storage.cloud_mirror import run_cloud_mirror
from niannian.storage.cloud_tier import CloudTierRouter
from niannian.storage.tier_chain import is_feimao_disabled
from niannian.storage.local import LocalMemoryStore
from niannian.storage.ops_error_log import OpsErrorLog
from niannian.storage.sync.cron import cron_matches
from niannian.storage.sync.dream_lock import DreamLock

logger = logging.getLogger("niannian.worker")

_DEFAULT_CRONS = {
    "cloud_sync": "*/15 * * * *",
    "cloud_mirror": "0 2 * * *",
    "baidu_backup": "0 2 * * *",
    "tier_health_check": "*/15 * * * *",
    "cloud_refresh": "0 */6 * * *",
    "feimao_probe": "0 */6 * * *",
    "monthly_archive": "0 3 1 * *",
    "dream": "0 3 * * *",
    "nas_archive": "0 4 * * 0",
}

_ENV_CRON_KEYS = {
    "cloud_sync": "NIANNIAN_WORKER_CLOUD_SYNC_CRON",
    "cloud_mirror": "NIANNIAN_WORKER_CLOUD_MIRROR_CRON",
    "baidu_backup": "NIANNIAN_WORKER_BAIDU_BACKUP_CRON",
    "tier_health_check": "NIANNIAN_WORKER_TIER_HEALTH_CRON",
    "cloud_refresh": "NIANNIAN_WORKER_CLOUD_REFRESH_CRON",
    "feimao_probe": "NIANNIAN_WORKER_FEIMAO_PROBE_CRON",
    "monthly_archive": "NIANNIAN_WORKER_MONTHLY_ARCHIVE_CRON",
    "dream": "NIANNIAN_WORKER_DREAM_CRON",
    "nas_archive": "NIANNIAN_WORKER_NAS_ARCHIVE_CRON",
}

_CONFIG_CRON_KEYS = {
    "cloud_sync": "cloud_sync_cron",
    "cloud_mirror": "cloud_mirror_cron",
    "baidu_backup": "baidu_backup_cron",
    "tier_health_check": "tier_health_check_cron",
    "cloud_refresh": "cloud_refresh_cron",
    "feimao_probe": "feimao_probe_cron",
    "monthly_archive": "monthly_archive_cron",
    "dream": "dream_cron",
    "nas_archive": "nas_archive_cron",
}


def resolve_cron(job: str, config: dict[str, Any]) -> str:
    env_key = _ENV_CRON_KEYS[job]
    if os.environ.get(env_key):
        return os.environ[env_key]
    worker_cfg = ((config.get("runtime") or {}).get("docker") or {}).get("worker") or {}
    tier_cfg = (config.get("storage") or {}).get("tier") or {}
    if job == "cloud_mirror":
        mirror_cron = (tier_cfg.get("mirror") or {}).get("cron")
        if mirror_cron:
            return str(mirror_cron)
    if job == "tier_health_check":
        interval = tier_cfg.get("health_check_interval_minutes")
        if interval == 15:
            return "*/15 * * * *"
    cred_cfg = (config.get("storage") or {}).get("credentials") or {}
    refresh_cron = (cred_cfg.get("refresh") or {}).get("worker_cron")
    if job == "cloud_refresh" and refresh_cron:
        return str(refresh_cron)
    feimao_cfg = tier_cfg.get("feimao") or {}
    probe_cron = (feimao_cfg.get("connectivity") or {}).get("probe_cron")
    if job == "feimao_probe" and probe_cron:
        return str(probe_cron)
    archive_cfg = (config.get("storage") or {}).get("monthly_archive") or {}
    if job == "monthly_archive" and archive_cfg.get("cron"):
        return str(archive_cfg["cron"])
    return str(worker_cfg.get(_CONFIG_CRON_KEYS[job]) or _DEFAULT_CRONS[job])


def storage_primary_mode(config: dict[str, Any]) -> str:
    return str(((config.get("storage") or {}).get("primary") or {}).get("mode", "tiered"))


def should_run_dream_on_node(config: dict[str, Any]) -> bool:
    mode = os.environ.get("NIANNIAN_MODE") or (config.get("runtime") or {}).get("mode", "auto")
    defer = ((config.get("runtime") or {}).get("mac") or {}).get("defer_dream_to_docker", True)
    if mode == "docker":
        return True
    if mode == "mac" and defer:
        return False
    return True


class DockerWorker:
    """Docker cron 循环：hourly sync · tier health · daily mirror · dream · archive。"""

    def __init__(
        self,
        *,
        home: Path | None = None,
        config: dict[str, Any] | None = None,
        project: str = "global",
    ) -> None:
        self._config_path = resolve_config_path()
        self._config = config if config is not None else load_config(self._config_path)
        self._home = Path(home) if home is not None else resolve_data_home()
        self._project = project
        self._providers = load_providers_config()
        self._store = LocalMemoryStore(
            self._home,
            config=self._config,
            providers=self._providers,
        )
        self._state_path = self._home / "sync" / "worker_state.json"
        self._dream_lock = DreamLock(self._home / "sync" / "dream.lock")
        self._poll_sec = int(os.environ.get("NIANNIAN_WORKER_POLL_SEC", "30"))
        self._crons = {job: resolve_cron(job, self._config) for job in _DEFAULT_CRONS}

    @property
    def store(self) -> LocalMemoryStore:
        return self._store

    def _load_state(self) -> dict[str, Any]:
        if not self._state_path.is_file():
            return {"last_runs": {}}
        return json.loads(self._state_path.read_text(encoding="utf-8"))

    def _save_state(self, state: dict[str, Any]) -> None:
        self._state_path.parent.mkdir(parents=True, exist_ok=True)
        self._state_path.write_text(
            json.dumps(state, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    def _minute_key(self, dt: datetime) -> str:
        return dt.strftime("%Y%m%d%H%M")

    def _should_run_job(self, job: str, now: datetime, *, force: bool = False) -> bool:
        if force:
            return True
        state = self._load_state()
        key = self._minute_key(now)
        if state.get("last_runs", {}).get(job) == key:
            return False
        return cron_matches(self._crons[job], now)

    def _mark_job_run(self, job: str, now: datetime) -> None:
        state = self._load_state()
        last_runs = state.setdefault("last_runs", {})
        last_runs[job] = self._minute_key(now)
        state[f"last_{job}"] = now.astimezone(timezone.utc).isoformat()
        self._save_state(state)

    def run_cloud_sync(self) -> dict[str, Any]:
        pull = self._store.pull_from_cloud()
        push = self._store.flush_to_cloud(force=True)
        return {
            "pull": pull.to_dict(),
            "push": push.to_dict(),
            "ok": pull.ok and push.ok,
        }

    def run_dream(self) -> dict[str, Any]:
        if not should_run_dream_on_node(self._config):
            return {"ok": True, "skipped": True, "reason": "defer_dream_to_docker"}

        if not self._dream_lock.acquire():
            return {"ok": True, "skipped": True, "reason": "dream lock held by another writer"}

        try:
            from niannian import NianNian

            nn = NianNian(home=self._home, project=self._project, config_path=self._config_path)
            stats = nn.dream(project=self._project)
            push = self._store.flush_to_cloud(force=True)
            return {"ok": True, "dream": stats, "push": push.to_dict()}
        finally:
            self._dream_lock.release()

    def run_baidu_backup(self) -> dict[str, Any]:
        return run_baidu_backup(self._store, config=self._config)

    def run_cloud_mirror(self) -> dict[str, Any]:
        return run_cloud_mirror(self._store, config=self._config)

    def run_tier_health_check(self) -> dict[str, Any]:
        cloud = self._store.cloud
        if not isinstance(cloud, CloudTierRouter):
            return {"ok": True, "skipped": True, "reason": "not_tiered_mode"}
        health = cloud.check_health(record_failure=True)
        return {
            "ok": True,
            "active_role": health.active_role,
            "message": health.message,
            "degraded": health.degraded,
        }

    def run_cloud_refresh(self) -> dict[str, Any]:
        from niannian.storage.cloud_auth.manager import CloudCredentialManager
        from niannian.storage.feimao_probe import run_feimao_probe

        mgr = CloudCredentialManager(self._home, config=self._config)
        tier_state_path = self._home / "sync" / "tier_state.json"
        feimao_reach = None
        if tier_state_path.is_file():
            try:
                tier_state = json.loads(tier_state_path.read_text(encoding="utf-8"))
                feimao_reach = (tier_state.get("providers") or {}).get("feimao", {}).get("reachability")
            except (json.JSONDecodeError, OSError):
                feimao_reach = None
        try:
            results = mgr.refresh_all_if_needed(feimao_reachability=feimao_reach)
        except Exception as exc:
            logger.warning("worker cloud_refresh failed: %s", exc)
            return {"ok": True, "error": str(exc)[:200], "providers": {}}
        return {
            "ok": all(CloudCredentialManager._refresh_result_ok(r) for r in results.values()),
            "providers": {k: {"ok": v.ok, "status": v.status, "message": v.message} for k, v in results.items()},
        }

    def run_monthly_archive(self) -> dict[str, Any]:
        from niannian.storage.monthly_archive import run_monthly_archive

        return run_monthly_archive(self._store, config=self._config)

    def run_feimao_probe(self) -> dict[str, Any]:
        from niannian.storage.cloud_auth.manager import CloudCredentialManager
        from niannian.storage.feimao_probe import run_feimao_probe

        mgr = CloudCredentialManager(self._home, config=self._config)
        return run_feimao_probe(self._home, config=self._config, credential_manager=mgr)

    def run_daily_cold_backup(self) -> dict[str, Any]:
        if storage_primary_mode(self._config) == "primary_backup":
            return self.run_baidu_backup()
        return self.run_cloud_mirror()

    def run_archive(self) -> dict[str, Any]:
        nas_cfg = ((self._config.get("sync") or {}).get("push_to_nas") or {})
        archive_cfg = (self._config.get("storage") or {}).get("archive") or {}
        if nas_cfg.get("require_dream_completed", True):
            meta_path = self._home / "meta.json"
            last_dream = None
            if meta_path.is_file():
                try:
                    last_dream = json.loads(meta_path.read_text(encoding="utf-8")).get("last_dream_at")
                except (json.JSONDecodeError, OSError):
                    last_dream = None
            if not last_dream:
                return {
                    "ok": True,
                    "chunks_archived": 0,
                    "skipped_reason": "dream_not_completed",
                }
        min_age = int(archive_cfg.get("min_age_days", 7))
        archive_dir = Path(os.environ.get("NIANNIAN_ARCHIVE_DIR", "/data/archive"))
        result = self._store.archive_distilled(min_age_days=min_age, archive_dir=archive_dir)
        return result.to_dict()

    def tick(self, now: datetime | None = None, *, force: bool = False) -> dict[str, Any]:
        now = now or datetime.now()
        results: dict[str, Any] = {"ts": now.astimezone(timezone.utc).isoformat(), "jobs": {}}
        tiered = storage_primary_mode(self._config) == "tiered"

        if self._should_run_job("cloud_sync", now, force=force):
            logger.info("worker: cloud_sync")
            results["jobs"]["cloud_sync"] = self.run_cloud_sync()
            self._mark_job_run("cloud_sync", now)

        if tiered and self._should_run_job("tier_health_check", now, force=force):
            logger.info("worker: tier_health_check")
            results["jobs"]["tier_health_check"] = self.run_tier_health_check()
            self._mark_job_run("tier_health_check", now)

        if tiered and self._should_run_job("cloud_refresh", now, force=force):
            logger.info("worker: cloud_refresh")
            results["jobs"]["cloud_refresh"] = self.run_cloud_refresh()
            self._mark_job_run("cloud_refresh", now)

        tier_cfg = (self._config.get("storage") or {}).get("tier") or {}
        if tiered and not is_feimao_disabled(tier_cfg) and self._should_run_job("feimao_probe", now, force=force):
            logger.info("worker: feimao_probe")
            results["jobs"]["feimao_probe"] = self.run_feimao_probe()
            self._mark_job_run("feimao_probe", now)

        archive_cfg = (self._config.get("storage") or {}).get("monthly_archive") or {}
        if archive_cfg.get("enabled", True) and self._should_run_job("monthly_archive", now, force=force):
            logger.info("worker: monthly_archive")
            results["jobs"]["monthly_archive"] = self.run_monthly_archive()
            self._mark_job_run("monthly_archive", now)

        mirror_job = "cloud_mirror" if tiered else "baidu_backup"
        if self._should_run_job(mirror_job, now, force=force):
            logger.info("worker: %s", mirror_job)
            results["jobs"][mirror_job] = self.run_daily_cold_backup()
            self._mark_job_run(mirror_job, now)

        if self._should_run_job("dream", now, force=force):
            logger.info("worker: dream")
            results["jobs"]["dream"] = self.run_dream()
            self._mark_job_run("dream", now)

        try:
            from niannian import NianNian
            from niannian.soul.learning_auto import run_skill_learning_tick

            nn = NianNian(home=self._home, project=self._project, config_path=self._config_path)
            skill = run_skill_learning_tick(nn, manual=False)
            if skill.get("ran"):
                logger.info("worker: skill_learning batch")
                results["jobs"]["skill_learning"] = skill
        except Exception as exc:
            logger.warning("worker: skill_learning tick failed: %s", exc)

        if self._should_run_job("nas_archive", now, force=force):
            logger.info("worker: nas_archive")
            results["jobs"]["nas_archive"] = self.run_archive()
            self._mark_job_run("nas_archive", now)

        self._record_job_errors(results.get("jobs") or {})
        return results

    def _record_job_errors(self, jobs: dict[str, Any]) -> None:
        log = OpsErrorLog(self._home, config=self._config)
        for job_name, payload in jobs.items():
            if not isinstance(payload, dict):
                continue
            if payload.get("ok", True) and not payload.get("errors"):
                continue
            if payload.get("skipped"):
                continue
            message = ""
            if payload.get("errors"):
                errs = payload.get("errors")
                if isinstance(errs, list):
                    message = "; ".join(str(e) for e in errs[:3])
                else:
                    message = str(errs)
            elif payload.get("message"):
                message = str(payload.get("message"))
            elif payload.get("reason"):
                message = str(payload.get("reason"))
            log.record(
                f"worker_{job_name}_failed",
                component="worker",
                message=message or job_name,
                detail={"job": job_name, "result": payload},
                check_key=f"worker:{job_name}",
            )

    def run_once(self) -> dict[str, Any]:
        """冒烟 / 手动：立即执行全部任务（忽略 cron）。"""
        return self.tick(force=True)

    def run(self) -> None:
        logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
        logger.info(
            "niannian worker started home=%s mode=%s crons=%s",
            self._home,
            os.environ.get("NIANNIAN_MODE", "auto"),
            self._crons,
        )
        while True:
            try:
                self.tick()
            except Exception as exc:
                logger.exception("worker tick failed")
                OpsErrorLog(self._home, config=self._config).record(
                    "worker_tick_failed",
                    component="worker",
                    message=str(exc)[:500],
                    check_key="worker:tick",
                )
            time.sleep(self._poll_sec)


def run() -> None:
    DockerWorker().run()


def main() -> None:
    run()


if __name__ == "__main__":
    main()
