"""
Storage Manager for NeuroFlow

Handles local file storage for models, runs, and logs.
"""

from __future__ import annotations

import json
import os
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

import joblib
import numpy as np
import pandas as pd


def make_json_serializable(obj: Any) -> Any:
    """Recursively convert non-JSON-serializable objects to serializable types."""
    if obj is None:
        return None
    if isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, pd.DataFrame):
        return obj.to_dict(orient='records')
    if isinstance(obj, pd.Series):
        return obj.tolist()
    if isinstance(obj, dict):
        return {k: make_json_serializable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [make_json_serializable(item) for item in obj]
    if hasattr(obj, 'to_dict'):
        return make_json_serializable(obj.to_dict())
    if hasattr(obj, '__dict__'):
        return make_json_serializable(obj.__dict__)
    # Fallback to string representation
    return str(obj)

class StorageManager:
    """
    Manages local storage for NeuroFlow.
    
    Directory structure:
    - runs/           # Execution run data
      - {run_id}/
        - result.json     # Execution result
        - logs.jsonl      # Event logs
        - artifacts/      # Intermediate artifacts (scalers, encoders)
        - models/         # Trained model files
    - data/           # Uploaded/referenced data files
    - exports/        # Exported workflows
    """
    
    def __init__(self, base_path: str = "./storage"):
        self.base_path = Path(base_path)
        self.runs_path = self.base_path / "runs"
        self.data_path = self.base_path / "data"
        self.exports_path = self.base_path / "exports"
        
        # Ensure directories exist
        self._ensure_directories()
    
    def _ensure_directories(self) -> None:
        """Create necessary directories if they don't exist."""
        self.runs_path.mkdir(parents=True, exist_ok=True)
        self.data_path.mkdir(parents=True, exist_ok=True)
        self.exports_path.mkdir(parents=True, exist_ok=True)
    
    def _get_run_path(self, run_id: str) -> Path:
        """Get the path for a specific run."""
        run_path = self.runs_path / run_id
        run_path.mkdir(parents=True, exist_ok=True)
        return run_path
    
    def _get_artifacts_path(self, run_id: str) -> Path:
        """Get the artifacts path for a run."""
        artifacts_path = self._get_run_path(run_id) / "artifacts"
        artifacts_path.mkdir(parents=True, exist_ok=True)
        return artifacts_path
    
    def _get_models_path(self, run_id: str) -> Path:
        """Get the models path for a run."""
        models_path = self._get_run_path(run_id) / "models"
        models_path.mkdir(parents=True, exist_ok=True)
        return models_path
    
    # ============== Data File Operations ==============
    
    def resolve_data_path(self, file_path: str) -> Path:
        """
        Resolve a data file path.
        
        Search order:
        1. Absolute path
        2. Storage data directory
        3. Bundled data directory (from package)
        4. Current working directory
        5. Relative to cwd/data/
        """
        path = Path(file_path)
        
        if path.is_absolute():
            if path.exists():
                return path
            raise FileNotFoundError(f"Data file not found: {file_path}")
        
        # Try in storage data directory
        data_file = self.data_path / file_path
        if data_file.exists():
            return data_file
        
        # Try in bundled data directory (for pip-installed package)
        bundled_data = Path(__file__).parent.parent / "data" / file_path
        if bundled_data.exists():
            return bundled_data
        
        # Also check NEUROFLOW_DATA_DIR environment variable
        env_data_dir = os.environ.get("NEUROFLOW_DATA_DIR")
        if env_data_dir:
            env_data_file = Path(env_data_dir) / file_path
            if env_data_file.exists():
                return env_data_file
        
        # Try as relative path from current directory
        if path.exists():
            return path
        
        # Try in cwd/data/
        cwd_data_file = Path.cwd() / "data" / file_path
        if cwd_data_file.exists():
            return cwd_data_file
        
        # Try just the filename in case path has directory prefix
        if "/" in file_path:
            filename = Path(file_path).name
            for search_path in [self.data_path, bundled_data.parent, Path.cwd() / "data"]:
                candidate = search_path / filename
                if candidate.exists():
                    return candidate
        
        raise FileNotFoundError(f"Data file not found: {file_path}")
    
    def save_uploaded_file(self, filename: str, content: bytes) -> Path:
        """Save an uploaded file to the data directory."""
        file_path = self.data_path / filename
        file_path.write_bytes(content)
        return file_path
    
    def list_data_files(self) -> list[str]:
        """List all files in the data directory."""
        return [f.name for f in self.data_path.iterdir() if f.is_file()]
    
    # ============== Model Operations ==============
    
    def save_model(self, run_id: str, model_name: str, model: Any) -> str:
        """
        Save a trained model using joblib.
        
        Returns the path to the saved model.
        """
        models_path = self._get_models_path(run_id)
        model_file = models_path / f"{model_name}.joblib"
        
        joblib.dump(model, model_file)
        
        return str(model_file)
    
    def load_model(self, run_id: str, model_name: str) -> Any:
        """Load a trained model."""
        models_path = self._get_models_path(run_id)
        model_file = models_path / f"{model_name}.joblib"
        
        if not model_file.exists():
            raise FileNotFoundError(f"Model not found: {model_file}")
        
        return joblib.load(model_file)
    
    def list_models(self, run_id: str) -> list[str]:
        """List all models for a run."""
        models_path = self._get_models_path(run_id)
        return [f.stem for f in models_path.glob("*.joblib")]
    
    # ============== Artifact Operations ==============
    
    def save_artifact(self, run_id: str, artifact_name: str, artifact: Any) -> str:
        """Save an intermediate artifact (scaler, encoder, etc.)."""
        artifacts_path = self._get_artifacts_path(run_id)
        artifact_file = artifacts_path / f"{artifact_name}.joblib"
        
        joblib.dump(artifact, artifact_file)
        
        return str(artifact_file)
    
    def load_artifact(self, run_id: str, artifact_name: str) -> Any:
        """Load an artifact."""
        artifacts_path = self._get_artifacts_path(run_id)
        artifact_file = artifacts_path / f"{artifact_name}.joblib"
        
        if not artifact_file.exists():
            raise FileNotFoundError(f"Artifact not found: {artifact_file}")
        
        return joblib.load(artifact_file)
    
    # ============== Run Result Operations ==============
    
    def save_run_result(self, run_id: str, result: Any) -> None:
        """Save the execution result for a run."""
        run_path = self._get_run_path(run_id)
        result_file = run_path / "result.json"
        
        # Convert to dict if Pydantic model
        # First try to get dict representation, then sanitize for JSON
        if hasattr(result, "model_dump"):
            try:
                result_dict = result.model_dump(mode="json")
            except Exception:
                # If model_dump fails, manually convert
                result_dict = result.model_dump()
                result_dict = make_json_serializable(result_dict)
        elif hasattr(result, "dict"):
            result_dict = result.dict()
            result_dict = make_json_serializable(result_dict)
        else:
            result_dict = make_json_serializable(result)
        
        with open(result_file, "w") as f:
            json.dump(result_dict, f, indent=2, default=str)
    
    def load_run_result(self, run_id: str) -> dict[str, Any]:
        """Load the execution result for a run."""
        run_path = self._get_run_path(run_id)
        result_file = run_path / "result.json"
        
        if not result_file.exists():
            raise FileNotFoundError(f"Run result not found: {run_id}")
        
        with open(result_file) as f:
            return json.load(f)
    
    def list_runs(self) -> list[dict[str, Any]]:
        """List all runs with basic info."""
        runs = []
        
        for run_dir in self.runs_path.iterdir():
            if run_dir.is_dir():
                result_file = run_dir / "result.json"
                if result_file.exists():
                    try:
                        with open(result_file) as f:
                            result = json.load(f)
                        runs.append({
                            "run_id": run_dir.name,
                            "workflow_id": result.get("workflow_id"),
                            "status": result.get("status"),
                            "started_at": result.get("started_at"),
                            "completed_at": result.get("completed_at"),
                        })
                    except Exception:
                        runs.append({
                            "run_id": run_dir.name,
                            "status": "unknown",
                        })
        
        return sorted(runs, key=lambda x: x.get("started_at", ""), reverse=True)
    
    def delete_run(self, run_id: str) -> bool:
        """Delete a run and all its artifacts."""
        run_path = self._get_run_path(run_id)
        
        if run_path.exists():
            shutil.rmtree(run_path)
            return True
        return False
    
    # ============== Logging Operations ==============
    
    def log_event(
        self,
        run_id: str,
        node_id: str,
        event: str,
        details: Optional[dict[str, Any]] = None,
    ) -> None:
        """Log an event during execution."""
        run_path = self._get_run_path(run_id)
        logs_file = run_path / "logs.jsonl"
        
        log_entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "node_id": node_id,
            "event": event,
            "details": details or {},
        }
        
        with open(logs_file, "a") as f:
            f.write(json.dumps(log_entry, default=str) + "\n")
    
    def get_run_logs(self, run_id: str) -> list[dict[str, Any]]:
        """Get all logs for a run."""
        run_path = self._get_run_path(run_id)
        logs_file = run_path / "logs.jsonl"
        
        if not logs_file.exists():
            return []
        
        logs = []
        with open(logs_file) as f:
            for line in f:
                if line.strip():
                    logs.append(json.loads(line))
        
        return logs
    
    # ============== Workflow Export/Import ==============
    
    def export_workflow(self, workflow_id: str, workflow_data: dict[str, Any]) -> str:
        """Export a workflow to a JSON file."""
        export_file = self.exports_path / f"{workflow_id}.json"
        
        with open(export_file, "w") as f:
            json.dump(workflow_data, f, indent=2, default=str)
        
        return str(export_file)
    
    def import_workflow(self, file_path: str) -> dict[str, Any]:
        """Import a workflow from a JSON file."""
        with open(file_path) as f:
            return json.load(f)
    
    def list_exported_workflows(self) -> list[str]:
        """List all exported workflow files."""
        return [f.stem for f in self.exports_path.glob("*.json")]
