#!/usr/bin/env python3
"""
Hugging Face adapter for Athena.
Fetches trending models and datasets via the official HF API — no auth required.

API: https://huggingface.co/api/models  (models)
     https://huggingface.co/api/datasets (datasets)

Sort options: likes, downloads, lastModified, createdAt
Note: "trending" is NOT a valid sort parameter (verified against live API).
Strategy: sort=likes for popularity, filter recent models by createdAt,
          also grab sort=lastModified for recently updated models.

Rate limits: HF is generous for unauthenticated reads; 2s spacing is polite.

============================================================================
SCORING DESIGN PRINCIPLE (do not violate)
============================================================================
Structural metadata (likes, downloads, age) is an ADOPTION signal that is
genuine and useful — unlike arXiv where there are no upvotes. However, raw
popularity must never dominate content-relevance. A 100K-download model that
is a fine-tune of an existing base is less interesting than a 500-download
model that proposes a novel architecture. Content-relevance signals (pipeline
type, library, tags matching AI/agent/reasoning themes) MUST carry weight.
Download/like counts are score_type: actual (real numbers from HF).
"""

import json
import math
import os
import re
import time
import urllib.request
import urllib.error
from datetime import datetime, timedelta, timezone

from adapters import SourceAdapter
from adapters._store import true_first_seen, upsert_entries


