import json
import re
import frappe

@frappe.whitelist()
def optimize_routes(pickup_date=None):
    """Optimize routes for all trucks on a given date."""
    if not pickup_date:
        pickup_date = frappe.utils.today()
    
    trucks = frappe.get_list("Truck Profile", 
        filters={"active": 1},
        fields=["name", "truck_name", "total_slots", "weight_capacity"]
    )
    
    pickups = frappe.get_list("Scheduled Pickup",
        filters={
            "pickup_date": pickup_date,
            "shipment_type": "Truck",
            "truck_profile": ["is", "not set"]
        },
        fields=["name", "customer_number", "company_name", "estimated_items", 
                "estimated_weight", "gaylord_count", "gaylord_sizes", "slots_needed",
                "latitude", "longitude"],
        order_by="stop_order asc"
    )
    
    for pickup in pickups:
        if not pickup.slots_needed and pickup.gaylord_sizes:
            pickup.slots_needed = _calculate_slots(pickup.gaylord_sizes)
        elif not pickup.slots_needed:
            pickup.slots_needed = pickup.gaylord_count or 1
    
    pickups.sort(key=lambda x: x.slots_needed or 0, reverse=True)
    
    routes = {}
    for truck in trucks:
        routes[truck.name] = {
            "truck": truck,
            "pickups": [],
            "used_slots": 0,
            "used_weight": 0
        }
    
    unassigned = []
    
    for pickup in pickups:
        assigned = False
        for truck_name, route in routes.items():
            truck = route["truck"]
            slots = pickup.slots_needed or 0
            weight = float(pickup.estimated_weight or 0)
            
            if route["used_slots"] + slots <= truck.total_slots:
                if truck.weight_capacity and route["used_weight"] + weight <= truck.weight_capacity:
                    route["pickups"].append(pickup)
                    route["used_slots"] += slots
                    route["used_weight"] += weight
                    assigned = True
                    frappe.db.set_value("Scheduled Pickup", pickup.name, "truck_profile", truck_name)
                    break
                elif not truck.weight_capacity:
                    route["pickups"].append(pickup)
                    route["used_slots"] += slots
                    route["used_weight"] += weight
                    assigned = True
                    frappe.db.set_value("Scheduled Pickup", pickup.name, "truck_profile", truck_name)
                    break
        if not assigned:
            unassigned.append(pickup)
    
    for truck_name, route in routes.items():
        if route["pickups"]:
            route["pickups"].sort(key=lambda p: (float(p.latitude or 0), float(p.longitude or 0)))
            for i, pickup in enumerate(route["pickups"], 1):
                frappe.db.set_value("Scheduled Pickup", pickup.name, "stop_order", i)
    
    frappe.db.commit()
    
    return {
        "success": True,
        "date": pickup_date,
        "trucks_assigned": len([r for r in routes.values() if r["pickups"]]),
        "total_pickups": len(pickups),
        "unassigned": len(unassigned),
        "routes": {
            truck_name: {
                "truck_name": route["truck"].truck_name,
                "slots_used": route["used_slots"],
                "slots_total": route["truck"].total_slots,
                "weight_used": route["used_weight"],
                "weight_capacity": route["truck"].weight_capacity,
                "stops": len(route["pickups"]),
                "pickups": [{"name": p.name, "company": p.company_name, "slots": p.slots_needed} for p in route["pickups"]]
            }
            for truck_name, route in routes.items() if route["pickups"]
        }
    }

def _calculate_slots(gaylord_sizes_text):
    if not gaylord_sizes_text:
        return 0
    size_map = {"small": 1, "medium": 2, "large": 3}
    total = 0
    matches = re.findall(r'(\d+)\s*(\w+)', gaylord_sizes_text.lower())
    for count, size in matches:
        slots = size_map.get(size, 1)
        total += int(count) * slots
    return total or 1


@frappe.whitelist()
def get_scheduled_pickups(pickup_date=None):
    """Get scheduled pickups for a given date."""
    if not pickup_date:
        pickup_date = frappe.utils.today()
    
    pickups = frappe.get_list("Scheduled Pickup",
        filters={"pickup_date": pickup_date},
        fields=["name", "customer_number", "company_name", "estimated_items",
                 "estimated_weight", "gaylord_count", "gaylord_sizes", "slots_needed",
                 "latitude", "longitude", "stop_order", "truck_profile", "status",
                 "contact_name", "contact_phone", "address_line", "city", "state", "zip_code"],
        order_by="stop_order asc"
    )
    
    return pickups