from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import pandas as pd
import gymnasium as gym
from rl_recommender.logger import RecommenderLogger
from rl_recommender.simulation import MABSimulation
from rl_recommender.UnifiedRecommendationEnv import UnifiedRecommendationEnv
from db import DataController, DataManager

app = FastAPI(title="MAB Recommendation API", version="1.0.0")

logger = RecommenderLogger(name="api", log_dir="/tmp/logs", enable_file=False)

# Global simulation instance
simulation = None

# Request/Response models
class RecommendationRequest(BaseModel):
    user_id: int
    n_recommendations: Optional[int] = 3
    algorithm_type: Optional[str] = "epsilon_greedy"
    algorithm_params: Optional[Dict[str, Any]] = {}

class FeedbackRequest(BaseModel):
    user_id: int
    item_id: int
    event_type: str
    timestamp: Optional[str] = None

class RecommendationResponse(BaseModel):
    user_id: int
    recommended_items: List[int]
    algorithm_type: str
    confidence_scores: Optional[List[float]] = None

class FeedbackResponse(BaseModel):
    status: str
    message: str
    updated_rewards: Optional[Dict[str, Any]] = None

def initialize_simulation():
    """Initialize the MAB simulation environment."""
    global simulation
    try:
        # Prefer DB-backed environment; fallback to simple CSV-like mock
        try:
            import os
            db_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "enhanced_recommender.db")
            env = UnifiedRecommendationEnv(
                data_source={"db_path": db_path},
                max_arms=50,
                n_suggestions=3,
                seed=42,
                use_user_context=True
            )
        except Exception as env_error:
            logger.warning(f"Failed to create DB environment: {str(env_error)}")
            # Minimal mock environment
            class MockEnvironment:
                def __init__(self):
                    self.max_arms = 50
                    self.n_suggestions = 3
                    self.n_products = 50
                def reset(self):
                    return {"user_id": 0}, {}
                def step(self, action):
                    return {}, 0.0, False, False, {"event_types": []}
            env = MockEnvironment()
        
        # Create simulation with default algorithm
        simulation = MABSimulation(env, algorithm_type="epsilon_greedy", algorithm_params={"epsilon": 0.1})
        
        logger.info("MAB Simulation initialized successfully")
        return True
    except Exception as e:
        logger.error(f"Failed to initialize simulation: {str(e)}")
        return False

@app.on_event("startup")
async def startup_event():
    """Initialize simulation on startup."""
    logger.info("Starting API initialization...")
    if not initialize_simulation():
        logger.error("Failed to initialize simulation on startup")
        # Try to initialize with a simpler approach
        logger.info("Attempting fallback initialization...")
        try:
            global simulation
            # Create a minimal simulation without environment for basic functionality
            from rl_recommender.mabalgorithms import EpsilonGreedyMAB
            from rl_recommender.simulation import MABSimulation
            
            # Create a mock environment class
            class MockEnvironment:
                def __init__(self):
                    self.max_arms = 50
                    self.n_suggestions = 3
                    self.n_products = 50
                
                def reset(self):
                    return {"user_id": 0}, {}
                
                def step(self, action):
                    return {}, 0.0, False, False, {"event_types": []}
            
            env = MockEnvironment()
            simulation = MABSimulation(env, algorithm_type="epsilon_greedy", algorithm_params={"epsilon": 0.1})
            logger.info("Fallback simulation initialized successfully")
        except Exception as fallback_error:
            logger.error(f"Fallback initialization also failed: {str(fallback_error)}")
            simulation = None

@app.get("/")
def root():
    """Root endpoint with API information."""
    return {
        "message": "MAB Recommendation API",
        "version": "1.0.0",
        "endpoints": {
            "GET /recommend_item": "Get recommendations for a user",
            "POST /feedback": "Submit feedback for recommendations",
            "GET /status": "Get API and simulation status",
            "POST /configure": "Configure algorithm parameters"
        }
    }

@app.get("/status")
def get_status():
    """Get the current status of the API and simulation."""
    global simulation
    
    if simulation is None:
        return {
            "status": "error",
            "message": "Simulation not initialized",
            "simulation_ready": False
        }
    
    try:
        algorithm_info = {
            "type": simulation.algorithm_type,
            "class": simulation.current_algorithm.__class__.__name__,
            "n_arms": simulation.current_algorithm.n_arms,
            "n_suggestions": simulation.current_algorithm.n_suggestions
        }
        
        return {
            "status": "ready",
            "message": "Simulation ready",
            "simulation_ready": True,
            "algorithm": algorithm_info,
            "environment": {
                "max_arms": getattr(simulation.environment, 'max_arms', 'N/A'),
                "n_suggestions": getattr(simulation.environment, 'n_suggestions', 'N/A')
            }
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error getting status: {str(e)}",
            "simulation_ready": False
        }

