"""Batch processing utilities for DataFrame operations with LangChain."""

import asyncio
import os
from pathlib import Path
from datetime import date, datetime, timezone
from dataclasses import asdict
from typing import Any, Optional, TYPE_CHECKING, Type, List

import pandas as pd
from langchain_core.runnables import RunnableLambda

from .logger import ParquetLogger
from .config import (
    JobConfig, StorageConfig, ProcessingConfig, ColumnConfig,
    S3Config, EventType, LLMConfig
)
from .tagging import with_tags


async def batch_run(
    df: pd.DataFrame,
    llm: Any,
    prompt_col: str = "prompt",
    config_col: str = "config",
    tools_col: Optional[str] = "tools",
    max_concurrency: int = 10,
    show_progress: bool = True,
    return_exceptions: bool = True,
    return_results: bool = True,
    row_timeout: Optional[float] = None,
) -> Optional[list]:
    """
    Low-level async batch processing for DataFrames.

    **Most users should use batch_process() instead**, which includes automatic logging.
    Use batch_run() only when you need direct control over logging.

    Args:
        df: DataFrame with prepared data
        llm: LangChain LLM (already configured with callbacks, structured output, etc.)
        prompt_col: Column name containing prompts (default: "prompt")
        config_col: Column name containing config dicts from with_tags() (default: "config")
        tools_col: Column name containing tools lists (default: "tools", set None to skip)
        max_concurrency: Maximum concurrent requests
        show_progress: Show progress bar (auto-detects notebook vs terminal)
        return_exceptions: Return exceptions instead of raising
        return_results: If False, don't keep results in memory (useful for huge DataFrames)
        row_timeout: Per-row timeout in seconds (None = no timeout). Useful with
            slow service tiers (e.g. OpenAI flex) to prevent indefinite hangs.

    Returns:
        List of results in same order as DataFrame rows, or None if return_results=False

    Example:
        >>> # When you already have logging configured:
        >>> with ParquetLogger('./logs') as logger:
        >>>     llm.callbacks = [logger]
        >>>     results = await batch_run(df, llm, max_concurrency=100)
        >>>
        >>> # For most users, use batch_process() instead:
        >>> results = await batch_process(df)  # Logging handled automatically
    """
    rows = df.to_dict('records')

    # Setup progress tracking if requested
    progress_bar = None
    if show_progress:
        try:
            from tqdm.auto import tqdm
            progress_bar = tqdm(total=len(rows), desc="Processing batch")
        except Exception:
            progress_bar = None
            print(f"Processing {len(rows)} rows...")

    async def process_row(row: dict) -> Any:
        """Process a single row from the DataFrame."""
        # Build invoke kwargs from DataFrame columns
        invoke_kwargs = {"input": row.get(prompt_col)}

        # Add config if column exists
        if config_col in row and row[config_col] is not None:
            invoke_kwargs["config"] = row[config_col]

        # Add tools if column exists and not None
        if tools_col and tools_col in row and row[tools_col] is not None:
            invoke_kwargs["tools"] = row[tools_col]

        try:
            if row_timeout:
                result = await asyncio.wait_for(
                    llm.ainvoke(**invoke_kwargs),
                    timeout=row_timeout
                )
            else:
                result = await llm.ainvoke(**invoke_kwargs)
            if progress_bar:
                progress_bar.update(1)
                progress_bar.refresh()
            return result
        except asyncio.TimeoutError:
            if progress_bar:
                progress_bar.update(1)
                progress_bar.refresh()
            if return_exceptions:
                return TimeoutError(f"Row timed out after {row_timeout}s")
            raise TimeoutError(f"Row timed out after {row_timeout}s")
        except Exception as e:
            if progress_bar:
                progress_bar.update(1)
                progress_bar.refresh()
            if return_exceptions:
                return e
            raise

    # Create runner and process batch
    runner = RunnableLambda(process_row)

    if return_results:
        # Normal mode: collect and return results
        results = await runner.abatch(
            rows,
            config={"max_concurrency": max_concurrency},
            return_exceptions=return_exceptions
        )

        if progress_bar:
            progress_bar.close()

        return results
    else:
        # Memory-efficient mode: don't collect results
        await runner.abatch(
            rows,
            config={"max_concurrency": max_concurrency},
            return_exceptions=return_exceptions
        )

        if progress_bar:
            progress_bar.close()

        return None


