# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re
import sys
import unicodedata
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


# Ensure nemo_gym is importable when run as a pre-commit hook outside the venv.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import yaml

from nemo_gym.server_metadata import ServerMetadata, visit_agent_server, visit_resources_server


README_PATH = Path("README.md")

RESOURCES_SERVERS_FOLDER = Path("resources_servers")
RESPONSES_API_AGENTS_FOLDER = Path("responses_api_agents")


@dataclass
class AgentDatasetsMetadata:
    """Metadata extracted from agent datasets configuration."""

    license: str | None = None
    types: list[str] = field(default_factory=list)
    huggingface_repo_id: Optional[str] = None

    def to_dict(self) -> dict[str, str | list[str] | None]:  # pragma: no cover
        """Convert to dict for backward compatibility."""
        return {
            "huggingface_repo_id": self.huggingface_repo_id,
            "license": self.license,
            "types": self.types,
        }


@dataclass
class ConfigMetadata:
    """Combined metadata from YAML configuration file."""

    huggingface_repo_id: Optional[str] = None
    domain: Optional[str] = None
    description: Optional[str] = None
    verified: bool = False
    verified_url: Optional[str] = None
    value: Optional[str] = None
    license: Optional[str] = None
    types: list[str] = field(default_factory=list)

    @classmethod
    def from_yaml_data(
        cls, resource: ServerMetadata, agent: AgentDatasetsMetadata
    ) -> "ConfigMetadata":  # pragma: no cover
        """Combine resources server and agent datasets metadata."""
        return cls(
            domain=resource.domain,
            description=resource.description,
            verified=resource.verified,
            verified_url=resource.verified_url,
            value=resource.value,
            huggingface_repo_id=agent.huggingface_repo_id,
            license=agent.license,
            types=agent.types,
        )


@dataclass
class ServerInfo:
    """Information about a resources server for table generation."""

    name: str
    display_name: str
    config_metadata: ConfigMetadata
    config_path: str
    config_filename: str
    readme_path: str
    yaml_file: Path
    base_folder: str = "resources_servers"

    @property
    def huggingface_repo_id(self) -> str | None:  # pragma: no cover
        return self.config_metadata.huggingface_repo_id

    @property
    def domain(self) -> str | None:  # pragma: no cover
        return self.config_metadata.domain

    @property
    def types(self) -> list[str]:  # pragma: no cover
        return self.config_metadata.types

    def get_domain_or_empty(self) -> str:  # pragma: no cover
        return self.config_metadata.domain or ""

    def get_description_or_dash(self) -> str:  # pragma: no cover
        return self.config_metadata.description or "-"

    def get_value_or_dash(self) -> str:  # pragma: no cover
        return self.config_metadata.value or "-"

    def get_license_or_dash(self) -> str:  # pragma: no cover
        return self.config_metadata.license or "-"

    def get_verified_mark(self) -> str:  # pragma: no cover
        if self.config_metadata.verified and self.config_metadata.verified_url:
            return f"<a href='{self.config_metadata.verified_url}'>✓</a>"
        elif self.config_metadata.verified:
            return "✓"
        else:
            return "-"

    def get_train_mark(self) -> str:  # pragma: no cover
        return "✓" if "train" in set(self.config_metadata.types) else "-"

    def get_validation_mark(self) -> str:  # pragma: no cover
        return "✓" if "validation" in set(self.config_metadata.types) else "-"

    def get_dataset_link(self) -> str:  # pragma: no cover
        if not self.config_metadata.huggingface_repo_id:
            return "-"
        repo_id = self.config_metadata.huggingface_repo_id
        dataset_name = repo_id.split("/")[-1]
        dataset_url = f"https://huggingface.co/datasets/{repo_id}"
        return f"<a href='{dataset_url}'>{dataset_name}</a>"

    def get_config_link(self, use_filename: bool = True) -> str:  # pragma: no cover
        return f"<a href='{self.config_path}'>{self.config_filename if use_filename else 'config'}</a>"

    def get_readme_link(self) -> str:  # pragma: no cover
        return f"<a href='{self.readme_path}'>README</a>"


