"""
Data tools for Stats Compass MCP server.

Handles data loading, listing, and management.
"""

from typing import Optional

import pandas as pd
from fastmcp import Context, FastMCP
from stats_compass_core import data as core_data
from stats_compass_core.data.load_dataset import LoadDatasetInput

from stats_compass_mcp.session import SessionManager, get_session


def register_data_tools(mcp: FastMCP, session_manager: SessionManager, storage=None):
    """Register all data management tools with the FastMCP server."""

    @mcp.tool()
    def ping() -> dict:
        """Health check - verify server is running."""
        return {
            "status": "ok",
            "server": "stats-compass",
            "message": "Server is running. Sessions are created automatically."
        }

    @mcp.tool()
    def session_info(ctx: Context) -> dict:
        """
        Get information about your current session.
        
        Returns:
            Session info including session_id, created_at, dataframes, models.
        """
        session = get_session(ctx, session_manager)
        return session.get_info()

    @mcp.tool()
    def list_dataframes(ctx: Context) -> dict:
        """
        List all DataFrames in your session.
        
        Returns:
            List of DataFrames with name, shape, columns, and active status.
        """
        session = get_session(ctx, session_manager)

        dataframes = session.state.list_dataframes()
        active = session.state.get_active()
        active_name = active.name if active is not None else None

        return {
            "dataframes": [
                {
                    "name": df.name,
                    "shape": list(df.shape),
                    "columns": list(df.columns),
                    "is_active": df.name == active_name
                }
                for df in dataframes
            ],
            "active_dataframe": active_name,
            "count": len(dataframes)
        }

    @mcp.tool()
    def load_dataset(
        ctx: Context,
        name: str,
        set_active: bool = True
    ) -> dict:
        """
        Load a built-in sample dataset.
        
        Available datasets: TATASTEEL, Housing, Bukayo_Saka_7322
        
        Args:
            name: Dataset name
            set_active: Whether to set as active DataFrame (default: True)
        
        Returns:
            DataFrame info with name, shape, columns.
        """
        session = get_session(ctx, session_manager)

        params = LoadDatasetInput(name=name, set_active=set_active)
        result = core_data.load_dataset(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def load_csv(
        ctx: Context,
        path: str,
        name: Optional[str] = None,
        delimiter: str = ",",
        encoding: str = "utf-8",
        set_active: bool = True
    ) -> dict:
        """
        Load a CSV file from a local path.
        
        Args:
            path: Absolute path to the CSV file. Supports ~ expansion.
            name: Name for the DataFrame (default: filename without extension)
            delimiter: Field delimiter (default: comma)
            encoding: File encoding (default: utf-8)
            set_active: Whether to set as active DataFrame (default: True)
        
        Returns:
            DataFrame info with name, shape, columns, dtypes.
        """
        session = get_session(ctx, session_manager)

        from stats_compass_core.data.load_csv import LoadCSVInput
        from stats_compass_core.data.load_csv import load_csv as core_load_csv
        params = LoadCSVInput(
            path=path,
            name=name,
            delimiter=delimiter,
            encoding=encoding,
            set_active=set_active
        )
        result = core_load_csv(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def load_excel(
        ctx: Context,
        path: str,
        name: Optional[str] = None,
        sheet_name: Optional[str] = None,
        set_active: bool = True
    ) -> dict:
        """
        Load an Excel file from a local path.
        
        Args:
            path: Absolute path to the Excel file. Supports ~ expansion.
            name: Name for the DataFrame (default: filename without extension)
            sheet_name: Sheet to load (default: first sheet)
            set_active: Whether to set as active DataFrame (default: True)
        
        Returns:
            DataFrame info with name, shape, columns, dtypes.
        """
        session = get_session(ctx, session_manager)

        from stats_compass_core.data.load_excel import LoadExcelInput
        from stats_compass_core.data.load_excel import load_excel as core_load_excel
        params = LoadExcelInput(
            path=path,
            name=name,
            sheet_name=sheet_name,
            set_active=set_active
        )
        result = core_load_excel(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def list_files(
        ctx: Context,
        directory: str = "."
    ) -> dict:
        """
        List files in a directory. Useful for finding data files.
        
        Args:
            directory: Directory path. Supports ~ expansion (e.g., ~/Downloads).
        
        Returns:
            List of files in the directory.
        """
        session = get_session(ctx, session_manager)
        from stats_compass_core.data.list_files import ListFilesInput
        from stats_compass_core.data.list_files import list_files as core_list_files
        params = ListFilesInput(directory=directory)
        result = core_list_files(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def get_sample(
        ctx: Context,
        dataframe_name: Optional[str] = None,
        n: int = 10,
        method: str = "head"
    ) -> dict:
        """
        Get sample rows from a DataFrame.
        
        Args:
            dataframe_name: Name of DataFrame (default: active)
            n: Number of rows (default: 10)
            method: 'head', 'tail', or 'random'
        
        Returns:
            Sample rows as records.
        """
        session = get_session(ctx, session_manager)

        from stats_compass_core.data.get_sample import GetSampleInput
        from stats_compass_core.data.get_sample import get_sample as core_get_sample
        params = GetSampleInput(dataframe_name=dataframe_name, n=n, method=method)
        result = core_get_sample(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def get_schema(
        ctx: Context,
        dataframe_name: Optional[str] = None,
        sample_values: int = 3
    ) -> dict:
        """
        Get the schema and metadata of a DataFrame.
        
        Args:
            dataframe_name: Name of DataFrame (default: active)
            sample_values: Number of sample values per column
        
        Returns:
            Schema with columns, dtypes, nulls, and sample values.
        """
        session = get_session(ctx, session_manager)

        from stats_compass_core.data.get_schema import GetSchemaInput
        from stats_compass_core.data.get_schema import get_schema as core_get_schema
        params = GetSchemaInput(dataframe_name=dataframe_name, sample_values=sample_values)
        result = core_get_schema(state=session.state, params=params)
        return result.model_dump()

    @mcp.tool()
    def save_csv(
        ctx: Context,
        dataframe_name: str,
        filepath: str,
        index: bool = False
    ) -> dict:
        """
        Save a DataFrame to a CSV file.
        
        Args:
            dataframe_name: Name of the DataFrame to save
            filepath: Path where the CSV file will be saved (filename only for remote mode)
            index: Whether to write row index (default: False)
        
        Returns:
            Save result with filepath and download_url (if remote).
        """
        session = get_session(ctx, session_manager)

        # Extract just the filename for export path
        from pathlib import Path as PathLib
        filename = PathLib(filepath).name
        if not filename.endswith('.csv'):
            filename = f"{filename}.csv"

        # Use session exports directory
        export_path = session.export_path("data", filename)

        from stats_compass_core.data.save_csv import SaveCSVInput
        from stats_compass_core.data.save_csv import save_csv as core_save_csv
        input_data = SaveCSVInput(dataframe_name=dataframe_name, filepath=str(export_path), index=index)
        result = core_save_csv(state=session.state, input_data=input_data)

        # Add download URL if available
        result_dict = result if isinstance(result, dict) else result.model_dump()
        download_url = session.download_url("data", filename)
        if download_url:
            result_dict["download_url"] = download_url

        return result_dict

    @mcp.tool()
    def save_model(
        ctx: Context,
        model_id: str,
        filepath: str
    ) -> dict:
        """
        Save a trained model to a file.
        
        Args:
            model_id: ID of the model to save
            filepath: Path where the model will be saved (filename only for remote mode)
        
        Returns:
            Save result with filepath and download_url (if remote).
        """
        session = get_session(ctx, session_manager)

        # Extract just the filename for export path
        from pathlib import Path as PathLib
        filename = PathLib(filepath).name
        if not filename.endswith('.joblib'):
            filename = f"{filename}.joblib"

        # Use session exports directory
        export_path = session.export_path("models", filename)

        from stats_compass_core.ml.save_model import SaveModelInput
        from stats_compass_core.ml.save_model import save_model as core_save_model
        input_data = SaveModelInput(model_id=model_id, filepath=str(export_path))
        result = core_save_model(state=session.state, input_data=input_data)

        # Add download URL if available
        result_dict = result if isinstance(result, dict) else result
        download_url = session.download_url("models", filename)
        if download_url:
            result_dict["download_url"] = download_url

        return result_dict

    @mcp.tool()
    def delete_session(ctx: Context) -> dict:
        """
        Delete your current session and all its data.
        
        Returns:
            Deletion result.
        """
        session = get_session(ctx, session_manager)
        session_id = session.session_id

        files_deleted = 0
        if storage:
            files_deleted = storage.delete_session_files(session_id)
        session_deleted = session_manager.delete(session_id)

        return {
            "success": session_deleted,
            "files_deleted": files_deleted,
            "message": "Session deleted" if session_deleted else "Session not found"
        }

    @mcp.tool()
    def server_stats() -> dict:
        """
        Get server statistics (admin tool).
        
        Returns:
            Active sessions count, configuration, and session details.
        """
        return session_manager.get_stats()

    # Remote-only tools (only if storage is provided)
    if storage is not None:
        @mcp.tool()
        def get_upload_url(
            ctx: Context,
            filename: str,
            content_type: str = "text/csv"
        ) -> dict:
            """
            Get a presigned URL for uploading a file.
            
            For S3 storage: Returns a presigned PUT URL.
            For local storage: Returns a file path.
            
            Args:
                filename: Desired filename (e.g., "my_data.csv")
                content_type: MIME type (default: text/csv)
            
            Returns:
                Upload info with url, method, headers, file_key.
            """
            session = get_session(ctx, session_manager)

            return storage.get_upload_url(
                session_id=session.session_id,
                filename=filename,
                content_type=content_type
            )

        @mcp.tool()
        def register_uploaded_file(
            ctx: Context,
            file_key: str,
            dataframe_name: Optional[str] = None,
            file_type: str = "csv"
        ) -> dict:
            """
            Register an uploaded file and load it as a DataFrame.
            
            After uploading to the URL from get_upload_url(), call this
            to load the file into your session.
            
            Args:
                file_key: The file_key returned from get_upload_url()
                dataframe_name: Name for the DataFrame (default: filename without extension)
                file_type: File type - "csv" or "excel"
            
            Returns:
                DataFrame info with name, shape, columns, dtypes.
            """
            session = get_session(ctx, session_manager)

            if not storage.file_exists(session.session_id, file_key):
                return {"error": f"File not found: {file_key}. Did you upload it?"}

            file_path = storage.get_file_path(session.session_id, file_key)

            # Determine DataFrame name
            if not dataframe_name:
                dataframe_name = file_key.rsplit(".", 1)[0]

            # Load file
            try:
                if file_type.lower() == "excel":
                    df = pd.read_excel(file_path)
                else:
                    df = pd.read_csv(file_path)
            except Exception as e:
                return {"error": f"Failed to load file: {str(e)}"}

            # Register in session
            session.state.set_dataframe(df, name=dataframe_name, operation="upload")

            return {
                "success": True,
                "dataframe_name": dataframe_name,
                "shape": list(df.shape),
                "columns": list(df.columns),
                "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
            }
