"""
Intelligent Data Enrichment Agent using LLM-driven approach.
"""

import logging
import traceback
import warnings
from typing import Any

import pandas as pd

warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=DeprecationWarning)
pd.options.mode.chained_assignment = None

from sfn_blueprint import BaseLangChainAgent, setup_logger

from .constants import (
    format_feature_engineering_system_prompt,
    format_feature_engineering_user_prompt,
)
from .models import EnrichmentConfig, EnrichmentResponse, FeatureSuggestion, FeatureSuggestionResponse
from .utils import (
    EnrichmentValidator,
    FeatureProfiler,
    extract_dataframe_profile,
    validate_llm_json_output_feature_engineering,
)


class EnrichmentAgent(BaseLangChainAgent):
    """
    Intelligent agent for data enrichment and feature engineering using LLMs.
    """
    def __init__(self, config: EnrichmentConfig | None = None):
        # Initialize configuration
        super().__init__(config or EnrichmentConfig())  # Use default config if none provided
        # Set up logging
        try:
            logger_result = setup_logger(__name__)
            self.logger = logger_result[0] if isinstance(logger_result, tuple) else logger_result
        except Exception as e:
            logging.basicConfig(level=logging.INFO)
            self.logger = logging.getLogger(__name__)
            self.logger.info(f"Using fallback logging due to: {e}")
        
      
        self.agent_type = "enrichment_agent"
        
        # Set log level from config
        
        # Initialize components
        self.feature_profiler = FeatureProfiler()
        self.validator = EnrichmentValidator()
        # Initialize sfn-blueprint components with error handling

    def _parse_input_data(self, data: pd.DataFrame | dict[str, Any] | str) -> pd.DataFrame:
        if isinstance(data, pd.DataFrame):
            return data.copy()
        elif isinstance(data, dict):
            return pd.DataFrame(data)
        elif isinstance(data, str):
            return pd.read_csv(data)
        else:
            raise ValueError("Unsupported input data type.")

    def suggest_features(
        self,
        df: pd.DataFrame,
        user_prompt: str,
        system_prompt: str,
        profile: Any,
        parameters: dict[str, Any] | None = None
    ) -> list[FeatureSuggestion]:
        """
        Generate feature suggestions based on the input data and context.
        
        Args:
            df: Input DataFrame
            user_prompt: User's description of the task
            system_prompt: System instructions for feature generation
            profile: Data profile information
            parameters: Additional parameters for feature generation
            
        Returns:
            List of FeatureSuggestion objects
            
        Raises:
            ValueError: If no LLM handler is available or if feature generation fails
        """
  
        suggestions, cost = self.route_with_langchain(system_prompt, user_prompt, FeatureSuggestionResponse)
        return suggestions.model_dump()
            
      
    def enrich_data(
        self,
        data: pd.DataFrame | dict[str, Any] | str,
        db_name: str = "database",
        schema: str = "public",
        table_name: str = "table",
        domain_metadata: dict[str, Any] | None = None,
        entity_metadata: dict[str, Any] | None = None,
        use_case: str | None = None,
        ml_approach: str | None = None,
        mode: str | None = None,
        sample_rows: int = 5,
        parameters: dict[str, Any] | None = None,
        data_store: str | None = None,
    ) -> EnrichmentResponse:
        """
            Main entry point for data enrichment. Suggests and applies new features using all provided context.
        """
        try:
            df = self._parse_input_data(data)
            # print("\n** parsed data\n", df)
            profile = extract_dataframe_profile(df, sample_size=sample_rows)
            # print("\n** profile\n", profile)
            # Prepare all context for prompt
            system_prompt = format_feature_engineering_system_prompt(data_store=data_store)
            # print("\n** system_prompt\n", system_prompt)


            user_prompt = format_feature_engineering_user_prompt(
            db_name=db_name,
            schema=schema,
            table_name=table_name,
            use_case=use_case or "N/A",
            ml_approach=ml_approach or "N/A",
            domain_metadata=domain_metadata or {},
            entity_metadata=entity_metadata or {},
            column_names=profile["columns"],
            column_dtypes=profile["data_types"],
            computed_stats=profile["computed_stats"],
            sample_data=profile["sample_rows"],
            data_store=data_store,
            )
            # print("\n** user_prompt\n", user_prompt)
            
            # In production, call LLM here. For now, stub with demo suggestions.
            suggestions = self.suggest_features(
                df, user_prompt, system_prompt, profile, parameters
            )


            suggestions_list = suggestions.get("suggestions", [])
            suggestions_dict = {
                item['feature_name']: item 
                for item in suggestions_list
            }
            
            return EnrichmentResponse(
                success=True,
                suggestions=suggestions_dict,
                message="Data enrichment completed successfully.",
                errors=[]
            )
        except Exception as e:
            self.logger.error(f"Failed to enrich data: {e}")
            print(traceback.format_exc()[:500])
            return EnrichmentResponse(
                success=False,
                suggestions=None,
                message="Data enrichment failed.",
                errors=[str(e)]
            )

    def __call__(self,
        data: pd.DataFrame | dict[str, Any] | str,
        db_name: str = "database",
        schema: str = "public",
        table_name: str = "table",
        domain_metadata: dict[str, Any] | None = None,
        entity_metadata: dict[str, Any] | None = None,
        use_case: str | None = None,
        ml_approach: str | None = None,
        mode: str | None = None,
        sample_rows: int = 5,
        parameters: dict[str, Any] | None = None,
        data_store: str | None = None,
        ) -> EnrichmentResponse:
        """
        Main entry point for data enrichment. Suggests and applies new features using all provided context.
        """
        return self.enrich_data(data, db_name, schema,
                                table_name, domain_metadata,
                                entity_metadata, use_case, ml_approach,
                                mode, sample_rows, parameters, data_store)