import os
import json
import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from src.prompts.templates.flght_prompt import flight_prompt      # You must save the prompt function here
from src.config.llm_config import LlmConfig
from config import Config
from src.components.data_extraction.weather_data import DroneWeatherDataExtractor

class DroneAssessmentAgent:
    def __init__(self):
        self.llm = ChatOpenAI(
            api_key=Config.OPENAI_API_KEY,
            model=LlmConfig.openai.models.gpt_4o,
            temperature=0.3
        )

        self.weather_extractor = DroneWeatherDataExtractor()

    async def run(self, booking_form: str) -> dict:
        """
        Run the drone environmental & safety assessment agent.

        Args:
            booking_form_text (str): Raw text of the booking form

        Returns:
            dict: JSON output of the AI assessment
        """
        weather_data = await self.weather_extractor.extract_booking_weather_data(booking_form)
        booking_form_text = json.dumps(booking_form)
       
        prompt = flight_prompt(booking_form_text, weather_data)  # Contains the instructions + JSON template
        messages = [
            SystemMessage(content=prompt),
            HumanMessage(content=booking_form_text)
        ]

        try:
            response = self.llm.invoke(messages)

            if hasattr(response, "content"):
                response_text = response.content.strip()
                print("Raw LLM Response:\n", response_text[:300], "...\n")

                # Attempt to parse JSON from response
                start_idx = response_text.find('{')
                end_idx = response_text.rfind('}') + 1

                if start_idx == -1 or end_idx == -1:
                    raise ValueError("No JSON object found in output")

                json_str = response_text[start_idx:end_idx]
                return json.loads(json_str)

            else:
                raise ValueError("LLM returned no usable content")

        except Exception as e:
            print(f"Error in DroneAssessmentAgent: {str(e)}")
            return {
                "error": str(e),
                "raw_response": response.content if 'response' in locals() else None
            }

async def main():
    """Async main function to properly handle the async agent."""
    from test2 import booking_form_input
   
    agent = DroneAssessmentAgent()
    result = await agent.run(booking_form_input)

    print("\nStructured Assessment Output:\n")
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    asyncio.run(main())