class HuggingFaceAdapter(SourceAdapter):
    """Hugging Face API adapter — models and datasets."""

    BASE = "https://huggingface.co/api"

    # AI relevance signals: pipeline types and tags worth tracking
    AI_PIPELINE_TAGS = {
        "text-generation", "text2text-generation", "conversational",
        "text-classification", "token-classification", "feature-extraction",
        "text-to-image", "image-to-text", "image-to-image",
        "zero-shot-classification", "translation", "summarization",
        "reinforcement-learning", "question-answering",
        "automatic-speech-recognition", "text-to-speech",
        "fill-mask", "sentence-similarity",
    }

    # High-signal tags indicating AI/ML relevance
    AI_TAGS = [
        "transformers", "diffusers", "peft", "trl", "accelerate",
        "text-generation", "conversational", "llama", "gemma",
        "mistral", "qwen", "deepseek", "phi",
        "agent", "agents", "rag", "retrieval",
        "reinforcement-learning", "rlhf", "dpo", "orpo",
        "alignment", "reasoning", "multimodal", "vision-language",
        "code", "code-generation",
    ]

    # Libraries that indicate AI/ML models (vs. random uploads)
    AI_LIBRARIES = {
        "transformers", "diffusers", "peft", "trl", "accelerate",
        "vllm", "text-generation-inference", "sentence-transformers",
        "timm", "onnxruntime", "openvino",
    }

    def __init__(self, token: str = None):
        self.token = token or os.environ.get("HUGGINGFACE_TOKEN", "")

    def name(self) -> str:
        return "huggingface"

    def _headers(self):
        headers = {"User-Agent": "athena/0.1"}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        return headers

    def _request(self, path: str, max_retries: int = 2) -> list | dict | None:
        """Make a GET request to the HF API."""
        url = f"{self.BASE}{path}"
        req = urllib.request.Request(url, headers=self._headers())

        for attempt in range(max_retries + 1):
            try:
                with urllib.request.urlopen(req, timeout=20) as resp:
                    return json.loads(resp.read().decode("utf-8"))
            except (urllib.error.HTTPError, urllib.error.URLError) as e:
                if attempt < max_retries:
                    time.sleep(3 * (attempt + 1))
                    continue
                print(f"    HF API error: {e}")
                return None
            except Exception as e:
                print(f"    Request error: {e}")
                return None
        return None

    def _is_ai_relevant(self, model: dict) -> bool:
        """Check if a model/dataset is AI/ML relevant.

        Uses multiple signals to avoid noise from random uploads:
        - Has an AI-related pipeline_tag
        - Has an AI-related library_name
        - Has AI-related tags
        """
        pipeline = model.get("pipeline_tag", "")
        library = model.get("library_name", "")
        tags = model.get("tags", [])

        # Pipeline tag is the strongest signal
        if pipeline in self.AI_PIPELINE_TAGS:
            return True

        # Library is also strong
        if library in self.AI_LIBRARIES:
            return True

        # Tags as fallback
        tags_lower = [t.lower() for t in tags]
        if any(tag in tags_lower for tag in self.AI_TAGS):
            return True

        return False

    def _score(self, model: dict) -> float:
        """Score: adoption signal (likes/downloads) + AI relevance.

        Unlike arXiv (no upvotes), HF has real popularity metrics.
        However, adoption ≠ relevance — a 1M-download base model fine-tune
        is less interesting than a 500-download novel architecture.

        Design:
        - Adoption score (likes + downloads) — log scale, capped
        - Relevance bonus (pipeline type, library, tags) — carries weight
        - Recency bonus — newer models get slight boost
        """
        likes = model.get("likes", 0)
        downloads = model.get("downloads", 0)

        # Adoption: log scale on likes (stronger signal than downloads)
        # Downloads can be inflated by programmatic pulls; likes are intentional
        adoption = min(math.log1p(likes) / 2.0, 6.0)
        # Small download bonus (log scale, capped lower)
        download_bonus = min(math.log1p(downloads) / 4.0, 2.0)

        # Relevance: pipeline tag indicates what the model actually does
        pipeline = model.get("pipeline_tag", "")
        pipeline_bonus = 1.5 if pipeline in self.AI_PIPELINE_TAGS else 0.0

        # Library bonus: transformers/diffusers are curated ecosystems
        library = model.get("library_name", "")
        library_bonus = 0.5 if library in self.AI_LIBRARIES else 0.0

        # Tag relevance: specific AI tags indicate focused models
        tags = model.get("tags", [])
        ai_tag_count = sum(1 for t in tags if t.lower() in self.AI_TAGS)
        tag_bonus = min(ai_tag_count * 0.15, 1.0)

        # Recency: newer models get a slight boost
        created = model.get("createdAt", "")
        try:
            created_dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
            now = datetime.now(timezone.utc)
            age_days = (now - created_dt).days
            recency = max(0, 1.0 - age_days / 180.0)  # full bonus for <180d
        except (ValueError, TypeError):
            recency = 0.0

        return min(round(adoption + download_bonus + pipeline_bonus +
                          library_bonus + tag_bonus + recency, 2), 10.0)

    def _tags(self, model: dict) -> list:
        """Generate category tags from HF metadata."""
        tags = ["huggingface"]

        # Type: model or dataset
        if "pipeline_tag" in model:
            tags.append("type:model")
        elif "cardData" in model or "dataset" in str(model.get("id", "")).lower():
            tags.append("type:dataset")

        # Pipeline type
        pipeline = model.get("pipeline_tag", "")
        if pipeline:
            tags.append(f"pipeline:{pipeline}")

        # Library
        library = model.get("library_name", "")
        if library:
            tags.append(f"library:{library}")

        # AI-specific tags from HF metadata
        hf_tags = model.get("tags", [])
        ai_tag_subsets = {
            "topic:agent": ["agent", "agents", "tool-use"],
            "topic:reasoning": ["reasoning", "cot", "chain-of-thought"],
            "topic:alignment": ["alignment", "rlhf", "dpo", "orpo", "safe"],
            "topic:multimodal": ["multimodal", "vision-language", "image-text"],
            "topic:code": ["code", "code-generation", "code-llama"],
            "topic:llm": ["llama", "gemma", "mistral", "qwen", "phi", "deepseek"],
        }
        hf_tags_lower = [t.lower() for t in hf_tags]
        for tag_label, keywords in ai_tag_subsets.items():
            if any(kw in hf_tags_lower for kw in keywords):
                tags.append(tag_label)

        # Organization signal
        org_id = model.get("id", "").split("/")[0] if "/" in model.get("id", "") else ""
        known_orgs = ["meta-llama", "openai", "anthropic", "deepseek-ai",
                       "google", "microsoft", "mistralai", "huggingface",
                       "stabilityai", "nvidia", "qwen", "01-ai"]
        if org_id.lower() in known_orgs:
            tags.append(f"org:{org_id}")

        # Adoption level
        likes = model.get("likes", 0)
        if likes > 5000:
            tags.append("adoption:high")
        elif likes > 500:
            tags.append("adoption:medium")
        elif likes > 50:
            tags.append("adoption:low")

        return tags

    def fetch(self, query: str = "", limit: int = 20) -> list[dict]:
        """
        Fetch AI models from Hugging Face.

        Strategy:
        1. Fetch top models by likes (popularity signal)
        2. Fetch recently modified models (fresh signal)
        3. Filter for AI relevance
        4. Deduplicate, score, sort, return top N
        """
        now = datetime.now(timezone.utc)

        all_models = []

        # 1. Popular models (sort=likes) — 60 to get enough AI-relevant ones
        popular = self._request("/models?sort=likes&limit=60")
        if popular and isinstance(popular, list):
            all_models.extend(popular)

        time.sleep(1)  # polite spacing

        # 2. Recently modified (sort=lastModified) — fresh models getting attention
        # Filter to models created in last 90 days to avoid noise
        cutoff = (now - timedelta(days=90)).strftime("%Y-%m-%dT%H:%M:%S.000Z")
        recent = self._request(f"/models?sort=lastModified&limit=40")
        if recent and isinstance(recent, list):
            all_models.extend(recent)

        # 3. If query provided, also search
        if query:
            search = self._request(f"/models?search={urllib.parse.quote(query)}&sort=likes&limit=20")
            if search and isinstance(search, list):
                all_models.extend(search)

        # Deduplicate by id
        seen = set()
        unique = []
        for m in all_models:
            mid = m.get("id", m.get("modelId", ""))
            if mid and mid not in seen:
                seen.add(mid)
                unique.append(m)
        all_models = unique

        # Filter for AI relevance
        ai_models = [m for m in all_models if self._is_ai_relevant(m)]

        # Score and sort
        for m in ai_models:
            m["_score"] = self._score(m)
        ai_models.sort(key=lambda m: m.get("_score", 0), reverse=True)
        ai_models = ai_models[:limit]

        # Convert to DB format
        entries = []
        for model in ai_models:
            score = model.pop("_score", 0)
            model_id = model.get("id", model.get("modelId", ""))
            source_id = model_id.replace("/", "__")

            # URL
            url = f"https://huggingface.co/{model_id}"

            # Title: model ID with context
            org = model_id.split("/")[0] if "/" in model_id else "unknown"
            name = model_id.split("/")[-1] if "/" in model_id else model_id
            pipeline = model.get("pipeline_tag", "")
            if pipeline:
                title = f"{name} ({pipeline}) by {org}"
            else:
                title = f"{name} by {org}"

            # Extracted text: model tags + metadata as text (no card fetch needed)
            tags_text = ", ".join(model.get("tags", []))
            library = model.get("library_name", "")
            likes = model.get("likes", 0)
            downloads = model.get("downloads", 0)
            extracted_text = f"Library: {library}. Likes: {likes}. Downloads: {downloads}. Tags: {tags_text}"

            tags = self._tags(model)

            # Structured metadata
            raw_meta = {
                "model_id": model_id,
                "org": org,
                "name": name,
                "likes": likes,
                "downloads": downloads,
                "pipeline_tag": pipeline,
                "library_name": library,
                "tags": model.get("tags", []),
                "createdAt": model.get("createdAt", ""),
                "lastModified": model.get("lastModified", ""),
                "score_type": "actual",  # real likes/downloads from HF
            }

            now_str = now.strftime("%Y-%m-%dT%H:%M:%SZ")
            first_seen = true_first_seen(raw_meta, "huggingface", now_str)
            entries.append({
                "source": "huggingface",
                "source_id": source_id,
                "url": url,
                "title": title,
                "extracted_text": extracted_text,
                "summary": None,
                "category_tags": json.dumps(tags),
                "signal_score": score,
                "raw_metadata": json.dumps(raw_meta),
                "first_seen": first_seen,
                "last_updated": now_str,
            })

        return entries