def visit_agent_datasets(data: dict) -> AgentDatasetsMetadata:  # pragma: no cover
    agent = AgentDatasetsMetadata()
    if not isinstance(data, dict):
        return agent
    for v1 in data.values():
        if not isinstance(v1, dict):
            continue
        v2 = v1.get("responses_api_agents")
        if not isinstance(v2, dict):
            continue
        for v3 in v2.values():
            if not isinstance(v3, dict):
                continue
            datasets = v3.get("datasets")
            if isinstance(datasets, list):
                for entry in datasets:
                    if isinstance(entry, dict):
                        agent.types.append(entry.get("type"))
                        if entry.get("type") == "train":
                            agent.license = entry.get("license")
                            hf_id = entry.get("huggingface_identifier")
                            if hf_id and isinstance(hf_id, dict):
                                agent.huggingface_repo_id = hf_id.get("repo_id")
            elif v3.get("harbor_datasets") or v3.get("vf_env_id"):
                agent.types.append("train")
    return agent


def agent_has_resources_server_ref(data: dict) -> bool:  # pragma: no cover
    if not isinstance(data, dict):
        return False
    for v1 in data.values():
        if not isinstance(v1, dict):
            continue
        v2 = v1.get("responses_api_agents")
        if not isinstance(v2, dict):
            continue
        for v3 in v2.values():
            if isinstance(v3, dict) and v3.get("resources_server"):
                return True
    return False


def extract_config_metadata(yaml_path: Path, from_agent: bool = False) -> ConfigMetadata:  # pragma: no cover
    """
    Domain:
        {name}_resources_server:
            resources_servers:
                {name}:
                    domain: {example_domain}
                    verified: {true/false}
                    description: {example_description}
                    value: {example_value}
                    ...
        {something}_simple_agent:
            responses_api_agents:
                simple_agent:
                    datasets:
                        - name: train
                          type: {example_type_1}
                          license: {example_license_1}
                          huggingface_identifier:
                            repo_id: {example_repo_id_1}
                            artifact_fpath: {example_artifact_fpath_1}
                        - name: validation
                          type: {example_type_2}
                          license: {example_license_2}
    """
    with yaml_path.open() as f:
        data = yaml.safe_load(f)

    resource_data = visit_agent_server(data) if from_agent else visit_resources_server(data)
    agent_data = visit_agent_datasets(data)

    return ConfigMetadata.from_yaml_data(resource_data, agent_data)


def get_training_server_info() -> list[ServerInfo]:  # pragma: no cover
    """Collect training-ready server metadata (skips example_* servers)."""
    training_servers = []

    for base_folder in (RESOURCES_SERVERS_FOLDER, RESPONSES_API_AGENTS_FOLDER):
        from_agent = base_folder == RESPONSES_API_AGENTS_FOLDER
        for subdir in base_folder.iterdir():
            if not subdir.is_dir():
                continue

            if subdir.name.startswith("example_"):
                continue

            configs_folder = subdir / "configs"
            if not (configs_folder.exists() and configs_folder.is_dir()):
                continue

            yaml_files = list(configs_folder.glob("*.yaml"))
            if not yaml_files:
                continue

            for yaml_file in yaml_files:
                if from_agent:
                    with yaml_file.open() as f:
                        raw = yaml.safe_load(f) or {}
                    if agent_has_resources_server_ref(raw):
                        continue

                yaml_data = extract_config_metadata(yaml_file, from_agent=from_agent)
                if not yaml_data.types:
                    continue

                server_name = subdir.name
                display_name = server_name.replace("_", " ").title()
                config_path = f"{base_folder.name}/{server_name}/configs/{yaml_file.name}"
                readme_path = f"{base_folder.name}/{server_name}/README.md"

                training_servers.append(
                    ServerInfo(
                        name=server_name,
                        display_name=display_name,
                        config_metadata=yaml_data,
                        config_path=config_path,
                        config_filename=yaml_file.name,
                        readme_path=readme_path,
                        yaml_file=yaml_file,
                        base_folder=base_folder.name,
                    )
                )

    return training_servers


