import sys
import os
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(script_dir, "../../.."))


if project_root not in sys.path:
    sys.path.insert(0, project_root)


from HeGFlow.graph.tool_node import Tool_node
from typing import Union
from typing import Any,Dict,List
from local_datasets.TravelPlanner.tools.flights.apis import Flights
import json
import asyncio
import pandas as pd


class Flights_tool(Tool_node):
    def __init__(self):
        super().__init__(
            name="Flights",
            description="a tool for flight'。"
        )
        self._city_set = set()
        self._state_set = set()
        self._state_to_cities = {} 

        try:
            with open("", "r", encoding="utf-8") as f:
                for line in f:
                    self._city_set.add(line.strip())
        except FileNotFoundError:
            print("Warning: citySet.txt not found. City lookup will be incomplete.")
            


        try:
            with open("", "r", encoding="utf-8") as f:
                for line in f:
                    self._state_set.add(line.strip())
        except FileNotFoundError:
            print("Warning: stateSet.txt not found. State lookup will be incomplete.")

        try:
            with open("", "r", encoding="utf-8") as f:
                for line in f:
                    parts = line.strip().split('\t')
                    if len(parts) == 2:
                        city, state = parts
                        if state not in self._state_to_cities:
                            self._state_to_cities[state] = []
                        self._state_to_cities[state].append(city)
                    else:
                        print(f"Warning: Malformed line in citySet_with_stateSet.txt: {line.strip()}")
        except FileNotFoundError:
            print("Warning: citySet_with_stateSet.txt not found. State-to-city mapping will be empty.")

        self.flights_api = Flights(path="")

    

    def _is_city(self,location):
       
        return location in self._city_set

    def _is_state(self,location):
       
        return location in self._state_set

    def _get_cities_in_state(self,state_name):
       
        return self._state_to_cities.get(state_name, []) 
    
    async def _execute_tool(self, task:Dict[str,str]) -> str:

       
        if self.flights_api.data is not None:  
            
            origin_city = task["org"]
            destination_city = task["dest"]
            flight_date = task["date"]  

            departure_date = flight_date[0]
            return_date = flight_date[-1]
            print(departure_date)
            print(return_date)

            actual_origins = []
            actual_destinations = []

            if self._is_city(origin_city): 
                actual_origins.append(origin_city)

            elif self._is_state(origin_city):
                actual_origins.extend(self._get_cities_in_state(origin_city))

            if self._is_city(destination_city):
                actual_destinations.append(destination_city)
            
            elif self._is_state(destination_city):
                actual_destinations.extend(self._get_cities_in_state(destination_city))

            
            all_results: List[Dict[str, Any]] = []
            all_results_return: List[Dict[str, Any]] = []

            if not actual_origins or not actual_destinations:
                print(f"Warning: Could not identify valid origin cities for '{origin_city}' or destination cities for '{destination_city}'. Returning empty results.")
                return "flight not found"

            #results = flights_api.run_for_annotation(origin_city, destination_city, flight_date)
            for org_c in actual_origins:
                for dest_c in actual_destinations:
                    
                    if org_c == dest_c:
                        continue
                    
                    print(f"Querying flights from {org_c} to {dest_c} on {departure_date}")
                   
                    current_results = self.flights_api.run(org_c, dest_c, departure_date)
                    current_results_return = self.flights_api.run(dest_c, org_c, return_date)
                    
                    if isinstance(current_results, pd.DataFrame):
                        if not current_results.empty:
                            flight_data = current_results.to_dict(orient='records')
                            all_results.extend(flight_data)
                        else:
                            
                             print(f"No flights found for {org_c} to {dest_c} on {departure_date} (empty DataFrame).")

                    if isinstance(current_results_return, pd.DataFrame):
                        if not current_results_return.empty:
                            flight_data_return = current_results_return.to_dict(orient='records')
                            all_results_return.extend(flight_data_return)
                        else:
                          
                             print(f"No flights found for {dest_c} to {org_c} on {return_date} (empty DataFrame).")   

            
            if not all_results:
               
                return "No flights matching your request were found!" 

            
            formatted_flight_strings = []
            for i, flight in enumerate(all_results):
               
                flight_number = flight.get('Flight Number', 'N/A')
                price = flight.get('Price', 'N/A')
                dep_time = flight.get('DepTime', 'N/A')
                arr_time = flight.get('ArrTime', 'N/A')
                actualelapsed_time = flight.get('ActualElapsedTime', 'N/A')
                origin_name = flight.get('OriginCityName', 'N/A')
                dest_name = flight.get('DestCityName', 'N/A')
                flight_date_str = flight.get('FlightDate', 'N/A')
                distance = flight.get('Distance', 'N/A')

               
                flight_string = (
                    f"{i+1}. Flight Number：{flight_number},Price：{price},"
                    f"DepTime：{dep_time},ArrTime：{arr_time},ActualElapsedTime：{actualelapsed_time},"
                    f"OriginCityName：{origin_name},DestCityName：{dest_name},FlightDate：{flight_date_str},Distance：{distance}"
                )
                formatted_flight_strings.append(flight_string)

           
            final_output_string = "\n".join(formatted_flight_strings)

            formatted_flight_return_strings = []
            for i, flight in enumerate(all_results_return):
              
                flight_number = flight.get('Flight Number', 'N/A')
                price = flight.get('Price', 'N/A')
                dep_time = flight.get('DepTime', 'N/A')
                arr_time = flight.get('ArrTime', 'N/A')
                actualelapsed_time = flight.get('ActualElapsedTime', 'N/A')
                origin_name = flight.get('OriginCityName', 'N/A')
                dest_name = flight.get('DestCityName', 'N/A')
                flight_date_str = flight.get('FlightDate', 'N/A')
                distance = flight.get('Distance', 'N/A')

               
                flight_string = (
                    f"{i+1}. Flight Number：{flight_number},Price：{price},"
                    f"DepTime：{dep_time},ArrTime：{arr_time},ActualElapsedTime：{actualelapsed_time},"
                    f"OriginCityName：{origin_name},DestCityName：{dest_name},FlightDate：{flight_date_str},Distance：{distance}"
                )
                formatted_flight_return_strings.append(flight_string)

           
            final_output_return_string = "\n".join(formatted_flight_return_strings)
            prompt_1 = "Below is a list of available departure flights"
            prompt_2 = "Below is a list of available return flights"

            result = prompt_1+"\n"+final_output_string + "\n" + prompt_2+"\n"+final_output_return_string

            return result
        else:
            print("Error: Flights API data could not be loaded. Returning empty results.")
            return "ERROR" 

    

            
