"""I/O operations tools for CSV Editor MCP Server."""

import pandas as pd
from typing import Dict, Any, Optional, List, Literal
from pathlib import Path
import tempfile
import json
from datetime import datetime

from fastmcp import Context
from ..models import (
    get_session_manager,
    OperationResult,
    OperationType,
    ExportFormat,
    SessionInfo
)
from ..utils.validators import validate_file_path, validate_url, sanitize_filename


async def load_csv(
    file_path: str,
    encoding: str = "utf-8",
    delimiter: str = ",",
    session_id: Optional[str] = None,
    header: Optional[int] = 0,
    na_values: Optional[List[str]] = None,
    parse_dates: Optional[List[str]] = None,
    ctx: Context = None
) -> Dict[str, Any]:
    """Load a CSV file into a session.
    
    Args:
        file_path: Path to the CSV file
        encoding: File encoding (default: utf-8)
        delimiter: Column delimiter (default: comma)
        session_id: Optional existing session ID to use
        header: Row number to use as header (default: 0)
        na_values: Additional strings to recognize as NA/NaN
        parse_dates: Columns to parse as dates
        ctx: FastMCP context
    
    Returns:
        Operation result with session ID and data info
    """
    try:
        # Validate file path
        is_valid, validated_path = validate_file_path(file_path)
        if not is_valid:
            return OperationResult(
                success=False,
                message=f"Invalid file path: {validated_path}",
                error=validated_path
            ).model_dump()
        
        if ctx:
            await ctx.info(f"Loading CSV file: {validated_path}")
            await ctx.report_progress(0.1, "Validating file...")
        
        # Get or create session
        session_manager = get_session_manager()
        session = session_manager.get_or_create_session(session_id)
        
        if ctx:
            await ctx.report_progress(0.3, "Reading file...")
        
        # Read CSV with pandas
        read_params = {
            "filepath_or_buffer": validated_path,
            "encoding": encoding,
            "delimiter": delimiter,
            "header": header,
        }
        
        if na_values:
            read_params["na_values"] = na_values
        if parse_dates:
            read_params["parse_dates"] = parse_dates
            
        df = pd.read_csv(**read_params)
        
        if ctx:
            await ctx.report_progress(0.8, "Processing data...")
        
        # Load into session
        session.load_data(df, validated_path)
        
        if ctx:
            await ctx.report_progress(1.0, "Complete!")
            await ctx.info(f"Loaded {len(df)} rows and {len(df.columns)} columns")
        
        return OperationResult(
            success=True,
            message=f"Successfully loaded CSV file",
            session_id=session.session_id,
            rows_affected=len(df),
            columns_affected=df.columns.tolist(),
            data={
                "shape": df.shape,
                "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
                "memory_usage_mb": df.memory_usage(deep=True).sum() / (1024 * 1024),
                "preview": df.to_dict('records')
            }
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to load CSV: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to load CSV file",
            error=str(e)
        ).model_dump()


async def load_csv_from_url(
    url: str,
    encoding: str = "utf-8",
    delimiter: str = ",",
    session_id: Optional[str] = None,
    ctx: Context = None
) -> Dict[str, Any]:
    """Load a CSV file from a URL.
    
    Args:
        url: URL of the CSV file
        encoding: File encoding
        delimiter: Column delimiter
        session_id: Optional existing session ID
        ctx: FastMCP context
    
    Returns:
        Operation result with session ID and data info
    """
    try:
        # Validate URL
        is_valid, validated_url = validate_url(url)
        if not is_valid:
            return OperationResult(
                success=False,
                message=f"Invalid URL: {validated_url}",
                error=validated_url
            ).model_dump()
        
        if ctx:
            await ctx.info(f"Loading CSV from URL: {url}")
            await ctx.report_progress(0.1, "Downloading file...")
        
        # Download CSV using pandas (it handles URLs directly)
        df = pd.read_csv(url, encoding=encoding, delimiter=delimiter)
        
        if ctx:
            await ctx.report_progress(0.8, "Processing data...")
        
        # Get or create session
        session_manager = get_session_manager()
        session = session_manager.get_or_create_session(session_id)
        session.load_data(df, url)
        
        if ctx:
            await ctx.report_progress(1.0, "Complete!")
            await ctx.info(f"Loaded {len(df)} rows and {len(df.columns)} columns")
        
        return OperationResult(
            success=True,
            message=f"Successfully loaded CSV from URL",
            session_id=session.session_id,
            rows_affected=len(df),
            columns_affected=df.columns.tolist(),
            data={
                "shape": df.shape,
                "source_url": url,
                "preview": df.head(5).to_dict('records')
            }
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to load CSV from URL: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to load CSV from URL",
            error=str(e)
        ).model_dump()


