# Import Hugging Face datasets
from datasets import load_dataset
from colorama import Fore
from mcp.server.fastmcp import FastMCP
import chromadb

# Create server
mcp = FastMCP("croptimizeserver")

# Load crop optimization dataset
dataset = load_dataset("DARJYO/sawotiQ29_crop_optimization")

# Initialize ChromaDB for crop data
chroma_client = chromadb.PersistentClient(path="crop_db")
collection = chroma_client.get_collection(name="crop_data")

# Add prompt function for crop recommendations
@mcp.prompt()
def crop_recommendation(crop_data: str) -> str:
    """Prompt template for generating crop recommendations"""
    return f"""You are an agricultural expert assistant designed to provide crop optimization advice.
            Analyze the following crop data and provide recommendations for optimal cultivation:
            {crop_data}"""

# Resource for searching crop information
@mcp.resource("crops://search/{query}")
def search_crops(query: str) -> str:
    """Search for crops based on growing conditions or characteristics"""
    results = collection.query(
        query_texts=[query],
        n_results=3,
        include=["documents", "metadatas"]
    )
    return str(results)

# Tool for getting crop details
@mcp.tool()
def crop_details(crop_name: str) -> str:
    """Get detailed information about a specific crop"""
    filtered_data = dataset['train'].filter(lambda x: x['crop_name'].lower() == crop_name.lower())
    if not filtered_data:
        return f"No information found for {crop_name}"
    return str(filtered_data[0])

# Tool for optimal growing conditions
@mcp.tool()
def optimal_conditions(crop_name: str) -> str:
    """Get optimal growing conditions for a specific crop"""
    crop_data = dataset['train'].filter(lambda x: x['crop_name'].lower() == crop_name.lower())
    if not crop_data:
        return f"No data available for {crop_name}"
    
    conditions = {
        'temperature': crop_data[0]['optimal_temperature'],
        'rainfall': crop_data[0]['annual_rainfall'],
        'soil_type': crop_data[0]['preferred_soil'],
        'altitude': crop_data[0]['optimal_altitude']
    }
    return str(conditions)

# Tool for yield prediction
@mcp.tool()
def yield_prediction(crop_name: str, region: str) -> str:
    """Predict yield for a crop in a specific region"""
    region_data = dataset['train'].filter(lambda x: 
        (x['region'].lower() == region.lower()) and 
        (x['crop_name'].lower() == crop_name.lower())
    )
    
    if not region_data:
        return f"No yield data available for {crop_name} in {region}"
    
    prediction = {
        'crop': crop_name,
        'region': region,
        'expected_yield': region_data[0]['average_yield'],
        'optimal_season': region_data[0]['best_season']
    }
    return str(prediction)

# Tool for pest/disease information
@mcp.tool()
def crop_protection(crop_name: str) -> str:
    """Get common pests and diseases for a crop"""
    crop_data = dataset['train'].filter(lambda x: x['crop_name'].lower() == crop_name.lower())
    if not crop_data:
        return f"No protection data available for {crop_name}"
    
    protection_info = {
        'common_pests': crop_data[0]['common_pests'],
        'common_diseases': crop_data[0]['common_diseases'],
        'prevention_methods': crop_data[0]['prevention_methods']
    }
    return str(protection_info)

if __name__ == "__main__":
    mcp.run(transport="stdio")