"""Watchtower — Scan HuggingFace for new popular open-weight models.

Continuously monitors the HuggingFace Hub for instruction-tuned,
open-license text-generation models that are gaining traction.  Tracks
which models have been seen, queued, and obliterated in a persistent
JSON state file (~/.obliteratus/watchtower_state.json).

Usage:
    from obliteratus.watchtower import Watchtower

    wt = Watchtower()
    new_models = wt.scan()          # returns list of newly discovered models
    trending  = wt.get_trending()   # returns current hot models sorted by downloads
    wt.start_scheduler(interval=3600)  # background hourly scan
    wt.stop_scheduler()
"""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from .service_contracts import (
    SchedulerEvent,
    SchedulerState,
    has_open_license,
    is_instruction_tuned,
    normalize_watchtower_candidate,
    scheduler_transition,
    validate_scheduler_interval,
    validate_watchtower_status,
)

logger = logging.getLogger(__name__)

# ── Constants ─────────────────────────────────────────────────────────

STATE_DIR = Path.home() / ".obliteratus"
STATE_FILE = STATE_DIR / "watchtower_state.json"

# Minimum downloads in last 7 days to qualify as "popular"
MIN_RECENT_DOWNLOADS = 1000

# Open-weight licenses we care about (lowercase, prefix-matched)
OPEN_LICENSES = frozenset({
    "apache-2.0", "mit", "bsd-2-clause", "bsd-3-clause",
    "llama2", "llama3", "llama3.1", "llama3.2", "llama3.3", "llama4",
    "gemma", "qwen", "deepseek",
    "cc-by-4.0", "cc-by-sa-4.0", "cc-by-nc-4.0",
    "openrail", "openrail++", "bigscience-openrail-m",
    "artistic-2.0", "wtfpl", "unlicense", "zlib",
    "other",  # many open models use "other" + a permissive custom license
})

# Keywords that indicate instruction/chat tuning
INSTRUCT_KEYWORDS = frozenset({
    "instruct", "chat", "it", "rlhf", "dpo", "sft",
    "aligned", "conversational", "assistant", "dialogue",
})

# ── Organizations to watch ────────────────────────────────────────────

MAJOR_LABS = {
    "google", "meta-llama", "mistralai", "Qwen", "deepseek-ai",
    "microsoft", "nvidia", "01-ai", "internlm", "THUDM",
    "tiiuae", "CohereForAI", "allenai", "openai", "openai-community",
    "moonshotai", "openbmb", "stabilityai", "stepfun-ai", "zai-org",
    "MiniMaxAI",
}

MEDIUM_LABS = {
    "teknium", "NousResearch", "OpenBuddy", "lmsys",
    "HuggingFaceH4", "mosaicml", "bigcode",
    "cognitivecomputations", "mlabonne", "huihui-ai",
    "Orenguteng", "WhiteRabbitNeo",
}

ALL_WATCHED_ORGS = MAJOR_LABS | MEDIUM_LABS


# ── Data classes ──────────────────────────────────────────────────────

@dataclass
class DiscoveredModel:
    """A model discovered by the Watchtower."""
    model_id: str                     # e.g. "meta-llama/Llama-3.1-8B-Instruct"
    name: str                         # short display name
    org: str                          # organization
    downloads_7d: int = 0             # downloads in last 7 days
    total_downloads: int = 0          # all-time downloads
    likes: int = 0                    # HF likes
    size_category: str = ""           # e.g. "7B", "13B", "70B"
    license: str = ""                 # license identifier
    pipeline_tag: str = ""            # should be "text-generation"
    discovered_at: str = ""           # ISO timestamp
    status: str = "new"               # new | queued | obliterating | obliterated | failed
    obliteration_metrics: dict = field(default_factory=dict)
    last_updated: str = ""            # ISO timestamp of last status change

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "DiscoveredModel":
        # Handle extra/missing keys gracefully
        if not isinstance(d, dict):
            raise ValueError("model state must be an object")
        known = {f.name for f in cls.__dataclass_fields__.values()}
        try:
            model = cls(**{k: v for k, v in d.items() if k in known})
        except TypeError as exc:
            raise ValueError("model state is missing required fields") from exc
        for name in (
            "model_id", "name", "org", "size_category", "license", "pipeline_tag",
            "discovered_at", "last_updated",
        ):
            if not isinstance(getattr(model, name), str):
                raise ValueError(f"model state {name} must be a string")
        for name in ("downloads_7d", "total_downloads", "likes"):
            value = getattr(model, name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"model state {name} must be a non-negative integer")
        model.status = validate_watchtower_status(model.status)
        if not isinstance(model.obliteration_metrics, dict):
            raise ValueError("model state obliteration_metrics must be an object")
        return model


