{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE OverloadedStrings     #-}

module Main where

import Control.Monad (foldM, when)
import Data.Aeson (Value(..), (.=))
import Data.Text (Text)
import Data.Vector (Vector)
import GHC.Exts (toList)
import OpenAI.V1
import OpenAI.V1.Chat.Completions
import OpenAI.V1.Tool
import OpenAI.V1.ToolCall

import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as ByteString.Lazy
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text.Encoding
import qualified Data.Text.IO as Text.IO
import qualified Data.Vector as Vector
import qualified System.Environment as Environment

-- | Mock weather data for different cities
mockWeatherData :: Text -> Text
mockWeatherData city = case Text.toLower city of
  "london" -> "London: 18°C, partly cloudy, 60% humidity, light rain expected"
  "paris" -> "Paris: 22°C, sunny, 45% humidity, clear skies"
  "tokyo" -> "Tokyo: 28°C, humid, 80% humidity, thunderstorms possible"
  "new york" -> "New York: 25°C, partly sunny, 55% humidity, pleasant weather"
  "san francisco" -> "San Francisco: 16°C, foggy, 70% humidity, typical coastal weather"
  _ -> city <> ": Weather data not available for this location. Try London, Paris, Tokyo, New York, or San Francisco."

-- | Process a weather tool call and return the result
processWeatherToolCall :: ToolCall -> IO (Message (Vector Content))
processWeatherToolCall (ToolCall_Function toolCallId function) = do
  let functionName = OpenAI.V1.ToolCall.name function
      arguments = OpenAI.V1.ToolCall.arguments function

  putStrLn $ "🔧 Processing tool call: " <> Text.unpack functionName
  putStrLn $ "   Arguments: " <> Text.unpack arguments

  -- Parse arguments to extract city name
  let weatherResult = case Aeson.decode (ByteString.Lazy.fromStrict $ Text.Encoding.encodeUtf8 arguments) of
        Just (Object obj) -> case lookup "city" (toList obj) of
          Just (String cityName) -> mockWeatherData cityName
          _ -> "Error: Could not parse city name from arguments"
        _ -> "Error: Invalid arguments format"

  putStrLn $ "   Result: " <> Text.unpack weatherResult
  putStrLn ""

  -- Return a Tool message with the result
  return $
    Tool
      { content = [Text weatherResult],
        tool_call_id = toolCallId
      }

-- | Define the weather tool
weatherTool :: Tool
weatherTool =
  Tool_Function $
    OpenAI.V1.Tool.Function
      { description = Just "Get current weather information for a city",
        name = "get_weather",
        parameters =
          Just $
            Aeson.object
              [ "type" .= ("object" :: Text),
                "properties"
                  .= Aeson.object
                    [ "city"
                        .= Aeson.object
                          [ "type" .= ("string" :: Text),
                            "description" .= ("The city name to get weather for" :: Text)
                          ]
                    ],
                "required" .= (["city"] :: [Text]),
                "additionalProperties" .= False
              ],
        strict = Just True
      }

-- | Convert a Message Text to Message (Vector Content)
convertMessage :: Message Text -> Message (Vector Content)
convertMessage (System {content, name}) = System {content = [Text content], name = name}
convertMessage (User {content, name}) = User {content = [Text content], name = name}
convertMessage (Assistant {assistant_content, refusal, name, assistant_audio, tool_calls}) =
  Assistant
    { assistant_content = fmap (\c -> [Text c]) assistant_content,
      refusal = refusal,
      name = name,
      assistant_audio = assistant_audio,
      tool_calls = tool_calls
    }
convertMessage (Tool {content, tool_call_id}) = Tool {content = [Text content], tool_call_id = tool_call_id}

-- | Process all tool calls in a message and return tool response messages
processToolCalls :: Vector ToolCall -> IO (Vector (Message (Vector Content)))
processToolCalls toolCalls = do
  when (not $ Vector.null toolCalls) $ do
    putStrLn $ "🤖 Processing " <> show (Vector.length toolCalls) <> " tool call(s)..."
  Vector.mapM processWeatherToolCall toolCalls

-- | Main chat loop
chatLoop :: (CreateChatCompletion -> IO ChatCompletionObject) -> Vector (Message (Vector Content)) -> IO ()
chatLoop createChatCompletion messages = do
  Text.IO.putStr "You: "
  userInput <- Text.IO.getLine

  -- Exit condition
  when (Text.toLower userInput == "quit" || Text.toLower userInput == "exit") $ do
    putStrLn "Goodbye! 👋"
    return ()

  -- Add user message to conversation
  let userMessage =
        User
          { content = [Text userInput],
            name = Nothing
          }
      updatedMessages = messages <> [userMessage]

  -- Make API call with tool support
  response <-
    createChatCompletion
      _CreateChatCompletion
        { messages = updatedMessages,
          model = "gpt-4o-mini",
          tools = Just [weatherTool],
          tool_choice = Just ToolChoiceAuto
        }

  let ChatCompletionObject {choices} = response

  -- Process each choice (usually just one)
  newMessages <- foldM (processChoice createChatCompletion) updatedMessages choices

  -- Continue the conversation
  chatLoop createChatCompletion newMessages

-- | Process a single choice, handling potential tool calls
processChoice ::
  (CreateChatCompletion -> IO ChatCompletionObject) ->
  Vector (Message (Vector Content)) ->
  Choice ->
  IO (Vector (Message (Vector Content)))
processChoice createChatCompletion messages choice = do
  let Choice {message = assistantMessage} = choice

  case assistantMessage of
    Assistant {assistant_content, tool_calls = Just toolCalls} -> do
      -- Display assistant's message if any
      case assistant_content of
        Just content -> putStrLn $ "Assistant: " <> Text.unpack content
        Nothing -> return ()

      -- Add assistant message to history
      let messagesWithAssistant = messages <> [convertMessage assistantMessage]

      -- Process tool calls
      toolResults <- processToolCalls toolCalls
      let messagesWithTools = messagesWithAssistant <> toolResults

      -- Make another API call to get the final response after tool calls
      finalResponse <-
        createChatCompletion
          _CreateChatCompletion
            { messages = messagesWithTools,
              model = "gpt-4o-mini",
              tools = Just [weatherTool],
              tool_choice = Just ToolChoiceAuto
            }

      let ChatCompletionObject {choices = finalChoices} = finalResponse

      -- Process the final response (this should not have tool calls)
      foldM processFinalChoice messagesWithTools finalChoices
    Assistant {assistant_content = Just content} -> do
      -- No tool calls, just display the response
      putStrLn $ "Assistant: " <> Text.unpack content
      return $ messages <> [convertMessage assistantMessage]
    _ -> do
      putStrLn "Assistant: (No response content)"
      return messages

-- | Process final choice after tool calls (should just be text response)
processFinalChoice :: Vector (Message (Vector Content)) -> Choice -> IO (Vector (Message (Vector Content)))
processFinalChoice messages choice = do
  let Choice {message = finalMessage} = choice
  case finalMessage of
    Assistant {assistant_content = Just content} -> do
      putStrLn $ "Assistant: " <> Text.unpack content
      return $ messages <> [convertMessage finalMessage]
    _ -> do
      putStrLn "Assistant: (Unexpected response format)"
      return messages

main :: IO ()
main = do
  putStrLn "🌤️  Weather Chatbot with Tool Calling"
  putStrLn "======================================"
  putStrLn "Ask me about the weather in different cities!"
  putStrLn "Try: 'What's the weather like in London?'"
  putStrLn "Available cities: London, Paris, Tokyo, New York, San Francisco"
  putStrLn "Type 'quit' or 'exit' to end the conversation."
  putStrLn ""

  -- Get OpenAI API key
  key <- Environment.getEnv "OPENAI_KEY"

  -- Set up client
  clientEnv <- getClientEnv "https://api.openai.com"
  let Methods {createChatCompletion} = makeMethods clientEnv (Text.pack key) Nothing Nothing

  -- Initial system message
  let systemMessage =
        System
          { content = [Text "You are a helpful weather assistant. Use the get_weather tool to provide current weather information when users ask about weather in specific cities. Be conversational and helpful."],
            name = Nothing
          }

  -- Start the chat loop
  chatLoop createChatCompletion [systemMessage]