def generate_training_table(servers: list[ServerInfo]) -> str:  # pragma: no cover
    """Generate table for training resources servers."""
    col_names = [
        "Environment",
        "Domain",
        "Description",
        "Value",
        "Train",
        "Validation",
        "License",
        "Config",
        "Dataset",
        # TODO: Add back in when we can verify resources servers
        # "Verified",
    ]
    if not servers:
        return handle_empty_table(col_names)

    rows = []

    for server in servers:
        rows.append(
            [
                server.display_name,
                server.get_domain_or_empty(),
                server.get_description_or_dash(),
                server.get_value_or_dash(),
                server.get_train_mark(),
                server.get_validation_mark(),
                server.get_license_or_dash(),
                server.get_config_link(use_filename=True),
                server.get_dataset_link(),
                # TODO: Add back in when we can verify resources servers
                # verified_mark,
            ]
        )

    rows.sort(
        key=lambda r: (
            normalize_str(r[0]),  # resources server name
            normalize_str(r[1]),  # domain
            tuple(normalize_str(cell) for cell in r),
        )
    )

    table = [col_names, ["-" for _ in col_names]] + rows
    return format_table(table)


def handle_empty_table(col_names: list[str]) -> str:  # pragma: no cover
    """Generate an empty table when there are no servers."""
    separator = ["-" * len(col_name) for col_name in col_names]
    return format_table([col_names, separator])


def normalize_str(s: str) -> str:  # pragma: no cover
    """
    Rows with identical domain values may get reordered differently
    between local and CI runs. We normalize text and
    use all columns as tie-breakers to ensure deterministic sorting.
    """
    if not s or not isinstance(s, str):
        return ""
    return unicodedata.normalize("NFKD", s).casefold().strip()


def format_table(table: list[list[str]]) -> str:  # pragma: no cover
    """Format grid of data into markdown table."""
    col_widths = []
    num_cols = len(table[0])

    for i in range(num_cols):
        max_len = 0
        for row in table:
            cell_len = len(str(row[i]))
            if cell_len > max_len:
                max_len = cell_len
        col_widths.append(max_len)

    # Pretty print cells for raw markdown readability
    formatted_rows = []
    for i, row in enumerate(table):
        formatted_cells = []
        for j, cell in enumerate(row):
            cell = str(cell)
            col_width = col_widths[j]
            pad_total = col_width - len(cell)
            if i == 1:  # header separater
                formatted_cells.append(cell * col_width)
            else:
                formatted_cells.append(cell + " " * pad_total)
        formatted_rows.append("| " + (" | ".join(formatted_cells)) + " |")

    return "\n".join(formatted_rows)


def main():  # pragma: no cover
    text = README_PATH.read_text()

    training_servers = get_training_server_info()
    training_table_str = generate_training_table(training_servers)

    training_pattern = re.compile(
        r"(<!-- START_TRAINING_SERVERS_TABLE -->)(.*?)(<!-- END_TRAINING_SERVERS_TABLE -->)",
        flags=re.DOTALL,
    )

    if not training_pattern.search(text):
        sys.stderr.write(
            "Error: README.md does not contain <!-- START_TRAINING_SERVERS_TABLE --> and <!-- END_TRAINING_SERVERS_TABLE --> markers.\n"
        )
        sys.exit(1)

    text = training_pattern.sub(lambda m: f"{m.group(1)}\n{training_table_str}\n{m.group(3)}", text)

    README_PATH.write_text(text)


if __name__ == "__main__":
    main()