async def load_csv_from_content(
    content: str,
    delimiter: str = ",",
    session_id: Optional[str] = None,
    has_header: bool = True,
    ctx: Context = None
) -> Dict[str, Any]:
    """Load CSV data from a string content.
    
    Args:
        content: CSV content as string
        delimiter: Column delimiter
        session_id: Optional existing session ID
        has_header: Whether first row is header
        ctx: FastMCP context
    
    Returns:
        Operation result with session ID and data info
    """
    try:
        if ctx:
            await ctx.info("Loading CSV from content string")
        
        # Parse CSV from string
        from io import StringIO
        df = pd.read_csv(
            StringIO(content),
            delimiter=delimiter,
            header=0 if has_header else None
        )
        
        # Get or create session
        session_manager = get_session_manager()
        session = session_manager.get_or_create_session(session_id)
        session.load_data(df, None)
        
        if ctx:
            await ctx.info(f"Loaded {len(df)} rows and {len(df.columns)} columns")
        
        return OperationResult(
            success=True,
            message=f"Successfully loaded CSV from content",
            session_id=session.session_id,
            rows_affected=len(df),
            columns_affected=df.columns.tolist(),
            data={
                "shape": df.shape,
                "preview": df.head(5).to_dict('records')
            }
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to parse CSV content: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to parse CSV content",
            error=str(e)
        ).model_dump()


async def export_csv(
    session_id: str,
    file_path: Optional[str] = None,
    format: ExportFormat = ExportFormat.CSV,
    encoding: str = "utf-8",
    index: bool = False,
    ctx: Context = None
) -> Dict[str, Any]:
    """Export session data to various formats.
    
    Args:
        session_id: Session ID to export
        file_path: Optional output file path (auto-generated if not provided)
        format: Export format (csv, tsv, json, excel, parquet, html, markdown)
        encoding: Output encoding
        index: Whether to include index in output
        ctx: FastMCP context
    
    Returns:
        Operation result with file path
    """
    try:
        # Get session
        session_manager = get_session_manager()
        session = session_manager.get_session(session_id)
        
        if not session or session.df is None:
            return OperationResult(
                success=False,
                message="Session not found or no data loaded",
                error="Invalid session ID"
            ).model_dump()
        
        if ctx:
            await ctx.info(f"Exporting data in {format.value} format")
            await ctx.report_progress(0.1, "Preparing export...")
        
        # Generate file path if not provided
        if not file_path:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"export_{session_id[:8]}_{timestamp}"
            
            # Determine extension based on format
            extensions = {
                ExportFormat.CSV: ".csv",
                ExportFormat.TSV: ".tsv",
                ExportFormat.JSON: ".json",
                ExportFormat.EXCEL: ".xlsx",
                ExportFormat.PARQUET: ".parquet",
                ExportFormat.HTML: ".html",
                ExportFormat.MARKDOWN: ".md"
            }
            
            file_path = tempfile.gettempdir() + "/" + filename + extensions[format]
        
        file_path = Path(file_path)
        df = session.df
        
        if ctx:
            await ctx.report_progress(0.5, f"Writing {format.value} file...")
        
        # Export based on format
        if format == ExportFormat.CSV:
            df.to_csv(file_path, encoding=encoding, index=index)
        elif format == ExportFormat.TSV:
            df.to_csv(file_path, sep='\t', encoding=encoding, index=index)
        elif format == ExportFormat.JSON:
            df.to_json(file_path, orient='records', indent=2)
        elif format == ExportFormat.EXCEL:
            df.to_excel(file_path, index=index, engine='openpyxl')
        elif format == ExportFormat.PARQUET:
            df.to_parquet(file_path, index=index)
        elif format == ExportFormat.HTML:
            df.to_html(file_path, index=index)
        elif format == ExportFormat.MARKDOWN:
            df.to_markdown(file_path, index=index)
        else:
            return OperationResult(
                success=False,
                message=f"Unsupported format: {format}",
                error="Invalid export format"
            ).model_dump()
        
        # Record operation
        session.record_operation(
            OperationType.EXPORT,
            {"format": format.value, "file_path": str(file_path)}
        )
        
        if ctx:
            await ctx.report_progress(1.0, "Export complete!")
            await ctx.info(f"Exported to {file_path}")
        
        return OperationResult(
            success=True,
            message=f"Successfully exported data to {format.value}",
            session_id=session_id,
            data={
                "file_path": str(file_path),
                "format": format.value,
                "rows_exported": len(df),
                "file_size_bytes": file_path.stat().st_size
            }
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to export data: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to export data",
            error=str(e)
        ).model_dump()


