db_utils.py

import os
from pymongo import MongoClient
from dotenv import load_dotenv

load_dotenv()

def get_db():
    mongo_url = os.getenv('MONGO_CONNECTION_URL')
    if not mongo_url:
        raise ValueError("MONGO_CONNECTION_URL not found in environment variables")
    
    client = MongoClient(mongo_url)
    db = client['AgriAI']
    return client, db


def setup_user_collections():
    """Create collections for user-specific data if they don't exist"""
    client, db = get_db()
    
    collections = ['farm_profiles', 'user_market_configs', 'user_financials', 'user_operations', 'user_reminders', 'user_strategic_plans']
    
    for coll_name in collections:
        if coll_name not in db.list_collection_names():
            db.create_collection(coll_name)
            print(f"Created collection: {coll_name}")
        else:
            print(f"Collection already exists: {coll_name}")
    
    # Create indexes for performance
    db.farm_profiles.create_index("user_id", unique=True)
    db.user_market_configs.create_index("user_id", unique=True)
    db.user_financials.create_index("user_id", unique=True)
    db.user_operations.create_index("user_id")
    db.user_reminders.create_index("user_id")
    db.user_strategic_plans.create_index("user_id")
    
    print("Indexes created successfully")
    client.close()


def save_strategic_plan(user_id, plan_data):
    """Save a generated strategic plan to the user's account"""
    client, db = get_db()
    from datetime import datetime
    
    # Store with metadata
    plan_document = {
        "user_id": user_id,
        "plan": plan_data,
        "saved_at": datetime.utcnow()
    }
    
    # Replace existing plan for user (always keep one active plan)
    result = db.user_strategic_plans.replace_one(
        {"user_id": user_id},
        plan_document,
        upsert=True
    )
    print(f"Saved strategic plan for user: {user_id}")
    
    client.close()
    # Return the upserted ID if new document, otherwise return success indicator
    return str(result.upserted_id) if result.upserted_id else "updated"