# ── Watchtower class ──────────────────────────────────────────────────

class Watchtower:
    """Scans HuggingFace Hub for new popular open-weight instruction models."""

    def __init__(
        self,
        state_file: Path | str | None = None,
        *,
        clock: Callable[[], datetime] | None = None,
        fetch_models: Callable[[], list[Any]] | None = None,
        thread_factory: Callable[..., threading.Thread] | None = None,
        event_factory: Callable[[], threading.Event] | None = None,
        scheduler_join_timeout: float = 5.0,
    ):
        self.state_file = Path(state_file) if state_file else STATE_FILE
        self._clock = clock or (lambda: datetime.now(timezone.utc))
        self._fetch_models = fetch_models or self._fetch_models_from_hf
        self._thread_factory = thread_factory or threading.Thread
        self._event_factory = event_factory or threading.Event
        self._scheduler_join_timeout = validate_scheduler_interval(scheduler_join_timeout)
        self._models: dict[str, DiscoveredModel] = {}
        self._last_scan: str | None = None
        self._scan_count: int = 0
        self._lock = threading.Lock()
        self._scheduler_lock = threading.Lock()
        self._scheduler_thread: threading.Thread | None = None
        self._scheduler_stop = self._event_factory()
        self._scheduler_state = SchedulerState.STOPPED
        self._on_new_model_callbacks: list = []

        # Load persisted state
        self._load_state()

    # ── Persistence ───────────────────────────────────────────────────

    def _load_state(self):
        """Load watchtower state from disk."""
        try:
            if self.state_file.exists():
                data = json.loads(self.state_file.read_text(encoding="utf-8"))
                if not isinstance(data, dict):
                    raise ValueError("state root must be an object")
                last_scan = data.get("last_scan")
                if last_scan is not None and not isinstance(last_scan, str):
                    raise ValueError("last_scan must be a string or null")
                scan_count = data.get("scan_count", 0)
                if isinstance(scan_count, bool) or not isinstance(scan_count, int) or scan_count < 0:
                    raise ValueError("scan_count must be a non-negative integer")
                raw_models = data.get("models", {})
                if not isinstance(raw_models, dict):
                    raise ValueError("models must be an object")
                loaded: dict[str, DiscoveredModel] = {}
                for mid, mdata in raw_models.items():
                    if not isinstance(mid, str) or not isinstance(mdata, dict):
                        raise ValueError("model state entries must be objects keyed by strings")
                    model = DiscoveredModel.from_dict(mdata)
                    if not model.model_id or model.model_id != mid:
                        raise ValueError("model state key must match model_id")
                    loaded[mid] = model
                self._last_scan = last_scan
                self._scan_count = scan_count
                self._models = loaded
                logger.info(
                    "Watchtower: loaded %d models from %s",
                    len(self._models), self.state_file,
                )
        except Exception as e:
            logger.warning("Watchtower: failed to load state: %s", e)

    def _save_state(self):
        """Persist watchtower state to disk."""
        tmp = self.state_file.with_suffix(".tmp")
        try:
            self.state_file.parent.mkdir(parents=True, exist_ok=True)
            with self._lock:
                data = {
                    "last_scan": self._last_scan,
                    "scan_count": self._scan_count,
                    "models": {mid: m.to_dict() for mid, m in self._models.items()},
                }
                # Serialize the shared temporary path as well as its snapshot.
                tmp.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
                tmp.replace(self.state_file)
        except Exception as e:
            try:
                tmp.unlink(missing_ok=True)
            except OSError:
                pass
            logger.warning("Watchtower: failed to save state: %s", e)

    # ── HuggingFace API helpers ───────────────────────────────────────

    @staticmethod
    def _fetch_models_from_hf(
        *,
        orgs: set[str] | None = None,
        limit_per_org: int = 50,
        min_downloads: int = MIN_RECENT_DOWNLOADS,
    ) -> list[Any]:
        """Fetch model metadata from HuggingFace Hub API.

        Returns a list of raw model info dicts. Gracefully returns []
        if the API is unreachable.
        """
        try:
            from huggingface_hub import HfApi
        except ImportError:
            logger.error("huggingface_hub not installed — cannot scan HF")
            return []

        api = HfApi()
        results = []
        search_orgs = sorted(orgs or ALL_WATCHED_ORGS, key=str.casefold)

        for org in search_orgs:
            try:
                models = api.list_models(
                    author=org,
                    pipeline_tag="text-generation",
                    sort="downloads",
                    direction=-1,
                    limit=limit_per_org,
                )
                for m in models:
                    if (getattr(m, "downloads", 0) or 0) >= min_downloads:
                        results.append(m)
            except Exception as e:
                logger.debug("Watchtower: error scanning org '%s': %s", org, e)
                continue

        return results

    @staticmethod
    def _is_instruction_tuned(model_id: str, tags: list[str] | None = None) -> bool:
        """Heuristic check if a model is instruction/chat tuned."""
        return is_instruction_tuned(model_id, tags, INSTRUCT_KEYWORDS)

    @staticmethod
    def _has_open_license(license_id: str | None) -> bool:
        """Check if the license is considered open-weight."""
        return has_open_license(license_id, OPEN_LICENSES)

    @staticmethod
    def _estimate_size(model_id: str, config: dict | None = None) -> str:
        """Estimate model size from the name or config."""
        # Try to extract a size like "7b", "70b", "1.5b", "397b"
        import re
        match = re.search(r'(\d+\.?\d*)\s*[bB]', model_id)
        if match:
            size = float(match.group(1))
            if size >= 1:
                return f"{match.group(1)}B"
            else:
                return f"{size * 1000:.0f}M"
        return "unknown"

    # ── Core scan logic ───────────────────────────────────────────────

    def scan(self, on_log=None) -> list[DiscoveredModel]:
        """Scan HuggingFace for new popular open-weight instruction models.

        Returns a list of *newly discovered* models (not previously seen).
        Thread-safe.
        """
        def _log(msg):
            logger.info(msg)
            if on_log:
                on_log(msg)

        _log("Watchtower: starting scan...")
        new_models = []

        try:
            raw_models = self._fetch_models()
            if not isinstance(raw_models, list):
                raise ValueError("HF API response must be a list")
        except Exception as e:
            _log(f"Watchtower: HF API error — {e}")
            return []

        _log(f"Watchtower: fetched {len(raw_models)} candidate models from HF Hub")

        now = self._now_iso()
        processed = 0

        for m in raw_models:
            candidate = normalize_watchtower_candidate(m)
            if candidate is None:
                continue
            model_id, downloads, likes, tags, license_id, pipeline_tag = candidate

            # Filter: must be text-generation
            if pipeline_tag and pipeline_tag != "text-generation":
                continue

            # Filter: minimum popularity
            if downloads < MIN_RECENT_DOWNLOADS:
                continue

            # Filter: open license
            if not self._has_open_license(license_id):
                continue

            # Filter: instruction-tuned (or from a major lab — they usually are)
            org = model_id.split("/")[0] if "/" in model_id else ""
            is_major = org in MAJOR_LABS
            if not is_major and not self._is_instruction_tuned(model_id, tags):
                continue

            processed += 1

            # Check if we've already seen this model
            with self._lock:
                if model_id in self._models:
                    # Update download counts
                    existing = self._models[model_id]
                    existing.downloads_7d = downloads  # HF "downloads" is ~recent
                    existing.total_downloads = downloads
                    existing.likes = likes
                    existing.last_updated = now
                    continue

                # New model!
                dm = DiscoveredModel(
                    model_id=model_id,
                    name=model_id.split("/")[-1] if "/" in model_id else model_id,
                    org=org,
                    downloads_7d=downloads,
                    total_downloads=downloads,
                    likes=likes,
                    size_category=self._estimate_size(model_id),
                    license=license_id,
                    pipeline_tag=pipeline_tag or "text-generation",
                    discovered_at=now,
                    status="new",
                    last_updated=now,
                )
                self._models[model_id] = dm
                new_models.append(dm)

        with self._lock:
            self._last_scan = now
            self._scan_count += 1

        self._save_state()

        _log(
            f"Watchtower: scan complete — {processed} qualifying models, "
            f"{len(new_models)} new discoveries"
        )

        # Fire callbacks for new models
        with self._lock:
            callbacks = list(self._on_new_model_callbacks)
        for dm in new_models:
            for cb in callbacks:
                try:
                    cb(dm)
                except Exception as e:
                    logger.warning("Watchtower callback error: %s", e)

        return new_models

    # ── Query methods ─────────────────────────────────────────────────

    def get_trending(self, limit: int = 50) -> list[DiscoveredModel]:
        """Return current hot models sorted by recent downloads (descending)."""
        with self._lock:
            models = list(self._models.values())
        models.sort(key=lambda m: m.downloads_7d, reverse=True)
        return models[:limit]

    def get_new_models(self) -> list[DiscoveredModel]:
        """Return models with status 'new' (not yet queued or obliterated)."""
        with self._lock:
            return [m for m in self._models.values() if m.status == "new"]

    def get_obliterated(self) -> list[DiscoveredModel]:
        """Return models that have been obliterated."""
        with self._lock:
            return [m for m in self._models.values() if m.status == "obliterated"]

    def get_all_models(self) -> list[DiscoveredModel]:
        """Return all tracked models."""
        with self._lock:
            return list(self._models.values())

    def get_model(self, model_id: str) -> DiscoveredModel | None:
        """Get a specific model by ID."""
        with self._lock:
            return self._models.get(model_id)

    def set_status(self, model_id: str, status: str, metrics: dict | None = None):
        """Update a model's status (new/queued/obliterating/obliterated/failed)."""
        status = validate_watchtower_status(status)
        with self._lock:
            if model_id not in self._models:
                return False
            m = self._models[model_id]
            m.status = status
            m.last_updated = self._now_iso()
            if metrics is not None:
                m.obliteration_metrics = metrics
        self._save_state()
        return True

    def get_stats(self) -> dict[str, Any]:
        """Return summary statistics."""
        with self._lock:
            total = len(self._models)
            by_status = {}
            for m in self._models.values():
                by_status[m.status] = by_status.get(m.status, 0) + 1
            last_scan = self._last_scan
            scan_count = self._scan_count
        return {
            "total_tracked": total,
            "last_scan": last_scan,
            "scan_count": scan_count,
            "by_status": by_status,
        }

    def get_model_choices(self) -> list[str]:
        """Return model IDs suitable for a dropdown, trending first."""
        trending = self.get_trending(limit=100)
        return [m.model_id for m in trending]

    # ── Callbacks ─────────────────────────────────────────────────────

    def on_new_model(self, callback):
        """Register a callback for when a new model is discovered.

        Signature: callback(model: DiscoveredModel)
        """
        with self._lock:
            self._on_new_model_callbacks.append(callback)

    # ── Scheduler ─────────────────────────────────────────────────────

    def start_scheduler(self, interval: int = 3600, on_log=None):
        """Start background scanning at the given interval (seconds).

        Safe to call multiple times — restarts with new interval.
        """
        interval_seconds = validate_scheduler_interval(interval)
        if not self.stop_scheduler():
            raise RuntimeError("Watchtower scheduler did not stop; refusing duplicate start")
        stop_event = self._event_factory()

        def _run():
            while not stop_event.is_set():
                try:
                    self.scan(on_log=on_log)
                except Exception as e:
                    logger.error("Watchtower scheduler error: %s", e)
                stop_event.wait(timeout=interval_seconds)

        with self._scheduler_lock:
            self._scheduler_stop = stop_event
            self._scheduler_state = scheduler_transition(
                self._scheduler_state,
                SchedulerEvent.START,
            )
            thread = self._thread_factory(
                target=_run, daemon=True, name="watchtower-scheduler"
            )
            self._scheduler_thread = thread
            try:
                thread.start()
            except BaseException:
                self._scheduler_state = scheduler_transition(
                    self._scheduler_state,
                    SchedulerEvent.REQUEST_STOP,
                )
                self._scheduler_state = scheduler_transition(
                    self._scheduler_state,
                    SchedulerEvent.STOP_CONFIRMED,
                )
                self._scheduler_thread = None
                raise
        logger.info("Watchtower: scheduler started (interval=%ss)", interval_seconds)

    def stop_scheduler(self) -> bool:
        """Stop the background scheduler if running."""
        with self._scheduler_lock:
            thread = self._scheduler_thread
            if thread is None:
                return True
            if self._scheduler_state is SchedulerState.RUNNING:
                self._scheduler_state = scheduler_transition(
                    self._scheduler_state,
                    SchedulerEvent.REQUEST_STOP,
                )
            self._scheduler_stop.set()
            if thread is threading.current_thread():
                logger.info(
                    "Watchtower: scheduler stop requested from its own thread; "
                    "confirmation deferred",
                )
                return False
            thread.join(timeout=self._scheduler_join_timeout)
            if thread.is_alive():
                self._scheduler_state = scheduler_transition(
                    self._scheduler_state,
                    SchedulerEvent.STOP_TIMEOUT,
                )
                logger.warning("Watchtower: scheduler did not stop before timeout")
                return False
            self._scheduler_state = scheduler_transition(
                self._scheduler_state,
                SchedulerEvent.STOP_CONFIRMED,
            )
            self._scheduler_thread = None
            logger.info("Watchtower: scheduler stopped")
            return True

    @property
    def is_scanning(self) -> bool:
        """True if the scheduler is actively running."""
        with self._scheduler_lock:
            return (
                self._scheduler_state is SchedulerState.RUNNING
                and self._scheduler_thread is not None
                and self._scheduler_thread.is_alive()
            )

    def _now_iso(self) -> str:
        """Return an aware clock reading normalized to UTC ISO format."""
        current = self._clock()
        if not isinstance(current, datetime) or current.tzinfo is None:
            raise ValueError("Watchtower clock must return a timezone-aware datetime")
        return current.astimezone(timezone.utc).isoformat()

    # ── Table formatting ──────────────────────────────────────────────

    def format_table(self, models: list[DiscoveredModel] | None = None) -> list[list[str]]:
        """Format models as a list of rows for a Gradio Dataframe.

        Columns: [Name, Org, Size, Downloads, Likes, License, Discovered, Status]
        """
        if models is None:
            models = self.get_trending()

        rows = []
        for m in models:
            discovered = ""
            if m.discovered_at:
                try:
                    dt = datetime.fromisoformat(m.discovered_at)
                    discovered = dt.strftime("%Y-%m-%d %H:%M")
                except Exception:
                    discovered = m.discovered_at[:16]

            status_emoji = {
                "new": "🆕",
                "queued": "⏳",
                "obliterating": "⚡",
                "obliterated": "✅",
                "failed": "❌",
            }.get(m.status, "❓")

            rows.append([
                m.model_id,
                m.org,
                m.size_category,
                f"{m.downloads_7d:,}",
                str(m.likes),
                m.license,
                discovered,
                f"{status_emoji} {m.status}",
            ])
        return rows

    TABLE_HEADERS = [
        "Model ID", "Org", "Size", "Downloads (7d)",
        "Likes", "License", "Discovered", "Status",
    ]


# ── Module-level singleton ────────────────────────────────────────────
# Lazy-initialized so importing the module doesn't trigger disk I/O.

_watchtower_instance: Watchtower | None = None
_watchtower_lock = threading.Lock()


def get_watchtower() -> Watchtower:
    """Get or create the global Watchtower singleton."""
    global _watchtower_instance
    if _watchtower_instance is None:
        with _watchtower_lock:
            if _watchtower_instance is None:
                _watchtower_instance = Watchtower()
    return _watchtower_instance
