from pydantic import BaseModel

from dataclasses import dataclass, field
from typing import List, Dict

@dataclass
class Planet:
    name: str
    description: str
    accommodation: str
    activities: List[str]
    photography: str

@dataclass
class Category:
    budget_range: str
    planets: List[Planet] = field(default_factory=list)

@dataclass
class TravelAgency(BaseModel):
    name: str
    categories: Dict[str, Category] = field(default_factory=dict)

    def add_planet(self, category_name: str, budget_range: str, planet: Planet):
        # Ensure category exists, create if not
        if category_name not in self.categories:
            self.categories[category_name] = Category(budget_range=budget_range)
        self.categories[category_name].planets.append(planet)


# class Planet(BaseModel):
#     def __init__(self, name, atmosphere, species, signature_dish, activities):
#         self.name = name
#         self.atmosphere = atmosphere
#         self.species = species
#         self.signature_dish = signature_dish
#         self.activities = activities

#     def describe(self):
#         return f"Planet: {self.name}\n" \
#                f"Atmosphere: {self.atmosphere}\n" \
#                f"Species: {self.species}\n" \
#                f"Signature Dish: {self.signature_dish}\n" \
#                f"Activities: {', '.join(self.activities)}"