app/models/trip_optimizer/optimizer.py

import json
import os
import datetime
from typing import Optional

from app.models.route_optimizer.distance_calculator import DistanceCalculator
from app.models.route_optimizer.route_planner import RoutePlanner
from app.models.budget_optimizer.cost_calculator import CostCalculator


class TripOptimizer:
    """Enhanced constraint-based trip planning optimizer with distance-aware scheduling,
    multi-destination support, and 3-tier budget plans."""

    def __init__(self, knowledge_base_path: str):
        self.destinations = []
        self.destinations_map = {}
        if os.path.exists(knowledge_base_path):
            with open(knowledge_base_path, 'r', encoding='utf-8-sig') as f:
                data = json.load(f)
                self.destinations = data.get('destinations', [])
                self.destinations_map = {d['id']: d for d in self.destinations}

        self.route_planner = RoutePlanner()

    def generate_plan(
        self,
        destination_id: str,
        duration: dict,
        budget: Optional[dict] = None,
        num_people: Optional[dict] = None,
        preferences: Optional[list] = None,
        language: str = 'vi',
        budget_priorities: Optional[dict] = None,
        dietary_restrictions: Optional[list] = None,
        cuisine_preferences: Optional[list] = None,
        user_profile: Optional[dict] = None,
        origin_id: Optional[str] = None,
    ) -> dict | None:
        """Generate an optimized trip plan with distance-aware scheduling and budget tiers."""
        dest = self.destinations_map.get(destination_id)
        if not dest:
            return None

        days = duration.get('days', 3)
        nights = duration.get('nights', days - 1)
        adults = (num_people or {}).get('adults', 2)
        children = (num_people or {}).get('children', 0)
        total_people = adults + children

        activities = dest.get('popular_activities', [])
        foods = dest.get('food_specialties', [])

        # Determine budget tier
        tier = self._determine_tier(budget, days, total_people, dest)

        # Score and sort activities
        scored_activities = self._score_activities(activities, preferences or [], user_profile)

        # Cluster activities by geographic proximity
        clusters = self._cluster_activities_for_days(scored_activities, days)

        # Build itinerary
        itinerary = []
        used_activities = set()
        total_cost = 0

        # Get daily costs for chosen tier
        daily_cost = CostCalculator.calculate_daily_cost(
            dest, tier, adults, children, datetime.datetime.now().month
        )

        for day in range(1, days + 1):
            day_plan = {'day': day, 'activities': []}

            # Pick cluster for this day
            cluster_idx = min(day - 1, len(clusters) - 1) if clusters else -1
            day_activities = clusters[cluster_idx] if cluster_idx >= 0 and clusters else scored_activities

            if day == 1:
                # Arrival day
                transport = dest.get('transport', {})
                arrival_info = self._get_arrival_info(transport, language, origin_id)
                if arrival_info:
                    day_plan['activities'].append({
                        'time': '08:00',
                        'type': 'transport',
                        'title': arrival_info['title'],
                        'description': arrival_info['description'],
                        'estimatedCost': 0
                    })

                # Check-in
                accom_cost = daily_cost['accommodation']
                day_plan['activities'].append({
                    'time': '10:00',
                    'type': 'hotel',
                    'title': 'Check-in khách sạn' if language == 'vi' else 'Hotel check-in',
                    'description': daily_cost.get('tier_description', ''),
                    'estimatedCost': accom_cost
                })
                total_cost += accom_cost

                # Afternoon: 1 light activity (distance-optimized)
                act = self._pick_activity(day_activities, used_activities, 'light')
                if act:
                    used_activities.add(act['name'])
                    act_cost = CostCalculator.calculate_activity_cost(act, adults, children)
                    travel_time = act.get('time_from_center_min', 15)
                    _coords = act.get('coordinates', {})
                    day_plan['activities'].append({
                        'time': '14:00',
                        'type': 'activity',
                        'title': act['name'] if language == 'vi' else act.get('name_en', act['name']),
                        'description': self._format_activity_desc(act, travel_time, language),
                        'estimatedCost': act_cost,
                        'travelTimeMin': travel_time,
                        'lat': _coords.get('lat'),
                        'lon': _coords.get('lon'),
                    })
                    total_cost += act_cost

            elif day == days:
                # Departure day
                act = self._pick_activity(day_activities, used_activities, 'light')
                if act:
                    used_activities.add(act['name'])
                    act_cost = CostCalculator.calculate_activity_cost(act, adults, children)
                    _coords = act.get('coordinates', {})
                    day_plan['activities'].append({
                        'time': '08:00',
                        'type': 'activity',
                        'title': act['name'] if language == 'vi' else act.get('name_en', act['name']),
                        'description': self._format_activity_desc(act, act.get('time_from_center_min', 10), language),
                        'estimatedCost': act_cost,
                        'lat': _coords.get('lat'),
                        'lon': _coords.get('lon'),
                    })
                    total_cost += act_cost

                day_plan['activities'].append({
                    'time': '14:00',
                    'type': 'hotel',
                    'title': 'Check-out' if language == 'en' else 'Trả phòng',
                    'description': '',
                    'estimatedCost': 0
                })

                day_plan['activities'].append({
                    'time': '16:00',
                    'type': 'transport',
                    'title': 'Khởi hành về' if language == 'vi' else 'Departure',
                    'description': '',
                    'estimatedCost': 0
                })
            else:
                # Full day
                accom_cost = daily_cost['accommodation']
                total_cost += accom_cost

                # Optimize order of activities in this day's cluster
                day_cluster = [a for a in day_activities if a['name'] not in used_activities]
                if len(day_cluster) > 2:
                    day_cluster = self.route_planner.optimize_daily_route(day_cluster)

                # Morning activity
                act_morning = self._pick_activity(day_cluster, used_activities, 'any')
                if act_morning:
                    used_activities.add(act_morning['name'])
                    act_cost = CostCalculator.calculate_activity_cost(act_morning, adults, children)
                    duration_h = act_morning.get('duration_hours', 3)
                    _coords_m = act_morning.get('coordinates', {})
                    day_plan['activities'].append({
                        'time': self._get_best_time_slot(act_morning, 'morning'),
                        'type': 'activity',
                        'title': act_morning['name'] if language == 'vi' else act_morning.get('name_en', act_morning['name']),
                        'description': self._format_activity_desc(act_morning, act_morning.get('time_from_center_min', 15), language),
                        'estimatedCost': act_cost,
                        'durationHours': duration_h,
                        'lat': _coords_m.get('lat'),
                        'lon': _coords_m.get('lon'),
                    })
                    total_cost += act_cost

                # Afternoon activity
                act_afternoon = self._pick_activity(day_cluster, used_activities, 'any')
                if act_afternoon:
                    used_activities.add(act_afternoon['name'])
                    act_cost = CostCalculator.calculate_activity_cost(act_afternoon, adults, children)
                    _coords_a = act_afternoon.get('coordinates', {})
                    day_plan['activities'].append({
                        'time': self._get_best_time_slot(act_afternoon, 'afternoon'),
                        'type': 'activity',
                        'title': act_afternoon['name'] if language == 'vi' else act_afternoon.get('name_en', act_afternoon['name']),
                        'description': self._format_activity_desc(act_afternoon, act_afternoon.get('time_from_center_min', 15), language),
                        'estimatedCost': act_cost,
                        'durationHours': act_afternoon.get('duration_hours', 2),
                        'lat': _coords_a.get('lat'),
                        'lon': _coords_a.get('lon'),
                    })
                    total_cost += act_cost

            # Add food for each day — use full daily food cost (3 meals, with child discount)
            if foods:
                food = foods[day % len(foods) - 1]
                food_cost = daily_cost['food']  # already includes 3 meals × adults + children w/ discount
                day_plan['activities'].append({
                    'time': '12:00' if day_plan['activities'] else '19:00',
                    'type': 'food',
                    'title': food['name'] if language == 'vi' else food.get('name_en', food['name']),
                    'description': f"{food.get('avg_price_vnd', 50000):,} VND/{'người' if language == 'vi' else 'person'}",
                    'estimatedCost': food_cost
                })
                total_cost += food_cost

            # Sort activities by time
            day_plan['activities'].sort(key=lambda x: x['time'])
            itinerary.append(day_plan)

        # Add misc costs
        daily_misc = daily_cost['misc'] * days
        transport_cost = daily_cost['transport'] * days
        total_cost += daily_misc + transport_cost

        dest_name = dest['name'] if language == 'vi' else dest.get('name_en', dest['name'])

        # Generate 3-tier cost comparison
        tier_plans = self._generate_tier_comparison(dest, days, adults, children)

        # Calculate route optimization score
        opt_score = self._calculate_optimization_score(itinerary, dest)

        # Compute consistent cost breakdown from the same daily_cost source
        accom_breakdown = int(daily_cost['accommodation'] * max(1, days - 1))
        food_breakdown = int(daily_cost['food'] * days)
        transport_breakdown = int(transport_cost)
        misc_breakdown = int(daily_misc)
        # Activities = total minus known costs (to ensure totalEstimatedCost == sum(breakdown))
        activities_breakdown = int(total_cost - accom_breakdown - food_breakdown - transport_breakdown - misc_breakdown)
        activities_breakdown = max(0, activities_breakdown)  # Guard against negative

        return {
            'destination': dest_name,
            'destination_id': destination_id,
            'duration': {'days': days, 'nights': nights},
            'travelers': num_people or {'adults': 2, 'children': 0, 'infants': 0},
            'itinerary': itinerary,
            'totalEstimatedCost': int(total_cost),
            'currency': 'VND',
            'selectedTier': tier,
            'costBreakdown': {
                'accommodation': accom_breakdown,
                'food': food_breakdown,
                'activities': activities_breakdown,
                'transport': transport_breakdown,
                'misc': misc_breakdown,
            },
            'perPerson': int(total_cost // total_people) if total_people > 0 else int(total_cost),
            'perPersonPerDay': int(total_cost // total_people // days) if total_people > 0 and days > 0 else 0,
            'tierComparison': tier_plans,
            'optimizationScore': opt_score,
            'tips': self._get_tips(dest, language, tier),
        }

    def generate_multi_destination_plan(
        self,
        destination_ids: list[str],
        total_days: int,
        budget: Optional[dict] = None,
        num_people: Optional[dict] = None,
        preferences: Optional[list] = None,
        language: str = 'vi',
    ) -> dict | None:
        """Generate a plan visiting multiple destinations with optimized routing."""
        dests = [self.destinations_map[did] for did in destination_ids if did in self.destinations_map]
        if not dests:
            return None

        # Optimize inter-city route order
        dests = self.route_planner.plan_multi_destination_route(dests)

        # Allocate days per destination (proportional, minimum 2 days each)
        allocation = self._allocate_days(dests, total_days)

        # Generate plan for each destination
        segments = []
        total_cost = 0
        current_day = 1

        for i, dest in enumerate(dests):
            days_for_dest = allocation[i]
            segment_plan = self.generate_plan(
                destination_id=dest['id'],
                duration={'days': days_for_dest, 'nights': days_for_dest - 1},
                budget=budget,
                num_people=num_people,
                preferences=preferences,
                language=language,
            )

            if segment_plan:
                # Adjust day numbers
                for day_plan in segment_plan['itinerary']:
                    day_plan['day'] = current_day
                    current_day += 1

                segments.append(segment_plan)
                total_cost += segment_plan['totalEstimatedCost']

            # Add inter-city travel (except after last destination)
            if i < len(dests) - 1:
                travel_info = self.route_planner.estimate_inter_city_travel(dest, dests[i + 1])
                travel_cost = self._estimate_travel_cost(travel_info)
                total_cost += travel_cost

        dest_names = [
            (d['name'] if language == 'vi' else d.get('name_en', d['name']))
            for d in dests
        ]

        adults = (num_people or {}).get('adults', 2)
        children = (num_people or {}).get('children', 0)
        total_people = adults + children

        return {
            'type': 'multi_destination',
            'destinations': dest_names,
            'destination_ids': [d['id'] for d in dests],
            'duration': {'days': total_days, 'nights': total_days - 1},
            'travelers': num_people or {'adults': 2, 'children': 0, 'infants': 0},
            'segments': segments,
            'totalEstimatedCost': int(total_cost),
            'perPerson': int(total_cost // total_people) if total_people > 0 else int(total_cost),
            'currency': 'VND',
            'routeOrder': dest_names,
            'tips': self._get_multi_dest_tips(dests, language),
        }

    def compare_destinations(
        self,
        dest_ids: list[str],
        days: int = 3,
        num_adults: int = 2,
        num_children: int = 0,
        language: str = 'vi'
    ) -> list[dict]:
        """Compare multiple destinations across cost, weather, activities, food."""
        results = []
        month = datetime.datetime.now().month

        for dest_id in dest_ids:
            dest = self.destinations_map.get(dest_id)
            if not dest:
                continue

            mid_cost = CostCalculator.calculate_trip_total(
                dest, days, 'mid_range', num_adults, num_children, month=month
            )

            name = dest['name'] if language == 'vi' else dest.get('name_en', dest['name'])
            is_best_month = month in dest.get('best_months', [])

            results.append({
                'destination_id': dest_id,
                'name': name,
                'region': dest.get('region', ''),
                'tags': dest.get('tags', []),
                'best_months': dest.get('best_months', []),
                'is_good_time_now': is_best_month,
                'avg_temp_c': dest.get('avg_temp_c', 0),
                'num_activities': len(dest.get('popular_activities', [])),
                'num_food_specialties': len(dest.get('food_specialties', [])),
                'cost_mid_range_total': mid_cost['total'],
                'cost_per_person': mid_cost['per_person'],
                'cost_per_person_usd': mid_cost['per_person_usd'],
                'description': dest.get(f'description_{language}', dest.get('description_vi', '')),
            })

        results.sort(key=lambda x: x['cost_mid_range_total'])
        return results

    # --- Private helpers ---

    def _determine_tier(self, budget: Optional[dict], days: int, total_people: int, dest: dict) -> str:
        """Determine the most appropriate budget tier."""
        if not budget:
            return 'mid_range'

        amount = budget.get('amount', 0)
        currency = budget.get('currency', 'VND')
        if currency == 'USD':
            amount = CostCalculator.convert_currency(amount, 'USD', 'VND')

        per_person_per_day = amount / max(1, total_people) / max(1, days)

        cost_tiers = dest.get('cost_tiers', {})
        for tier_name in ['luxury', 'mid_range', 'budget']:
            tier = cost_tiers.get(tier_name, {})
            accom = tier.get('accommodation_vnd', 800000)
            food = tier.get('food_per_meal_vnd', 80000) * 3
            transport = tier.get('transport_per_day_vnd', 200000)
            daily = accom + food + transport + 100000
            if per_person_per_day >= daily:
                return tier_name

        return 'budget'

    def _score_activities(self, activities: list, preferences: list, user_profile: Optional[dict] = None) -> list:
        """Score activities based on user preferences and profile."""
        scored = []
        for act in activities:
            score = 1.0
            act_type = act.get('type', '')
            act_name = (act.get('name', '') + ' ' + act.get('name_en', '')).lower()

            if preferences:
                for pref in preferences:
                    if pref.lower() in act_type.lower() or pref.lower() in act_name:
                        score += 0.5

            if user_profile:
                from app.models.user_profiler.profiler import UserProfiler
                profiler = UserProfiler()
                score += profiler.score_activity_for_user(act, user_profile) - 1.0

            scored.append({**act, '_score': score})

        scored.sort(key=lambda x: x['_score'], reverse=True)
        return scored

    def _cluster_activities_for_days(self, activities: list, days: int) -> list[list]:
        """Cluster activities into day-groups by proximity."""
        activity_days = max(1, days - 2) + 2
        try:
            clusters = self.route_planner.cluster_by_proximity(activities, activity_days)
            return clusters if clusters else [activities]
        except Exception:
            return [activities]

    def _pick_activity(self, scored_activities: list, used: set, mode: str) -> dict | None:
        """Pick next best activity not yet used."""
        for act in scored_activities:
            if act['name'] not in used:
                if mode == 'light' and act.get('duration_hours', 3) > 4:
                    continue
                return act
        return None

    def _get_best_time_slot(self, activity: dict, slot: str) -> str:
        """Get the best time slot based on activity's best_time."""
        best = activity.get('best_time', 'any')
        if slot == 'morning':
            return '06:00' if best == 'early_morning' else '08:00'
        elif slot == 'afternoon':
            return '17:00' if best == 'evening' else '14:00'
        return '08:00' if slot == 'morning' else '14:00'

    def _format_activity_desc(self, act: dict, travel_time: int, language: str) -> str:
        """Format activity description with duration and travel time."""
        duration_h = act.get('duration_hours', 2)
        act_type = act.get('type', 'sightseeing')
        hours = act.get('opening_hours', '')

        if language == 'vi':
            desc = f"~{duration_h}h - {act_type}"
            if travel_time > 0:
                desc += f" | Di chuyển ~{travel_time} phút"
            if hours and hours != '00:00-23:59':
                desc += f" | Mở cửa: {hours}"
        else:
            desc = f"~{duration_h}h - {act_type}"
            if travel_time > 0:
                desc += f" | Travel ~{travel_time}min"
            if hours and hours != '00:00-23:59':
                desc += f" | Open: {hours}"

        return desc

    # Map origin IDs to transport keys and display names
    ORIGIN_TRANSPORT_KEYS = {
        'hcmc': ('from_hcm', 'TP.HCM', 'HCMC'),
        'hanoi': ('from_hanoi', 'Hà Nội', 'Hanoi'),
        'da-nang': ('from_da_nang', 'Đà Nẵng', 'Da Nang'),
    }

    def _get_arrival_info(
        self, transport: dict, language: str, origin_id: Optional[str] = None
    ) -> dict | None:
        """Get arrival transport info based on actual origin city."""
        # Determine which transport route to look up
        if origin_id and origin_id in self.ORIGIN_TRANSPORT_KEYS:
            key, name_vi, name_en = self.ORIGIN_TRANSPORT_KEYS[origin_id]
        else:
            # Fallback: try all known keys, prefer from_hcm
            for oid in ['hcmc', 'hanoi', 'da-nang']:
                key, name_vi, name_en = self.ORIGIN_TRANSPORT_KEYS[oid]
                if transport.get(key):
                    break
            else:
                key, name_vi, name_en = 'from_hcm', 'TP.HCM', 'HCMC'

        route = transport.get(key, {})
        flight_time = route.get('flight', '')
        bus_time = route.get('bus', '')

        if flight_time:
            if language == 'vi':
                return {
                    'title': f'Bay từ {name_vi} ({flight_time})',
                    'description': 'Di chuyển bằng máy bay'
                }
            else:
                return {
                    'title': f'Flight from {name_en} ({flight_time})',
                    'description': 'Travel by airplane'
                }
        elif bus_time:
            if language == 'vi':
                return {
                    'title': f'Di chuyển từ {name_vi} ({bus_time})',
                    'description': 'Di chuyển bằng xe khách'
                }
            else:
                return {
                    'title': f'Travel from {name_en} ({bus_time})',
                    'description': 'Travel by bus'
                }

        return {
            'title': 'Di chuyển đến điểm đến' if language == 'vi' else 'Travel to destination',
            'description': ''
        }

    def _generate_tier_comparison(self, dest: dict, days: int, adults: int, children: int) -> list[dict]:
        """Generate cost comparison for 3 tiers."""
        month = datetime.datetime.now().month
        tiers = []
        for tier_name in ['budget', 'mid_range', 'luxury']:
            cost = CostCalculator.calculate_trip_total(
                dest, days, tier_name, adults, children, month=month
            )
            tier_data = dest.get('cost_tiers', {}).get(tier_name, {})
            tiers.append({
                'tier': tier_name,
                'total': cost['total'],
                'perPerson': cost['per_person'],
                'perPersonPerDay': cost['per_person_per_day'],
                'totalUsd': cost['total_usd'],
                'description': tier_data.get('description_vi', ''),
                'descriptionEn': tier_data.get('description_en', ''),
            })
        return tiers

    def _calculate_optimization_score(self, itinerary: list, dest: dict) -> float:
        """Calculate how well the itinerary is optimized (0-1)."""
        score = 0.7
        total_activities = sum(
            1 for day in itinerary for act in day['activities'] if act['type'] == 'activity'
        )
        if total_activities > 0:
            score += 0.1
        score += 0.05

        types = set()
        for day in itinerary:
            for act in day['activities']:
                if act['type'] == 'activity':
                    desc = act.get('description', '')
                    types.add(desc.split(' - ')[-1] if ' - ' in desc else '')
        if len(types) >= 3:
            score += 0.1

        return min(1.0, round(score, 2))

    def _get_tips(self, dest: dict, language: str, tier: str = 'mid_range') -> list:
        """Get travel tips for destination."""
        tips = []
        best_months = dest.get('best_months', [])
        if best_months:
            month_str = ', '.join(str(m) for m in best_months)
            if language == 'vi':
                tips.append(f"Thời điểm đẹp nhất: tháng {month_str}")
            else:
                tips.append(f"Best time to visit: months {month_str}")

        local_transport = dest.get('transport', {}).get('local', [])
        if local_transport:
            if language == 'vi':
                tips.append(f"Di chuyển nội thành: {', '.join(local_transport)}")
            else:
                tips.append(f"Local transport: {', '.join(local_transport)}")

        if tier == 'budget':
            tips.append(
                "Mẹo tiết kiệm: đặt phòng sớm, ăn quán bình dân, dùng xe bus/Grab" if language == 'vi'
                else "Budget tip: book early, eat at local stalls, use bus/Grab"
            )
        elif tier == 'luxury':
            tips.append(
                "Gợi ý: đặt resort có bể bơi, thử fine dining, thuê xe riêng" if language == 'vi'
                else "Tip: book resort with pool, try fine dining, hire private car"
            )

        return tips

    def _allocate_days(self, dests: list, total_days: int) -> list[int]:
        """Allocate days across destinations proportionally."""
        n = len(dests)
        if n == 0:
            return []

        min_per_dest = 2
        remaining = total_days - (min_per_dest * n)
        allocation = [min_per_dest] * n

        if remaining > 0:
            activity_counts = [len(d.get('popular_activities', [])) for d in dests]
            total_acts = sum(activity_counts) or 1
            for i in range(n):
                extra = int(remaining * activity_counts[i] / total_acts)
                allocation[i] += extra
            leftover = total_days - sum(allocation)
            for i in range(leftover):
                allocation[i % n] += 1

        return allocation

    def _estimate_travel_cost(self, travel_info: dict) -> int:
        """Estimate inter-city travel cost."""
        transport = travel_info.get('recommended_transport', 'bus')
        distance = travel_info.get('distance_km', 0)
        cost_per_km = {
            'grab': 15000, 'taxi': 15000, 'bus': 3000,
            'train': 5000, 'flight': 8000, 'car': 10000,
        }
        return int(distance * cost_per_km.get(transport, 5000))

    def _get_multi_dest_tips(self, dests: list, language: str) -> list:
        """Tips for multi-destination trips."""
        tips = []
        if language == 'vi':
            tips.append(f"Lộ trình tối ưu: {' → '.join(d['name'] for d in dests)}")
            for i in range(len(dests) - 1):
                info = self.route_planner.estimate_inter_city_travel(dests[i], dests[i + 1])
                tips.append(
                    f"{dests[i]['name']} → {dests[i + 1]['name']}: "
                    f"~{info['distance_km']}km, ~{info['travel_time_min']} phút ({info['recommended_transport']})"
                )
        else:
            tips.append(f"Optimized route: {' → '.join(d.get('name_en', d['name']) for d in dests)}")
            for i in range(len(dests) - 1):
                info = self.route_planner.estimate_inter_city_travel(dests[i], dests[i + 1])
                tips.append(
                    f"{dests[i].get('name_en', dests[i]['name'])} → {dests[i + 1].get('name_en', dests[i + 1]['name'])}: "
                    f"~{info['distance_km']}km, ~{info['travel_time_min']}min ({info['recommended_transport']})"
                )
        return tips