async def get_session_info(
    session_id: str,
    ctx: Context = None
) -> Dict[str, Any]:
    """Get information about a specific session.
    
    Args:
        session_id: Session ID
        ctx: FastMCP context
    
    Returns:
        Session information
    """
    try:
        session_manager = get_session_manager()
        session = session_manager.get_session(session_id)
        
        if not session:
            return OperationResult(
                success=False,
                message="Session not found",
                error="Invalid session ID"
            ).model_dump()
        
        if ctx:
            await ctx.info(f"Retrieved info for session {session_id}")
        
        info = session.get_info()
        return OperationResult(
            success=True,
            message="Session info retrieved",
            session_id=session_id,
            data=info.model_dump()
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to get session info: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to get session info",
            error=str(e)
        ).model_dump()


async def list_sessions(ctx: Context = None) -> Dict[str, Any]:
    """List all active sessions.
    
    Args:
        ctx: FastMCP context
    
    Returns:
        List of active sessions
    """
    try:
        session_manager = get_session_manager()
        sessions = session_manager.list_sessions()
        
        if ctx:
            await ctx.info(f"Found {len(sessions)} active sessions")
        
        return {
            "success": True,
            "message": f"Found {len(sessions)} active sessions",
            "sessions": [s.model_dump() for s in sessions]
        }
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to list sessions: {str(e)}")
        return {
            "success": False,
            "message": "Failed to list sessions",
            "error": str(e)
        }


async def close_session(
    session_id: str,
    ctx: Context = None
) -> Dict[str, Any]:
    """Close and clean up a session.
    
    Args:
        session_id: Session ID to close
        ctx: FastMCP context
    
    Returns:
        Operation result
    """
    try:
        session_manager = get_session_manager()
        removed = await session_manager.remove_session(session_id)
        
        if not removed:
            return OperationResult(
                success=False,
                message="Session not found",
                error="Invalid session ID"
            ).model_dump()
        
        if ctx:
            await ctx.info(f"Closed session {session_id}")
        
        return OperationResult(
            success=True,
            message=f"Session {session_id} closed successfully",
            session_id=session_id
        ).model_dump()
        
    except Exception as e:
        if ctx:
            await ctx.error(f"Failed to close session: {str(e)}")
        return OperationResult(
            success=False,
            message="Failed to close session",
            error=str(e)
        ).model_dump()