-- PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
-- Copyright 2025 TimeHexOn & foxhop & russell@unturf
-- https://www.permacomputer.com

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- UncloseAI Haskell Library
-- OpenAI-compatible API client with streaming support
-- Compatible with vLLM, Ollama, and OpenAI-compatible endpoints

import Network.HTTP.Simple
import Network.HTTP.Client (responseBody)
import Data.Aeson
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import qualified Data.Text.Encoding as TE
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as BS
import GHC.Generics
import Control.Exception
import Control.Monad (unless)
import Control.Monad.IO.Class (liftIO)
import System.IO
import System.Environment (lookupEnv)
import Data.Maybe (fromMaybe, isJust)
import Data.Conduit
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.Combinators as CC

-- Model info type
data ModelInfo = ModelInfo
  { modelId :: Text
  , modelEndpoint :: String
  , modelMaxTokens :: Int
  } deriving (Show)

-- UncloseAI Client type
data UncloseAIClient = UncloseAIClient
  { clientModels :: [ModelInfo]
  , clientTtsEndpoints :: [String]
  , clientTimeout :: Int
  } deriving (Show)

-- Message types
data ChatMessage = ChatMessage
  { role :: Text
  , content :: Text
  } deriving (Generic, Show)

instance ToJSON ChatMessage

data ChatRequest = ChatRequest
  { model :: Text
  , messages :: [ChatMessage]
  , max_tokens :: Int
  , stream :: Maybe Bool
  } deriving (Generic, Show)

instance ToJSON ChatRequest where
  toJSON (ChatRequest m msgs mt s) = object $
    [ "model" .= m
    , "messages" .= msgs
    , "max_tokens" .= mt
    ] ++ case s of
      Just True -> ["stream" .= True]
      _ -> []

data ChatResponse = ChatResponse
  { choices :: [Choice]
  } deriving (Generic, Show)

data Choice = Choice
  { message :: ResponseMessage
  } deriving (Generic, Show)

data ResponseMessage = ResponseMessage
  { respContent :: Text
  } deriving (Generic, Show)

instance FromJSON ChatResponse
instance FromJSON Choice
instance FromJSON ResponseMessage where
  parseJSON = withObject "ResponseMessage" $ \v ->
    ResponseMessage <$> v .: "content"

data TTSRequest = TTSRequest
  { tts_model :: Text
  , voice :: Text
  , input :: Text
  } deriving (Show)

instance ToJSON TTSRequest where
  toJSON (TTSRequest m v i) = object
    [ "model" .= m
    , "voice" .= v
    , "input" .= i
    ]

-- Models discovery types
data ModelData = ModelData
  { mdId :: Text
  , mdMaxModelLen :: Maybe Int
  } deriving (Generic, Show)

instance FromJSON ModelData where
  parseJSON = withObject "ModelData" $ \v ->
    ModelData
      <$> v .: "id"
      <*> v .:? "max_model_len"

data ModelsResponse = ModelsResponse
  { modelsData :: [ModelData]
  } deriving (Generic, Show)

instance FromJSON ModelsResponse where
  parseJSON = withObject "ModelsResponse" $ \v ->
    ModelsResponse <$> v .: "data"

-- Initialize client with auto-discovery
initClient :: Int -> IO UncloseAIClient
initClient timeout = do
  putStrLn "Initializing UncloseAI client..."

  -- Discover chat/code models
  models <- discoverModelsLoop 1 []

  -- Discover TTS endpoints
  ttsEndpoints <- discoverTtsLoop 1 []

  putStrLn $ "Discovered " ++ show (length models) ++ " models, " ++
             show (length ttsEndpoints) ++ " TTS endpoints\n"

  return $ UncloseAIClient
    { clientModels = models
    , clientTtsEndpoints = ttsEndpoints
    , clientTimeout = timeout
    }

