"""
Model card generator for merged/stacked models.

Generates HuggingFace-compatible model cards (README.md) that document
the merge configuration, base models, and usage instructions.

Usage::

    from model_osmosis.utils.model_card import ModelCardGenerator, MergeCardConfig

    config = MergeCardConfig(
        model_name="my-merged-model",
        base_models=["model-a", "model-b"],
        merge_strategy="weighted_average",
        merge_weights=[0.6, 0.4],
    )

    generator = ModelCardGenerator()
    card_text = generator.generate(config)
    generator.save(config, "README.md")
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import json


@dataclass
class MergeCardConfig:
    """Configuration for generating a model card.

    Args:
        model_name: Name of the merged model.
        base_models: List of base model IDs that were merged.
        merge_strategy: Strategy used for merging (e.g., 'weighted_average', 'slerp').
        merge_weights: Optional weights used for each base model.
        description: Optional description of the merged model.
        license: License for the merged model.
        tags: Tags for the model card.
        language: Language(s) the model supports.
        pipeline_tag: HuggingFace pipeline tag (e.g., 'text-generation').
        extra_metadata: Additional metadata to include.
    """

    model_name: str
    base_models: List[str] = field(default_factory=list)
    merge_strategy: str = "unknown"
    merge_weights: Optional[List[float]] = None
    description: str = ""
    license: str = "apache-2.0"
    tags: List[str] = field(default_factory=lambda: ["merge", "model-osmosis"])
    language: Optional[List[str]] = None
    pipeline_tag: Optional[str] = None
    extra_metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class StackCardConfig:
    """Configuration for generating a stacked model card.

    Args:
        model_name: Name of the stacked model.
        expert_models: List of expert model IDs.
        expert_names: Human-readable names for each expert.
        top_k: Number of experts activated per input.
        description: Optional description.
        license: License for the stacked model.
        tags: Tags for the model card.
    """

    model_name: str
    expert_models: List[str] = field(default_factory=list)
    expert_names: Optional[List[str]] = None
    top_k: int = 2
    description: str = ""
    license: str = "apache-2.0"
    tags: List[str] = field(default_factory=lambda: ["moe", "stacked", "model-osmosis"])
    extra_metadata: Dict[str, Any] = field(default_factory=dict)


class ModelCardGenerator:
    """Generate HuggingFace model cards for merged/stacked models."""

    MERGE_TEMPLATE = '''---
license: {license}
tags:
{tags}
base_model:
{base_models}
library_name: model-osmosis
{extra_yaml}
---

# {model_name}

{description}

## Model Details

This model was created by merging the following models using [ModelOsmosis](https://github.com/model-osmosis/model-osmosis):

| Base Model | Weight |
|------------|--------|
{model_table}

### Merge Configuration

- **Strategy**: `{merge_strategy}`
- **Weights**: `{merge_weights}`
- **Created**: {created_date}
- **Library**: ModelOsmosis v{version}

## Usage

```python
from model_osmosis import HuggingFaceAdapter

adapter = HuggingFaceAdapter()
model = adapter.load("{model_name}")

# Or load directly with transformers
from transformers import AutoModel
model = AutoModel.from_pretrained("{model_name}")
```

## Merge Recipe

To reproduce this merge:

```python
from model_osmosis import OsmosisPipeline

pipeline = OsmosisPipeline()
merged = pipeline.merge(
    models={base_model_list},
    strategy="{merge_strategy}",
    model_weights={merge_weights},
)
```

## Base Models

{base_model_details}

---

*Generated with [ModelOsmosis](https://github.com/model-osmosis/model-osmosis)*
'''

    STACK_TEMPLATE = '''---
license: {license}
tags:
{tags}
base_model:
{base_models}
library_name: model-osmosis
{extra_yaml}
---

# {model_name}

{description}

## Model Details

This is a **Mixture-of-Experts (MoE)** model created by stacking the following expert models using [ModelOsmosis](https://github.com/model-osmosis/model-osmosis):

| Expert | Model |
|--------|-------|
{expert_table}

### Stack Configuration

- **Architecture**: Sparse Mixture-of-Experts
- **Top-K**: {top_k} experts per input
- **Total Experts**: {num_experts}
- **Created**: {created_date}
- **Library**: ModelOsmosis v{version}

## How It Works

1. A learned **gating network** routes each input to the top-{top_k} most relevant experts
2. Only the selected experts are computed (sparse activation)
3. Expert outputs are weighted by gate scores and combined
4. This preserves 100% of each expert's knowledge while enabling specialization

## Usage

```python
from model_osmosis import StackedModel
import torch

# Load the stacked model
model = StackedModel.from_pretrained("{model_name}")

# Forward pass
x = torch.randn(batch_size, input_dim)
output = model(x)  # Automatically routes to best experts
```

## Stack Recipe

To reproduce this stack:

```python
from model_osmosis import OsmosisPipeline, HuggingFaceAdapter

pipeline = OsmosisPipeline()
adapter = HuggingFaceAdapter()

experts = adapter.load_multiple({expert_list})
stacked = pipeline.stack(
    experts,
    sample_data=your_data,
    top_k={top_k},
)
```

## Expert Models

{expert_details}

---

*Generated with [ModelOsmosis](https://github.com/model-osmosis/model-osmosis)*
'''

    def __init__(self, version: str = "0.4.0") -> None:
        self.version = version

    def generate(self, config: Union[MergeCardConfig, StackCardConfig]) -> str:
        """Generate a model card string.

        Args:
            config: Configuration for the model card.

        Returns:
            The model card as a string (markdown format).
        """
        if isinstance(config, StackCardConfig):
            return self._generate_stack_card(config)
        return self._generate_merge_card(config)

    def _generate_merge_card(self, config: MergeCardConfig) -> str:
        """Generate a merge model card."""
        # Format tags
        tags = "\n".join(f"- {tag}" for tag in config.tags)

        # Format base models
        base_models = "\n".join(f"- {m}" for m in config.base_models)

        # Format weights
        weights = config.merge_weights or [1.0 / len(config.base_models)] * len(config.base_models)
        weights_str = ", ".join(f"{w:.4f}" for w in weights)

        # Model table
        model_rows = []
        for model, weight in zip(config.base_models, weights):
            model_rows.append(f"| [{model}](https://huggingface.co/{model}) | {weight:.4f} |")
        model_table = "\n".join(model_rows)

        # Base model details
        base_details = []
        for model in config.base_models:
            base_details.append(f"- **[{model}](https://huggingface.co/{model})**")
        base_model_details = "\n".join(base_details)

        # Extra YAML
        extra_yaml_parts = []
        if config.language:
            extra_yaml_parts.append(f"language:\n" + "\n".join(f"- {l}" for l in config.language))
        if config.pipeline_tag:
            extra_yaml_parts.append(f"pipeline_tag: {config.pipeline_tag}")
        extra_yaml = "\n".join(extra_yaml_parts)

        # Description
        description = config.description or f"A merged model combining {len(config.base_models)} base models using the {config.merge_strategy} strategy."

        return self.MERGE_TEMPLATE.format(
            license=config.license,
            tags=tags,
            base_models=base_models,
            extra_yaml=extra_yaml,
            model_name=config.model_name,
            description=description,
            model_table=model_table,
            merge_strategy=config.merge_strategy,
            merge_weights=f"[{weights_str}]",
            created_date=datetime.now().strftime("%Y-%m-%d"),
            version=self.version,
            base_model_list=json.dumps(config.base_models),
            base_model_details=base_model_details,
        )

    def _generate_stack_card(self, config: StackCardConfig) -> str:
        """Generate a stacked model card."""
        # Format tags
        tags = "\n".join(f"- {tag}" for tag in config.tags)

        # Format base models
        base_models = "\n".join(f"- {m}" for m in config.expert_models)

        # Expert names
        expert_names = config.expert_names or [f"Expert {i}" for i in range(len(config.expert_models))]

        # Expert table
        expert_rows = []
        for name, model in zip(expert_names, config.expert_models):
            expert_rows.append(f"| {name} | [{model}](https://huggingface.co/{model}) |")
        expert_table = "\n".join(expert_rows)

        # Expert details
        expert_details = []
        for name, model in zip(expert_names, config.expert_models):
            expert_details.append(f"- **{name}**: [{model}](https://huggingface.co/{model})")
        expert_details_str = "\n".join(expert_details)

        # Extra YAML
        extra_yaml_parts = []
        for key, value in config.extra_metadata.items():
            extra_yaml_parts.append(f"{key}: {value}")
        extra_yaml = "\n".join(extra_yaml_parts)

        # Description
        description = config.description or f"A Mixture-of-Experts model combining {len(config.expert_models)} specialized experts with top-{config.top_k} routing."

        return self.STACK_TEMPLATE.format(
            license=config.license,
            tags=tags,
            base_models=base_models,
            extra_yaml=extra_yaml,
            model_name=config.model_name,
            description=description,
            expert_table=expert_table,
            top_k=config.top_k,
            num_experts=len(config.expert_models),
            created_date=datetime.now().strftime("%Y-%m-%d"),
            version=self.version,
            expert_list=json.dumps(config.expert_models),
            expert_details=expert_details_str,
        )

    def save(
        self,
        config: Union[MergeCardConfig, StackCardConfig],
        path: Union[str, Path],
    ) -> Path:
        """Save the model card to a file.

        Args:
            config: Configuration for the model card.
            path: Output file path (typically README.md).

        Returns:
            The path to the saved file.
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        card_text = self.generate(config)
        path.write_text(card_text, encoding="utf-8")

        return path

    def generate_config_json(
        self,
        config: Union[MergeCardConfig, StackCardConfig],
    ) -> Dict[str, Any]:
        """Generate a config.json for the merged/stacked model.

        This creates a minimal config that can be used with transformers.

        Args:
            config: Configuration for the model.

        Returns:
            Dict that can be saved as config.json.
        """
        if isinstance(config, StackCardConfig):
            return {
                "architectures": ["StackedModel"],
                "model_type": "stacked_moe",
                "num_experts": len(config.expert_models),
                "top_k": config.top_k,
                "expert_models": config.expert_models,
                "expert_names": config.expert_names,
                "library_name": "model-osmosis",
            }
        else:
            return {
                "architectures": ["MergedModel"],
                "model_type": "merged",
                "base_models": config.base_models,
                "merge_strategy": config.merge_strategy,
                "merge_weights": config.merge_weights,
                "library_name": "model-osmosis",
            }

    def save_config_json(
        self,
        config: Union[MergeCardConfig, StackCardConfig],
        path: Union[str, Path],
    ) -> Path:
        """Save config.json for the model.

        Args:
            config: Configuration for the model.
            path: Output file path.

        Returns:
            The path to the saved file.
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        config_dict = self.generate_config_json(config)
        path.write_text(json.dumps(config_dict, indent=2), encoding="utf-8")

        return path
