"""
Feature connection manager for linking pre-extracted features to datasets.

This module provides functionality to:
- Connect feature directories to datasets
- Scan for available features across samples
- Plan extraction for missing features with dependency resolution
- Cache scan results for performance

Configuration is stored in ~/.sltk/feature_configs/{dataset}.json
Scan results are cached in ~/.sltk/feature_cache/{dataset}.json
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# Feature types and their dependencies
FEATURE_DEPENDENCIES: dict[str, list[str]] = {
    "mediapipe": [],  # No dependencies
    "wilor": [],  # No dependencies
    "smpl": [],  # No dependencies (NLF extraction)
    "angles": ["mediapipe"],  # Requires mediapipe body poses
    "hamer": ["wilor"],  # Requires WiLoR MANO params
    "segments": ["angles", "hamer"],  # Requires both angles and hamer
}

# Feature types that can be extracted vs computed
EXTRACTABLE_FEATURES = {"mediapipe", "wilor", "smpl"}
COMPUTED_FEATURES = {"angles", "hamer", "segments"}

# Standard naming patterns
NAMING_PATTERNS = {
    "id_based": "{sample_id}_{feature_type}.h5",
    "gloss_based": "{source}-{gloss}_{feature_type}.h5",
    "flat": "{sample_id}.h5",
}


def get_config_dir() -> Path:
    """Get the SLTK config directory."""
    config_dir = Path.home() / ".sltk" / "feature_configs"
    config_dir.mkdir(parents=True, exist_ok=True)
    return config_dir


def get_cache_dir() -> Path:
    """Get the SLTK cache directory."""
    cache_dir = Path.home() / ".sltk" / "feature_cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


@dataclass
class FeatureSource:
    """
    A feature source directory with naming pattern.

    Attributes:
        path: Absolute path to feature directory
        feature_types: List of feature types in this directory (e.g., ["wilor", "smpl"])
        naming_pattern: Naming pattern for files (e.g., "{sample_id}_{feature_type}.h5")
    """

    path: str
    feature_types: list[str] = field(default_factory=list)
    naming_pattern: str = "{sample_id}_{feature_type}.h5"

    def to_dict(self) -> dict[str, Any]:
        return {
            "path": self.path,
            "feature_types": self.feature_types,
            "naming_pattern": self.naming_pattern,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> FeatureSource:
        return cls(
            path=data["path"],
            feature_types=data.get("feature_types", []),
            naming_pattern=data.get("naming_pattern", "{sample_id}_{feature_type}.h5"),
        )


@dataclass
class FeatureConnectionConfig:
    """
    Configuration for connecting features to a dataset.

    Attributes:
        dataset_name: Name of the dataset
        feature_sources: List of feature source directories
        last_scan: Timestamp of last scan (ISO format)
    """

    dataset_name: str
    feature_sources: list[FeatureSource] = field(default_factory=list)
    last_scan: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "dataset_name": self.dataset_name,
            "feature_sources": [s.to_dict() for s in self.feature_sources],
            "last_scan": self.last_scan,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> FeatureConnectionConfig:
        return cls(
            dataset_name=data["dataset_name"],
            feature_sources=[FeatureSource.from_dict(s) for s in data.get("feature_sources", [])],
            last_scan=data.get("last_scan"),
        )

    def save(self) -> Path:
        """Save config to disk."""
        config_path = get_config_dir() / f"{self.dataset_name}.json"
        with open(config_path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)
        return config_path

    @classmethod
    def load(cls, dataset_name: str) -> FeatureConnectionConfig | None:
        """Load config from disk."""
        config_path = get_config_dir() / f"{dataset_name}.json"
        if not config_path.exists():
            return None
        with open(config_path) as f:
            return cls.from_dict(json.load(f))


@dataclass
class FeatureAvailability:
    """
    Feature availability for a single sample.

    Attributes:
        sample_id: Sample identifier
        available: Dict of feature_type -> file path for available features
        missing: List of feature types that are missing
    """

    sample_id: str
    available: dict[str, str] = field(default_factory=dict)
    missing: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "sample_id": self.sample_id,
            "available": self.available,
            "missing": self.missing,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> FeatureAvailability:
        return cls(
            sample_id=data["sample_id"],
            available=data.get("available", {}),
            missing=data.get("missing", []),
        )


@dataclass
class FeatureSummary:
    """
    Summary of feature availability across a dataset.

    Attributes:
        dataset_name: Name of the dataset
        total_samples: Total number of samples scanned
        feature_counts: Dict of feature_type -> count of samples with that feature
        last_scan: Timestamp of scan
    """

    dataset_name: str
    total_samples: int
    feature_counts: dict[str, int] = field(default_factory=dict)
    last_scan: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "dataset_name": self.dataset_name,
            "total_samples": self.total_samples,
            "feature_counts": self.feature_counts,
            "last_scan": self.last_scan,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> FeatureSummary:
        return cls(
            dataset_name=data["dataset_name"],
            total_samples=data.get("total_samples", 0),
            feature_counts=data.get("feature_counts", {}),
            last_scan=data.get("last_scan"),
        )


@dataclass
class ExtractionStep:
    """A single step in an extraction plan."""

    order: int
    feature_type: str
    sample_ids: list[str]
    depends_on: list[str]
    is_extraction: bool  # True if extraction, False if computed from other features

    def to_dict(self) -> dict[str, Any]:
        return {
            "order": self.order,
            "feature_type": self.feature_type,
            "sample_count": len(self.sample_ids),
            "sample_ids": self.sample_ids[:100],  # Limit for API response
            "depends_on": self.depends_on,
            "is_extraction": self.is_extraction,
        }


@dataclass
class ExtractionPlan:
    """
    Plan for extracting missing features.

    Includes dependency resolution and step ordering.
    """

    target_features: list[str]
    steps: list[ExtractionStep]
    total_samples: int
    total_extractions: int

    def to_dict(self) -> dict[str, Any]:
        return {
            "target_features": self.target_features,
            "steps": [s.to_dict() for s in self.steps],
            "total_samples": self.total_samples,
            "total_extractions": self.total_extractions,
        }


class FeatureConnectionManager:
    """
    Manager for connecting and scanning features for datasets.

    Example:
        >>> manager = FeatureConnectionManager()
        >>> # Connect a feature source
        >>> source = FeatureSource(
        ...     path="/vol/research/SignFeaturePool/features2/BSLCP2/wilor",
        ...     feature_types=["wilor"],
        ...     naming_pattern="{sample_id}_wilor.h5"
        ... )
        >>> config = manager.connect("BSLCP", [source])
        >>> # Scan for features
        >>> availability = manager.scan("BSLCP", sample_ids)
        >>> # Plan extraction for missing features
        >>> plan = manager.plan_extraction("BSLCP", ["segments"], availability)
    """

    def __init__(self):
        self._configs: dict[str, FeatureConnectionConfig] = {}
        self._cache: dict[str, dict[str, FeatureAvailability]] = {}

    def connect(
        self,
        dataset_name: str,
        sources: list[FeatureSource],
    ) -> FeatureConnectionConfig:
        """
        Configure feature sources for a dataset.

        Args:
            dataset_name: Name of the dataset
            sources: List of FeatureSource configurations

        Returns:
            FeatureConnectionConfig that was saved
        """
        config = FeatureConnectionConfig(
            dataset_name=dataset_name,
            feature_sources=sources,
            last_scan=None,
        )
        config.save()
        self._configs[dataset_name] = config
        return config

    def get_config(self, dataset_name: str) -> FeatureConnectionConfig | None:
        """Get the feature connection config for a dataset."""
        if dataset_name in self._configs:
            return self._configs[dataset_name]
        config = FeatureConnectionConfig.load(dataset_name)
        if config:
            self._configs[dataset_name] = config
        return config

    def scan(
        self,
        dataset_name: str,
        sample_ids: list[str],
        feature_types: list[str] | None = None,
    ) -> dict[str, FeatureAvailability]:
        """
        Scan for feature availability across samples.

        Args:
            dataset_name: Name of the dataset
            sample_ids: List of sample IDs to scan
            feature_types: Feature types to check (default: all known types)

        Returns:
            Dict of sample_id -> FeatureAvailability
        """
        config = self.get_config(dataset_name)
        if config is None:
            raise ValueError(f"No feature configuration found for dataset: {dataset_name}")

        if feature_types is None:
            feature_types = list(FEATURE_DEPENDENCIES.keys())

        from sltk.data.feature_discovery import discover_features

        results: dict[str, FeatureAvailability] = {}

        for sample_id in sample_ids:
            available: dict[str, str] = {}

            # Check each feature source
            for source in config.feature_sources:
                source_path = Path(source.path)
                if not source_path.exists():
                    logger.warning(f"Feature source path does not exist: {source_path}")
                    continue

                # Discover features for this sample
                found = discover_features(
                    sample_id=sample_id,
                    source_path=source_path,
                    feature_types=source.feature_types,
                    naming_pattern=source.naming_pattern,
                )
                available.update(found)

            # Determine missing features
            missing = [ft for ft in feature_types if ft not in available]

            results[sample_id] = FeatureAvailability(
                sample_id=sample_id,
                available=available,
                missing=missing,
            )

        # Update config with scan timestamp
        config.last_scan = datetime.now().isoformat()
        config.save()

        # Cache results
        self._cache[dataset_name] = results
        self._save_cache(dataset_name, results)

        return results

    def get_sample_features(
        self,
        dataset_name: str,
        sample_id: str,
    ) -> dict[str, str]:
        """
        Get all feature file paths for a specific sample.

        Args:
            dataset_name: Name of the dataset
            sample_id: Sample identifier

        Returns:
            Dict of feature_type -> file path
        """
        # Check cache first
        if dataset_name in self._cache and sample_id in self._cache[dataset_name]:
            return self._cache[dataset_name][sample_id].available

        # Try loading from disk cache
        cache = self._load_cache(dataset_name)
        if cache and sample_id in cache:
            return cache[sample_id].available

        # Scan just this sample
        config = self.get_config(dataset_name)
        if config is None:
            return {}

        results = self.scan(dataset_name, [sample_id])
        return results.get(sample_id, FeatureAvailability(sample_id=sample_id)).available

    def get_summary(self, dataset_name: str) -> FeatureSummary | None:
        """
        Get a summary of feature availability from cache.

        Returns None if no scan has been performed.
        """
        cache = self._load_cache(dataset_name)
        if not cache:
            return None

        config = self.get_config(dataset_name)
        feature_counts: dict[str, int] = {}

        for availability in cache.values():
            for feature_type in availability.available:
                feature_counts[feature_type] = feature_counts.get(feature_type, 0) + 1

        return FeatureSummary(
            dataset_name=dataset_name,
            total_samples=len(cache),
            feature_counts=feature_counts,
            last_scan=config.last_scan if config else None,
        )

    def plan_extraction(
        self,
        dataset_name: str,
        target_features: list[str],
        availability: dict[str, FeatureAvailability] | None = None,
    ) -> ExtractionPlan:
        """
        Create an extraction plan for missing features with dependency resolution.

        Args:
            dataset_name: Name of the dataset
            target_features: List of feature types to extract
            availability: Pre-computed availability (loads from cache if None)

        Returns:
            ExtractionPlan with ordered steps
        """
        if availability is None:
            availability = self._load_cache(dataset_name)
            if not availability:
                raise ValueError(f"No scan data available for {dataset_name}. Run scan first.")

        # Resolve all required features including dependencies
        required = self._resolve_dependencies(target_features)

        # Build extraction steps
        steps: list[ExtractionStep] = []
        order = 1

        for feature_type in required:
            # Find samples missing this feature
            missing_samples = [
                sample_id
                for sample_id, avail in availability.items()
                if feature_type not in avail.available
            ]

            if not missing_samples:
                continue

            steps.append(
                ExtractionStep(
                    order=order,
                    feature_type=feature_type,
                    sample_ids=missing_samples,
                    depends_on=FEATURE_DEPENDENCIES.get(feature_type, []),
                    is_extraction=feature_type in EXTRACTABLE_FEATURES,
                )
            )
            order += 1

        total_extractions = sum(len(s.sample_ids) for s in steps)

        return ExtractionPlan(
            target_features=target_features,
            steps=steps,
            total_samples=len(availability),
            total_extractions=total_extractions,
        )

    def _resolve_dependencies(self, target_features: list[str]) -> list[str]:
        """
        Resolve feature dependencies and return ordered list.

        Returns features in order such that dependencies come before dependents.
        """
        resolved: list[str] = []
        seen: set[str] = set()

        def resolve(feature: str):
            if feature in seen:
                return
            seen.add(feature)

            # Resolve dependencies first
            for dep in FEATURE_DEPENDENCIES.get(feature, []):
                resolve(dep)

            resolved.append(feature)

        for feature in target_features:
            resolve(feature)

        return resolved

    def _save_cache(
        self,
        dataset_name: str,
        availability: dict[str, FeatureAvailability],
    ):
        """Save scan results to disk cache."""
        cache_path = get_cache_dir() / f"{dataset_name}.json"
        data = {sample_id: avail.to_dict() for sample_id, avail in availability.items()}
        with open(cache_path, "w") as f:
            json.dump(data, f)

    def _load_cache(self, dataset_name: str) -> dict[str, FeatureAvailability] | None:
        """Load scan results from disk cache."""
        if dataset_name in self._cache:
            return self._cache[dataset_name]

        cache_path = get_cache_dir() / f"{dataset_name}.json"
        if not cache_path.exists():
            return None

        with open(cache_path) as f:
            data = json.load(f)

        cache = {
            sample_id: FeatureAvailability.from_dict(avail_data)
            for sample_id, avail_data in data.items()
        }
        self._cache[dataset_name] = cache
        return cache


# Module-level instance for convenience
_manager: FeatureConnectionManager | None = None


def get_feature_manager() -> FeatureConnectionManager:
    """Get the global feature connection manager instance."""
    global _manager
    if _manager is None:
        _manager = FeatureConnectionManager()
    return _manager