-- Model discovery
discoverModelsFromEndpoint :: String -> IO [ModelInfo]
discoverModelsFromEndpoint ep = do
  putStrLn $ "Endpoint: " ++ ep

  result <- try $ do
    request <- parseRequest $ "GET " ++ ep ++ "/models"
    response <- httpLBS request
    return $ getResponseBody response

  case result of
    Left (e :: SomeException) -> return []
    Right body ->
      case decode body :: Maybe ModelsResponse of
        Nothing -> return []
        Just modelsResp -> do
          let modelsList = modelsData modelsResp
              -- Filter out modelperm-* entries
              filtered = filter (\md -> not $ T.isPrefixOf "modelperm-" (mdId md)) modelsList
          mapM (\md -> do
            let maxToks = fromMaybe 8192 (mdMaxModelLen md)
            return $ ModelInfo (mdId md) ep maxToks
            ) filtered

discoverModelsLoop :: Int -> [ModelInfo] -> IO [ModelInfo]
discoverModelsLoop i acc | i > 9999 = return $ reverse acc
discoverModelsLoop i acc = do
  maybeEndpoint <- lookupEnv $ "MODEL_ENDPOINT_" ++ show i
  case maybeEndpoint of
    Nothing -> return $ reverse acc
    Just ep -> do
      newModels <- discoverModelsFromEndpoint ep
      discoverModelsLoop (i + 1) (reverse newModels ++ acc)

discoverTtsLoop :: Int -> [String] -> IO [String]
discoverTtsLoop i acc | i > 9999 = return $ reverse acc
discoverTtsLoop i acc = do
  maybeEndpoint <- lookupEnv $ "TTS_ENDPOINT_" ++ show i
  case maybeEndpoint of
    Nothing -> return $ reverse acc
    Just ep -> do
      putStrLn $ "Discovering TTS from: " ++ ep
      discoverTtsLoop (i + 1) (ep : acc)

-- Streaming response types
data StreamDelta = StreamDelta
  { deltaContent :: Maybe Text
  } deriving (Generic, Show)

instance FromJSON StreamDelta where
  parseJSON = withObject "StreamDelta" $ \v ->
    StreamDelta <$> v .:? "content"

data StreamChoice = StreamChoice
  { delta :: StreamDelta
  } deriving (Generic, Show)

instance FromJSON StreamChoice

data StreamChunk = StreamChunk
  { streamChoices :: [StreamChoice]
  } deriving (Generic, Show)

instance FromJSON StreamChunk where
  parseJSON = withObject "StreamChunk" $ \v ->
    StreamChunk <$> v .: "choices"

-- Non-streaming chat completion
chat :: UncloseAIClient -> [ChatMessage] -> Maybe Int -> Maybe Int -> Maybe Double -> IO (Either String Text)
chat client msgs maybeModelIdx maybeMaxToks maybeTemp = do
  let modelIdx = fromMaybe 0 maybeModelIdx
      maxToks = fromMaybe 100 maybeMaxToks
      temp = fromMaybe 0.7 maybeTemp
      models = clientModels client

  if modelIdx >= length models
    then return $ Left "Invalid model index"
    else do
      let modelInfo = models !! modelIdx
      let req = ChatRequest
            { model = modelId modelInfo
            , messages = msgs
            , max_tokens = maxToks
            , stream = Nothing
            }

      result <- try $ do
        request <- parseRequest $ "POST " ++ modelEndpoint modelInfo ++ "/chat/completions"
        let request' = setRequestBodyJSON req request
        response <- httpLBS request'
        return $ getResponseBody response

      case result of
        Right body ->
          case decode body :: Maybe ChatResponse of
            Just resp ->
              case choices resp of
                (c:_) -> return $ Right $ respContent $ message c
                [] -> return $ Left "No response choices"
            Nothing -> return $ Left "Could not parse response"
        Left (e :: SomeException) -> return $ Left $ show e

-- Streaming chat completion - yields content via IO action
chatStream :: UncloseAIClient -> [ChatMessage] -> Maybe Int -> Maybe Int -> Maybe Double -> IO (Either String ())
chatStream client msgs maybeModelIdx maybeMaxToks maybeTemp = do
  let modelIdx = fromMaybe 0 maybeModelIdx
      maxToks = fromMaybe 500 maybeMaxToks
      temp = fromMaybe 0.7 maybeTemp
      models = clientModels client

  if modelIdx >= length models
    then return $ Left "Invalid model index"
    else do
      let modelInfo = models !! modelIdx
      let req = ChatRequest
            { model = modelId modelInfo
            , messages = msgs
            , max_tokens = maxToks
            , stream = Just True
            }

      result <- try $ do
        request <- parseRequest $ "POST " ++ modelEndpoint modelInfo ++ "/chat/completions"
        let request' = setRequestBodyJSON req request
        withResponse request' $ \response ->
          runConduit $ getResponseBody response
            .| CC.linesUnboundedAscii
            .| CL.mapM_ processSSELine

      case result of
        Right () -> return $ Right ()
        Left (e :: SomeException) -> return $ Left $ show e