def seed_default_data(user_id):
    """Seed default data for a new user"""
    client, db = get_db()
    from datetime import datetime
    
    # Check if user already has data
    existing_profile = db.farm_profiles.find_one({"user_id": user_id})
    if existing_profile:
        print(f"User {user_id} already has profile data")
        client.close()
        return
    
    # Default farm profile
    farm_profile = {
        "user_id": user_id,
        "location": {"state": "Maharashtra", "district": "Pune"},
        "farm_name": "My Farm",
        "total_acres": 42,
        "active_crops": [
            {"name": "Vegetables", "acres": 12, "icon": "🥦"},
            {"name": "Wheat", "acres": 20, "icon": "🌾"},
            {"name": "Banana Multi", "acres": 10, "icon": "🍌"}
        ],
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }
    
    # Default market config (crops user tracks)
    market_config = {
        "user_id": user_id,
        "tracked_crops": [
            {"name": "Wheat", "api_name": "Wheat", "state": "Maharashtra", "district": "Nashik", "mandi": "Nashik Regulated Mandi", "fallback_price": 2450, "fallback_change": 45},
            {"name": "Soybean", "api_name": "Soybean", "state": "Madhya Pradesh", "district": "Indore", "mandi": "Estimated (Last seen Indore)", "fallback_price": 4820, "fallback_change": -112},
            {"name": "Gram", "api_name": "Gram", "state": "Maharashtra", "district": "Pune", "mandi": "Closing Price", "fallback_price": 5100, "fallback_change": 18}
        ],
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }
    
    # Default financial data
    financial_data = {
        "user_id": user_id,
        "fiscal_year": "2025-26",
        "net_yield": 14200,
        "net_yield_formatted": "₹14,200",
        "yoy_change_pct": 12.4,
        "yoy_change_text": "+12.4%",
        "monthly": [
            {"month": "Jan", "net": 1200, "amount_text": "+₹1,200", "type": "surplus", "height_pct": 35},
            {"month": "Feb", "net": -450, "amount_text": "-₹450", "type": "deficit", "height_pct": 15},
            {"month": "Mar", "net": 2800, "amount_text": "+₹2,800", "type": "surplus", "height_pct": 60},
            {"month": "Apr", "net": 850, "amount_text": "+₹850", "type": "surplus", "height_pct": 25},
            {"month": "May", "net": -1100, "amount_text": "-₹1,100", "type": "deficit", "height_pct": 30},
            {"month": "Jun", "net": 3400, "amount_text": "+₹3,400", "type": "surplus", "height_pct": 75},
            {"month": "Jul", "net": 150, "amount_text": "+₹150", "type": "surplus", "height_pct": 10},
            {"month": "Aug", "net": -2300, "amount_text": "-₹2,300", "type": "deficit", "height_pct": 50},
            {"month": "Sep", "net": 920, "amount_text": "+₹920", "type": "surplus", "height_pct": 28},
            {"month": "Oct", "net": 1650, "amount_text": "+₹1,650", "type": "surplus", "height_pct": 45},
            {"month": "Nov", "net": 4200, "amount_text": "+₹4,200", "type": "surplus", "height_pct": 90},
            {"month": "Dec", "net": 2780, "amount_text": "+₹2,780", "type": "surplus", "height_pct": 65}
        ],
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }
    
    # Default operations
    operations_data = {
        "user_id": user_id,
        "operations": [
            {
                "id": 1, "date": "March 15, 2026", "title": "Apply Nitrogen Fertilizer",
                "description": "Strategic enrichment for Sector B Wheat fields based on soil sensor readings.",
                "icon": "agriculture", "color": "primary", "status": "active", "priority": "high"
            },
            {
                "id": 2, "date": "March 20, 2026", "title": "Pest Control Cycle",
                "description": "Preventative spray cycle initiated by localized insect swarm detection alerts.",
                "icon": "bug_report", "color": "tertiary", "status": "active", "priority": "high"
            },
            {
                "id": 3, "date": "April 02, 2026", "title": "Soil Quality Audit",
                "description": "Quarterly physical verification of IoT sensor calibration.",
                "icon": "science", "color": "outline", "status": "upcoming", "priority": "medium"
            }
        ],
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }
    
    # Default reminders
    reminders_data = {
        "user_id": user_id,
        "reminders": [
            {"task": "Apply Fertilizer", "date": "Mar 15"},
            {"task": "Pest Control", "date": "Mar 20"},
            {"task": "Soil Audit", "date": "Apr 02"}
        ],
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }
    
    # Insert all default data
    db.farm_profiles.insert_one(farm_profile)
    db.user_market_configs.insert_one(market_config)
    db.user_financials.insert_one(financial_data)
    db.user_operations.insert_one(operations_data)
    db.user_reminders.insert_one(reminders_data)
    
    print(f"Seed data created for user: {user_id}")
    client.close()


def get_user_profile(user_id):
    """Get user's farm profile"""
    client, db = get_db()
    profile = db.farm_profiles.find_one({"user_id": user_id})
    client.close()
    return profile


def get_user_market_config(user_id):
    """Get user's market configuration"""
    client, db = get_db()
    config = db.user_market_configs.find_one({"user_id": user_id})
    client.close()
    return config


def get_user_financials(user_id):
    """Get user's financial data"""
    client, db = get_db()
    financials = db.user_financials.find_one({"user_id": user_id})
    client.close()
    return financials


def get_user_operations(user_id):
    """Get user's operations"""
    client, db = get_db()
    operations = db.user_operations.find_one({"user_id": user_id})
    client.close()
    return operations


def get_user_reminders(user_id):
    """Get user's reminders"""
    client, db = get_db()
    reminders = db.user_reminders.find_one({"user_id": user_id})
    client.close()
    return reminders


if __name__ == "__main__":
    # Test connection and create collections
    print("Setting up MongoDB collections...")
    setup_user_collections()
    
    # For testing - seed data for a test user (use actual ObjectId from your users collection)
    # Example: from bson import ObjectId; seed_default_data(ObjectId("69c7d7f0b46dd3a714fbf7e9"))
    print("\nTo seed data for a user, import db_utils and call seed_default_data(user_id)")