async def batch_process(
    df: pd.DataFrame,
    llm_config: LLMConfig,
    job_config: Optional[JobConfig] = None,
    storage_config: Optional[StorageConfig] = None,
    processing_config: Optional[ProcessingConfig] = None,
    column_config: Optional[ColumnConfig] = None,
) -> Optional[List]:
    """
    Batch process DataFrame through LLM with automatic logging to Parquet.

    Args:
        df: DataFrame with prepared data
        llm_config: LLM configuration including class, kwargs, and structured output
        job_config: Job metadata configuration
        storage_config: Storage configuration for output files
        processing_config: Processing configuration for batch operations
        column_config: DataFrame column name configuration

    Returns:
        List of results if processing_config.return_results=True, None otherwise

    Examples:
        >>> # Simple usage with LLMConfig
        >>> from langchain_openai import ChatOpenAI
        >>> await batch_process(
        ...     df,
        ...     llm_config=LLMConfig(
        ...         llm_class=ChatOpenAI,
        ...         llm_kwargs={'model': 'gpt-4', 'temperature': 0.7}
        ...     ),
        ...     job_config=JobConfig(category='analysis', version='2.0')
        ... )

        >>> # With structured output
        >>> from pydantic import BaseModel
        >>> class EmailInfo(BaseModel):
        ...     email: str
        ...     valid: bool
        >>>
        >>> await batch_process(
        ...     df,
        ...     llm_config=LLMConfig(
        ...         llm_class=ChatOpenAI,
        ...         llm_kwargs={'model': 'gpt-4'},
        ...         structured_output=EmailInfo
        ...     ),
        ...     job_config=JobConfig(category='email_validation')
        ... )
    """
    # Initialize configs with defaults if not provided
    job_config = job_config or JobConfig()
    storage_config = storage_config or StorageConfig()
    processing_config = processing_config or ProcessingConfig()
    column_config = column_config or ColumnConfig()

    # Validate DataFrame has required columns
    if column_config.prompt not in df.columns:
        raise ValueError(f"DataFrame missing required column: {column_config.prompt}")

    # LLM will be created inside the logger context with callbacks

    # Format path templates with job metadata
    # Sanitize version for directory names (replace dots with underscores)
    version_safe = (job_config.version or 'unversioned').replace('.', '_')

    template_vars = {
        'job_category': job_config.category,
        'job_subcategory': job_config.subcategory or 'default',
        'environment': job_config.environment or 'production',
        'job_version': job_config.version,        # Keep original for backward compatibility
        'job_version_safe': version_safe,         # New sanitized version for paths
        'date': date.today().isoformat(),
    }

    # Build local output path
    local_path = Path(storage_config.output_dir) / storage_config.path_template.format(**template_vars)
    local_path.mkdir(parents=True, exist_ok=True)

    # Check for S3 bucket in environment if not configured
    if not storage_config.s3_config and os.environ.get('LANGCHAIN_S3_BUCKET'):
        storage_config.s3_config = S3Config(bucket=os.environ['LANGCHAIN_S3_BUCKET'])

    # Format S3 prefix if using S3
    if storage_config.s3_config:
        # Combine S3 prefix with formatted path (mirrors local structure)
        base_prefix = storage_config.s3_config.prefix.rstrip('/')
        formatted_path = storage_config.path_template.format(**template_vars).lstrip('/')
        s3_prefix = f"{base_prefix}/{formatted_path}/" if base_prefix else f"{formatted_path}/"
        storage_config.s3_config.prefix = s3_prefix

    # Build comprehensive logger metadata
    logger_metadata = {
        # Legacy flat fields (for backward compatibility in queries)
        'job_category': job_config.category,
        'job_subcategory': job_config.subcategory,
        'environment': job_config.environment,
        'job_description': job_config.description,
        'job_version': job_config.version,

        # Complete batch-level configs (NEW structure)
        'batch_config': {
            'job': asdict(job_config) if job_config else None,
            'storage': {
                'output_dir': storage_config.output_dir,
                'path_template': storage_config.path_template,
                's3': asdict(storage_config.s3_config) if storage_config.s3_config else None
            },
            'processing': asdict(processing_config) if processing_config else None,
            'column': asdict(column_config) if column_config else None,
            'llm': llm_config.to_metadata_dict(),
        },

        # Batch execution metadata
        'batch_started_at': datetime.now(timezone.utc).isoformat(),
        'batch_size': len(df),

        # Custom metadata from job_config (if any)
        **(job_config.metadata or {})
    }

    # Print status messages
    if processing_config.show_progress:
        print(f"🚀 Starting processing of {len(df)} rows...")
        print(f"📁 Local output: {local_path}")
        if storage_config.s3_config:
            print(f"☁️  S3 upload: s3://{storage_config.s3_config.bucket}/{storage_config.s3_config.prefix}")

    # Process with context manager for automatic cleanup
    with ParquetLogger(
        log_dir=str(local_path),
        buffer_size=processing_config.buffer_size,
        logger_metadata=logger_metadata,
        partition_on=processing_config.partition_on,
        event_types=processing_config.event_types,
        s3_config=storage_config.s3_config
    ) as logger:
        # Create LLM with logger as callback
        llm = llm_config.create_llm(callbacks=[logger])

        # Run batch processing
        results = await batch_run(
            df, llm,
            prompt_col=column_config.prompt,
            config_col=column_config.config,
            tools_col=column_config.tools,
            max_concurrency=processing_config.max_concurrency,
            show_progress=processing_config.show_progress,
            return_exceptions=processing_config.return_exceptions,
            return_results=processing_config.return_results,
            row_timeout=processing_config.row_timeout
        )

    # Print completion message
    if processing_config.show_progress:
        print("✅ Processing complete!")
        print(f"📍 Local files: {local_path}")
        if storage_config.s3_config:
            print(f"☁️  S3 location: s3://{storage_config.s3_config.bucket}/{storage_config.s3_config.prefix}")

    return results