if __name__ == "__main__":
    import argparse
    import sqlite3
    import urllib.parse

    parser = argparse.ArgumentParser(description="Hugging Face adapter for Athena")
    parser.add_argument("--limit", type=int, default=20, help="Max entries")
    parser.add_argument("--db", default=os.path.join(os.path.dirname(__file__), "..", "oracle.db"), help="SQLite DB")
    parser.add_argument("--schema", default=os.path.join(os.path.dirname(__file__), "..", "schema.sql"), help="Schema file")
    parser.add_argument("--dry-run", action="store_true", help="Don't store in DB")
    args = parser.parse_args()

    print(f"=== Hugging Face Adapter ===")
    print(f"  Limit: {args.limit}")
    print()

    adapter = HuggingFaceAdapter()
    entries = adapter.fetch(limit=args.limit)

    print(f"  Fetched {len(entries)} entries")

    if not args.dry_run:
        conn = sqlite3.connect(args.db)
        if os.path.exists(args.schema):
            with open(args.schema) as f:
                conn.executescript(f.read())
            conn.commit()

        cur = conn.cursor()
        stored = upsert_entries(conn, entries)
        conn.commit()
        conn.close()
        print(f"  Stored {stored} entries")

    # Print top 5
    print(f"\n  Top entries:")
    for i, e in enumerate(entries[:5]):
        meta = json.loads(e["raw_metadata"]) if isinstance(e["raw_metadata"], str) else e["raw_metadata"]
        print(f"    [{i+1}] score={e['signal_score']:.2f} likes={meta.get('likes', '?'):,} downloads={meta.get('downloads', '?'):,}")
        print(f"        {e['title'][:90]}")
        print(f"        {e['url']}")

    print(f"\n  Done.")