-- Text-to-speech generation
tts :: UncloseAIClient -> Text -> Maybe Text -> Maybe String -> IO (Either String String)
tts client text maybeVoice maybeOutputFile = do
  let voice = fromMaybe "alloy" maybeVoice
      outputFile = fromMaybe "/tmp/speech.mp3" maybeOutputFile
      ttsEndpoints = clientTtsEndpoints client

  if null ttsEndpoints
    then return $ Left "No TTS endpoints available"
    else do
      let endpoint = head ttsEndpoints
      let req = TTSRequest
            { tts_model = "tts-1"
            , voice = voice
            , input = text
            }

      result <- try $ do
        request <- parseRequest $ "POST " ++ endpoint ++ "/audio/speech"
        let request' = setRequestBodyJSON req request
        response <- httpLBS request'
        let body = getResponseBody response
        BL.writeFile outputFile body
        return outputFile

      case result of
        Right file -> return $ Right file
        Left (e :: SomeException) -> return $ Left $ show e

-- Process SSE line
processSSELine :: BS.ByteString -> IO ()
processSSELine line
  | BS.isPrefixOf "data: " line = do
      let dataStr = BS.drop 6 line
      unless (dataStr == "[DONE]") $ do
        case decode (BL.fromStrict dataStr) :: Maybe StreamChunk of
          Just chunk ->
            case streamChoices chunk of
              (c:_) ->
                case deltaContent (delta c) of
                  Just content -> TIO.putStr content >> hFlush stdout
                  Nothing -> return ()
              [] -> return ()
          Nothing -> return ()
  | otherwise = return ()

-- Demo program showing library usage
main :: IO ()
main = do
  hSetBuffering stdout NoBuffering
  putStrLn "=== UncloseAI Haskell Client (with Streaming) ===\n"

  -- Initialize client
  client <- initClient 30

  if null (clientModels client)
    then do
      putStrLn "ERROR: No models discovered"
    else do
      let models = clientModels client
      let firstModel = head models

      -- Non-streaming chat example
      putStrLn "=== Non-Streaming Chat ==="
      putStrLn $ "Model: " ++ T.unpack (modelId firstModel)

      let messages = [ChatMessage "user" "Explain quantum computing in one sentence"]
      result <- chat client messages Nothing Nothing Nothing
      case result of
        Right response -> putStrLn $ "Response: " ++ T.unpack response ++ "\n"
        Left err -> putStrLn $ "Error: " ++ err ++ "\n"

      -- Streaming chat example
      let modelIdx = if length models >= 2 then 1 else 0
      let streamModel = models !! modelIdx

      putStrLn "=== Streaming Chat ==="
      putStrLn $ "Model: " ++ T.unpack (modelId streamModel)
      putStr "Response: "

      let streamMessages = [ChatMessage "user" "Write a hello world program in Haskell"]
      streamResult <- chatStream client streamMessages (Just modelIdx) Nothing Nothing
      case streamResult of
        Right () -> putStrLn "\n"
        Left err -> putStrLn $ "\nError: " ++ err ++ "\n"

      -- TTS example
      if not (null (clientTtsEndpoints client))
        then do
          putStrLn "=== TTS Speech Generation ==="
          putStrLn "Model: tts-1"

          ttsResult <- tts client "Hello from UncloseAI Haskell client!" Nothing (Just "/tmp/speech.mp3")
          case ttsResult of
            Right file -> putStrLn $ "Audio saved to " ++ file
            Left err -> putStrLn $ "TTS failed: " ++ err
        else return ()

      putStrLn "\n=== Examples Complete ==="
