"""
AdamOps Model Registry Module

Provides model versioning, metadata tracking, and comparison.
"""

import json
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
import joblib

from adamops.utils.logging import get_logger
from adamops.utils.config import get_config

logger = get_logger(__name__)


class ModelVersion:
    """Represents a model version."""
    
    def __init__(self, name: str, version: str, path: str, metadata: Dict):
        self.name = name
        self.version = version
        self.path = path
        self.metadata = metadata
        self.created_at = metadata.get("created_at", datetime.now().isoformat())
    
    def to_dict(self) -> Dict:
        return {
            "name": self.name, "version": self.version, "path": self.path,
            "metadata": self.metadata, "created_at": self.created_at
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> "ModelVersion":
        return cls(data["name"], data["version"], data["path"], data["metadata"])


class ModelRegistry:
    """Model registry for versioning and tracking."""
    
    def __init__(self, path: Optional[str] = None, backend: str = "json"):
        """
        Initialize registry.
        
        Args:
            path: Registry storage path.
            backend: 'json' or 'sqlite'.
        """
        config = get_config()
        self.path = Path(path or config.registry_path)
        self.backend = backend
        self.path.mkdir(parents=True, exist_ok=True)
        
        if backend == "sqlite":
            self._init_sqlite()
        else:
            self._init_json()
    
    def _init_json(self):
        self.registry_file = self.path / "registry.json"
        if not self.registry_file.exists():
            self._save_json({})
    
    def _init_sqlite(self):
        self.db_file = self.path / "registry.db"
        conn = sqlite3.connect(self.db_file)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS models (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT, version TEXT, path TEXT,
                metadata TEXT, created_at TEXT,
                UNIQUE(name, version)
            )
        """)
        conn.commit()
        conn.close()
    
    def _load_json(self) -> Dict:
        with open(self.registry_file) as f:
            return json.load(f)
    
    def _save_json(self, data: Dict):
        with open(self.registry_file, "w") as f:
            json.dump(data, f, indent=2)
    
    def register(self, name: str, model: Any, metadata: Optional[Dict] = None,
                 version: Optional[str] = None) -> ModelVersion:
        """
        Register a model.
        
        Args:
            name: Model name.
            model: Model object to save.
            metadata: Model metadata (algorithm, metrics, etc).
            version: Version string (auto-incremented if None).
        
        Returns:
            ModelVersion: Registered model version.
        """
        # Get next version
        if version is None:
            existing = self.list_versions(name)
            version = f"v{len(existing) + 1}"
        
        # Save model
        model_path = self.path / "models" / name / f"{version}.joblib"
        model_path.parent.mkdir(parents=True, exist_ok=True)
        joblib.dump(model, model_path)
        
        # Create version record
        meta = metadata or {}
        meta["created_at"] = datetime.now().isoformat()
        model_version = ModelVersion(name, version, str(model_path), meta)
        
        # Store in registry
        if self.backend == "json":
            data = self._load_json()
            if name not in data:
                data[name] = {}
            data[name][version] = model_version.to_dict()
            self._save_json(data)
        else:
            conn = sqlite3.connect(self.db_file)
            conn.execute(
                "INSERT OR REPLACE INTO models (name, version, path, metadata, created_at) VALUES (?, ?, ?, ?, ?)",
                (name, version, str(model_path), json.dumps(meta), meta["created_at"])
            )
            conn.commit()
            conn.close()
        
        logger.info(f"Registered model {name} version {version}")
        return model_version
    
    def get(self, name: str, version: str = "latest") -> Optional[ModelVersion]:
        """Get a model version."""
        if version == "latest":
            versions = self.list_versions(name)
            if not versions:
                return None
            version = versions[-1].version
        
        if self.backend == "json":
            data = self._load_json()
            if name in data and version in data[name]:
                return ModelVersion.from_dict(data[name][version])
        else:
            conn = sqlite3.connect(self.db_file)
            cur = conn.execute(
                "SELECT name, version, path, metadata FROM models WHERE name=? AND version=?",
                (name, version)
            )
            row = cur.fetchone()
            conn.close()
            if row:
                return ModelVersion(row[0], row[1], row[2], json.loads(row[3]))
        
        return None
    
    def load(self, name: str, version: str = "latest") -> Any:
        """Load a model."""
        model_version = self.get(name, version)
        if model_version is None:
            raise ValueError(f"Model {name} version {version} not found")
        return joblib.load(model_version.path)
    
    def list_versions(self, name: str) -> List[ModelVersion]:
        """List all versions of a model."""
        if self.backend == "json":
            data = self._load_json()
            if name not in data:
                return []
            return [ModelVersion.from_dict(v) for v in data[name].values()]
        else:
            conn = sqlite3.connect(self.db_file)
            cur = conn.execute(
                "SELECT name, version, path, metadata FROM models WHERE name=? ORDER BY created_at",
                (name,)
            )
            versions = [ModelVersion(r[0], r[1], r[2], json.loads(r[3])) for r in cur.fetchall()]
            conn.close()
            return versions
    
    def list_models(self) -> List[str]:
        """List all registered model names."""
        if self.backend == "json":
            return list(self._load_json().keys())
        else:
            conn = sqlite3.connect(self.db_file)
            cur = conn.execute("SELECT DISTINCT name FROM models")
            names = [r[0] for r in cur.fetchall()]
            conn.close()
            return names
    
    def compare(self, name: str, metric: str = "accuracy") -> Dict:
        """Compare all versions of a model by a metric."""
        versions = self.list_versions(name)
        comparison = []
        for v in versions:
            metrics = v.metadata.get("metrics", {})
            comparison.append({
                "version": v.version,
                "created_at": v.created_at,
                metric: metrics.get(metric),
                **{k: val for k, val in metrics.items() if k != metric}
            })
        return sorted(comparison, key=lambda x: x.get(metric, 0) or 0, reverse=True)
    
    def delete(self, name: str, version: Optional[str] = None):
        """Delete a model or version."""
        if version:
            model_version = self.get(name, version)
            if model_version and Path(model_version.path).exists():
                Path(model_version.path).unlink()
            
            if self.backend == "json":
                data = self._load_json()
                if name in data and version in data[name]:
                    del data[name][version]
                    self._save_json(data)
            else:
                conn = sqlite3.connect(self.db_file)
                conn.execute("DELETE FROM models WHERE name=? AND version=?", (name, version))
                conn.commit()
                conn.close()
        else:
            # Delete all versions
            for v in self.list_versions(name):
                self.delete(name, v.version)


# Global registry instance
_registry: Optional[ModelRegistry] = None

def get_registry() -> ModelRegistry:
    """Get global registry instance."""
    global _registry
    if _registry is None:
        _registry = ModelRegistry()
    return _registry

def register_model(name: str, model: Any, **kwargs) -> ModelVersion:
    """Register a model in the global registry."""
    return get_registry().register(name, model, **kwargs)

def load_model(name: str, version: str = "latest") -> Any:
    """Load a model from the global registry."""
    return get_registry().load(name, version)
