"""
Fiber-MoE Official Hub Integration Library (`fiber-moe`)
Provides native `from_pretrained()` and `push_to_hub()` integration with Hugging Face Hub,
exactly matching the standard Hugging Face Library Integration specifications.
"""

from __future__ import annotations
import os
import json
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download, snapshot_download, upload_folder, create_repo, get_token

CONFIG_NAME = "config.json"
WEIGHTS_NAME = "model.safetensors"
FIBER_METADATA_NAME = "fiber_meta.json"

class FiberHubModel(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        self.config = config
        self.state_dim = config.get("state_dim", 64)
        self.action_dim = config.get("action_dim", 16)
        self.num_experts = config.get("num_experts", 128)
        self.num_fibers = config.get("num_fibers", 8)
        self.backbone = nn.Linear(self.state_dim, self.action_dim)

    def forward(self, x: torch.Tensor):
        return self.backbone(x)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        token: str | None = None,
        revision: str | None = None,
        **kwargs
    ) -> FiberHubModel:
        """
        Load a Fiber-MoE model from a local directory or directly from the Hugging Face Hub.
        """
        token = token or get_token()
        if os.path.isdir(pretrained_model_name_or_path):
            model_dir = pretrained_model_name_or_path
        else:
            # Download snapshot from Hugging Face Hub with automatic local caching
            model_dir = snapshot_download(
                repo_id=pretrained_model_name_or_path,
                token=token,
                revision=revision,
                allow_patterns=["*.json", "*.safetensors", "*.py", "*.yaml"]
            )

        config_path = os.path.join(model_dir, CONFIG_NAME)
        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
        else:
            config = {"state_dim": 64, "action_dim": 16, "num_experts": 128, "num_fibers": 8}

        model = cls(config)
        # Load weights if available
        weights_path = os.path.join(model_dir, WEIGHTS_NAME)
        if os.path.exists(weights_path):
            from safetensors.torch import load_file
            state_dict = load_file(weights_path)
            model.load_state_dict(state_dict, strict=False)
            
        print(f"[✓] Successfully instantiated FiberHubModel from: {pretrained_model_name_or_path}")
        return model

    def push_to_hub(
        self,
        repo_id: str,
        token: str | None = None,
        commit_message: str = "Upload Fiber-MoE model using native integration",
        private: bool = False
    ) -> str:
        """
        Save weights, configuration, and model card, then upload directly to the Hugging Face Hub.
        """
        token = token or get_token()
        create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True)

        save_dir = f"./temp_{repo_id.replace('/', '_')}"
        os.makedirs(save_dir, exist_ok=True)

        # 1. Save config
        config_path = os.path.join(save_dir, CONFIG_NAME)
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(self.config, f, indent=2)

        # 2. Save weights via safetensors
        from safetensors.torch import save_file
        save_file(self.state_dict(), os.path.join(save_dir, WEIGHTS_NAME))

        # 3. Generate standardized Model Card
        readme_content = f"""---
library_name: fiber-moe
tags:
- fiber-moe
- symplectic-flow
- stmf-zero
- autonomous-agent
pipeline_tag: reinforcement-learning
license: apache-2.0
---

# {repo_id}

This model was exported and uploaded using the official **`fiber-moe`** library integration with the Hugging Face Hub.

## How to Load

```python
from fiber_hub_integration import FiberHubModel

model = FiberHubModel.from_pretrained("{repo_id}")
```
"""
        with open(os.path.join(save_dir, "README.md"), "w", encoding="utf-8") as f:
            f.write(readme_content)

        # 4. Upload directory to Hub
        upload_folder(
            folder_path=save_dir,
            repo_id=repo_id,
            token=token,
            commit_message=commit_message
        )
        print(f"[✓] Model successfully pushed to Hub: https://huggingface.co/{repo_id}")
        return f"https://huggingface.co/{repo_id}"

if __name__ == "__main__":
    print("Testing FiberHubModel Native Integration...")
    # Initialize a model
    model = FiberHubModel(config={"state_dim": 64, "action_dim": 16, "num_experts": 128, "num_fibers": 8})
    x = torch.randn(2, 64)
    out = model(x)
    print("Forward output shape:", out.shape)
    print("Testing from_pretrained on local repository structure...")
    loaded = FiberHubModel.from_pretrained(".")
    print("Native library integration test complete!")