# Copyright (C) 2026 GoodData Corporation

"""Rule storage and indexing for the rules retrieval system."""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, cast

import orjson
import yaml

from aida_mcp.tools.rules.vocab import harvest_section_tokens


def _compute_rule_hash(
    *,
    path: str,
    description: str,
    always_apply: bool,
    required: list[str],
    related: list[str],
    origin: str,
    editable: bool,
    content: str,
) -> str:
    """
    Compute a stable SHA-256 hash for a rule.

    This is used as a "fingerprint" to detect when rule content/metadata changes.
    The hash is computed from a canonical JSON representation of the fields that
    affect retrieval + output (not from raw file bytes), so it is stable across
    equivalent serializations.
    """
    payload = {
        "path": path,
        "description": description,
        "always_apply": always_apply,
        "required": required,
        "related": related,
        "origin": origin,
        "editable": editable,
        "content": content,
    }
    # orjson.dumps returns bytes, use OPT_SORT_KEYS for canonical ordering
    raw = orjson.dumps(payload, option=orjson.OPT_SORT_KEYS).decode("utf-8")
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


@dataclass
class Rule:
    """Represents a single rule file with metadata and content."""

    path: str  # Relative path from the rules directory (e.g., "services/example-service.mdc")
    description: str  # From YAML frontmatter
    content: str  # Full rule content (without frontmatter)
    always_apply: bool = False  # Whether this rule should always be included

    # Dependencies from YAML frontmatter
    required: list[str] = field(default_factory=lambda: list[str]())  # Rules that must be included
    related: list[str] = field(default_factory=lambda: list[str]())  # Related rules (informational)

    # Fingerprints (computed at load time / post-init)
    rule_id: str = ""  # Human-readable stable identifier (defaults to path)
    rule_hash: str = ""  # SHA-256 fingerprint (hex)
    origin: str = "repo"
    editable: bool = True

    # Preprocessed for search
    tokens: set[str] = field(default_factory=lambda: set[str]())
    boost_tokens: set[str] = field(default_factory=lambda: set[str]())

    def __post_init__(self) -> None:
        """Tokenize content for keyword search (with automatic vocabulary harvesting)."""
        if not self.rule_id:
            self.rule_id = self.path
        if not self.rule_hash:
            self.rule_hash = _compute_rule_hash(
                path=self.path,
                description=self.description,
                always_apply=self.always_apply,
                required=self.required,
                related=self.related,
                origin=self.origin,
                editable=self.editable,
                content=self.content,
            )
        harvested = harvest_section_tokens(heading_path=None, text=f"{self.description}\n{self.content}")
        self.tokens = harvested.tokens
        self.boost_tokens = harvested.boost_tokens


class RuleStore:
    """Loads and indexes rules from a rules directory."""

    def __init__(self, rules_dir: Path):
        self.rules_dir = rules_dir
        self._rules: list[Rule] = []
        self._loaded = False

    def load(self) -> None:
        """Load all .mdc rule files from the rules directory."""
        if self._loaded:
            return

        if not self.rules_dir.exists():
            self._loaded = True
            return

        source_manifest = self._load_source_manifest()

        for mdc_file in self.rules_dir.rglob("*.mdc"):
            rule = self._parse_rule_file(mdc_file, source_manifest=source_manifest)
            if rule:
                self._rules.append(rule)

        self._loaded = True

    def _load_source_manifest(self) -> dict[str, dict[str, object]]:
        manifest_path = self.rules_dir / ".aida-rule-sources.json"
        if not manifest_path.exists():
            return {}
        try:
            raw = json.loads(manifest_path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            return {}
        if not isinstance(raw, dict):
            return {}
        return {str(k): cast(dict[str, object], v) for k, v in raw.items() if isinstance(v, dict)}

    def _parse_rule_file(self, file_path: Path, *, source_manifest: dict[str, dict[str, object]]) -> Rule | None:
        """Parse a single .mdc file into a Rule object."""
        try:
            content = file_path.read_text(encoding="utf-8")
        except OSError:
            return None

        # Parse YAML frontmatter
        description: str = ""
        always_apply: bool = False
        required: list[str] = []
        related: list[str] = []
        body = content

        if content.startswith("---"):
            parts = content.split("---", 2)
            if len(parts) >= 3:
                try:
                    frontmatter = yaml.safe_load(parts[1])
                    if frontmatter is not None and isinstance(frontmatter, dict):
                        fm = cast(dict[str, Any], frontmatter)
                        description = str(fm.get("description") or "")
                        always_apply = fm.get("alwaysApply") is True
                        # Parse dependencies
                        raw_required = fm.get("required")
                        if isinstance(raw_required, list):
                            required = [str(item) for item in cast(list[Any], raw_required) if item]
                        raw_related = fm.get("related")
                        if isinstance(raw_related, list):
                            related = [str(item) for item in cast(list[Any], raw_related) if item]
                    body = parts[2].strip()
                except yaml.YAMLError:
                    body = content

        # Get relative path from rules dir
        rel_path = str(file_path.relative_to(self.rules_dir))
        source_info = source_manifest.get(rel_path, {})
        origin = source_info.get("origin")
        editable = source_info.get("editable")
        origin_value = origin if isinstance(origin, str) else "repo"
        editable_value = editable if isinstance(editable, bool) else origin_value == "repo"

        return Rule(
            path=rel_path,
            rule_id=rel_path,
            rule_hash=_compute_rule_hash(
                path=rel_path,
                description=description,
                always_apply=always_apply,
                required=required,
                related=related,
                origin=origin_value,
                editable=editable_value,
                content=body,
            ),
            description=description,
            content=body,
            always_apply=always_apply,
            required=required,
            related=related,
            origin=origin_value,
            editable=editable_value,
        )

    def get_all_rules(self) -> list[Rule]:
        """Get all loaded rules."""
        self.load()
        return self._rules

    def get_always_apply_rules(self) -> list[Rule]:
        """Get rules marked as alwaysApply."""
        self.load()
        return [r for r in self._rules if r.always_apply]

    def get_rule_by_path(self, path: str) -> Rule | None:
        """Get a rule by its path (with or without .mdc extension)."""
        self.load()
        # Normalize path - add .mdc if missing
        if not path.endswith(".mdc"):
            path = f"{path}.mdc"
        for rule in self._rules:
            if rule.path == path:
                return rule
        return None


# Process-wide cache of rule stores keyed by materialized rules directory.
_RULE_STORES: dict[Path, RuleStore] = {}


def get_rule_store(rules_dir: Path) -> RuleStore:
    """Get or create the cached rule store for a materialized rules directory."""
    key = rules_dir.resolve()
    store = _RULE_STORES.get(key)
    if store is None:
        store = RuleStore(key)
        _RULE_STORES[key] = store
    return store
