"""
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


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
        if hasattr(result, "model_dump"):
            result_dict = result.model_dump(mode="json")
        elif hasattr(result, "dict"):
            result_dict = result.dict()
        else:
            result_dict = 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")]