@app.post("/recommend_item", response_model=RecommendationResponse)
def recommend_item(request: RecommendationRequest):
    """
    Get recommended items for a user.
    """
    global simulation
    
    if simulation is None:
        raise HTTPException(status_code=500, detail="Simulation not initialized")
    
    try:
        # Configure algorithm if different from current
        if request.algorithm_type != simulation.algorithm_type or request.algorithm_params:
            simulation.set_algorithm(request.algorithm_type, request.algorithm_params)
        
        # Get recommendations using the simulation
        recommendations = simulation.invoke_agent(
            algorithm=simulation.current_algorithm,
            user_id=request.user_id,
            n_recommendations=request.n_recommendations
        )
        
        # Calculate confidence scores (simple implementation)
        confidence_scores = []
        if hasattr(simulation.current_algorithm, 'arm_avg_rewards'):
            for item_id in recommendations:
                if item_id < len(simulation.current_algorithm.arm_avg_rewards):
                    confidence = max(0.1, min(1.0, simulation.current_algorithm.arm_avg_rewards[item_id]))
                    confidence_scores.append(confidence)
                else:
                    confidence_scores.append(0.5)
        else:
            confidence_scores = [0.7] * len(recommendations)
        
        logger.info(f"Generated {len(recommendations)} recommendations for user {request.user_id}")
        
        return RecommendationResponse(
            user_id=request.user_id,
            recommended_items=recommendations,
            algorithm_type=simulation.algorithm_type,
            confidence_scores=confidence_scores
        )
        
    except Exception as e:
        logger.error(f"Error generating recommendations for user {request.user_id}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error generating recommendations: {str(e)}")

@app.post("/feedback", response_model=FeedbackResponse)
def feedback(request: FeedbackRequest):
    """
    Submit feedback for a recommended item.
    """
    global simulation
    
    if simulation is None:
        raise HTTPException(status_code=500, detail="Simulation not initialized")
    
    try:
        # Create event DataFrame
        event_data = {
            'user_id': [request.user_id],
            'item_id': [request.item_id],
            'event_type': [request.event_type],
            'timestamp': [pd.Timestamp.now() if request.timestamp is None else pd.Timestamp(request.timestamp)]
        }
        
        event_df = pd.DataFrame(event_data)
        
        # Update the algorithm with the feedback
        simulation.update_algorithm(event_df)
        
        # Get updated reward information
        updated_rewards = None
        if hasattr(simulation.current_algorithm, 'arm_avg_rewards'):
            item_idx = int(request.item_id)
            if 0 <= item_idx < len(simulation.current_algorithm.arm_avg_rewards):
                updated_rewards = {
                    "item_id": request.item_id,
                    "current_avg_reward": float(simulation.current_algorithm.arm_avg_rewards[item_idx]),
                    "total_pulls": int(simulation.current_algorithm.arm_counts[item_idx]),
                    "total_reward": float(simulation.current_algorithm.arm_rewards[item_idx])
                }
        
        logger.info(f"Processed feedback from user {request.user_id} for item {request.item_id}: {request.event_type}")
        
        return FeedbackResponse(
            status="success",
            message="Feedback processed successfully",
            updated_rewards=updated_rewards
        )
        
    except Exception as e:
        logger.error(f"Error processing feedback from user {request.user_id}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing feedback: {str(e)}")

class ConfigureRequest(BaseModel):
    algorithm_type: str
    algorithm_params: Optional[Dict[str, Any]] = {}

@app.post("/configure", response_model=Dict[str, Any])
def configure_algorithm(request: ConfigureRequest):
    """
    Configure the MAB algorithm.
    """
    global simulation
    
    if simulation is None:
        raise HTTPException(status_code=500, detail="Simulation not initialized")
    
    try:
        # Set new algorithm
        new_algorithm = simulation.set_algorithm(request.algorithm_type, request.algorithm_params or {})
        
        logger.info(f"Algorithm configured to {request.algorithm_type}")
        
        return {
            "status": "success",
            "message": f"Algorithm configured to {request.algorithm_type}",
            "algorithm_info": {
                "type": request.algorithm_type,
                "class": new_algorithm.__class__.__name__,
                "n_arms": new_algorithm.n_arms,
                "n_suggestions": new_algorithm.n_suggestions
            }
        }
        
    except Exception as e:
        logger.error(f"Error configuring algorithm: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error configuring algorithm: {str(e)}")

@app.get("/health")
def health_check():
    """Health check endpoint."""
    return {"status": "healthy", "timestamp": pd.Timestamp.now().isoformat()}