#! /usr/bin/env nix-shell
#! nix-shell --keep --pure -i runghc
#! nix-shell --keep --pure -p cacert curl git
#! nix-shell --keep --pure -p '(haskellPackages.extend (self: super: {pin = self.callHackageDirect { pkg = "A-gent"; ver = "0.11.0.18"; sha256 = "BsvlGOludmvVvfjLYvjBkQumj2VuIe6M8kB/9CYkLlo=";} {};})).ghcWithPackages (ps: with ps; [ pin ])'

{-# OPTIONS_GHC -Wall -Werror -fno-warn-orphans #-}

{-# OPTIONS_GHC -frefinement-level-hole-fits=1 #-}

{-# LANGUAGE NoGeneralizedNewtypeDeriving #-}
{-# LANGUAGE Safe                         #-}

--------------------------------------------------------------------------------
--
-- Λ-gent, (c) 2026 SPISE MISU ApS, https://spdx.org/licenses/SSPL-1.0
--
--------------------------------------------------------------------------------

import           Prelude             hiding ( error, lines, mod )

import           Data.Either         ( partitionEithers, rights )
import           Data.List           ( findIndex, isPrefixOf, tails )

import qualified Agent.IO.Effects    as EFF

import           Agent.Data.JSON     ( Data )
import qualified Agent.Data.JSON     as JSON
import           Agent.IO.Restricted ( RIO )
import qualified Agent.IO.Restricted as RIO
import           Agent.LLM
  ( Context (Context, load, mode)
  , Eval
  , Mode (Code)
  , replWithMode
  )
import qualified Agent.LLM.Action    as ACT
import qualified Agent.LLM.Message   as MSG

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

data Message = Message
  { role    :: !String
  , content :: !String
  } deriving (Data, Show)

--------------------------------------------------------------------------------

data Request = Request
  { model       :: !String
  , max_tokens  :: !Int
  , temperature :: !Double
  , messages    :: ![Message]
  } deriving (Data, Show)

--------------------------------------------------------------------------------

data Response = Response
  { created :: !Int
  , choices :: ![Choice]
  , usage   :: !Usage
  } deriving (Data, Show)

data Choice = Choice
  { index         :: !Int
  , finish_reason :: !String
  , message       :: !Message
  } deriving (Data, Show)

data Usage = Usage
  { prompt_tokens     :: !Int
  , completion_tokens :: !Int
  , total_tokens      :: !Int
  } deriving (Data, Show)

data Error = Error
  { error :: String
  } deriving (Data, Show)

data Communication = Communication
  [Either Error Message] deriving (Data, Show)

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

eval :: Eval Communication
eval ctx msg =
  case ctx of
    Context { mode = Code } -> code    ctx           msg
    _______________________ -> return (ctx, ACT.text err)
  where
    mod = mode ctx
    err = show mod ++ " mode is not available for this agent"

--------------------------------------------------------------------------------

code
  ::
    ( EFF.LlmCodeRead rio
    , EFF.LlmCodeTmpl rio
    , EFF.LlmCodePost rio
    , EFF.LlmCodeSave rio
    )
  => Context Communication
  -> MSG.Message
  -> rio (Context Communication, ACT.Action)
code ctx msg =

  case msg of

    MSG.Atom txt fs ->
      EFF.llmCodePut txt fs >>= \ ecf ->
      case ecf of
        Right o ->
          return
            ( ctx
            , ACT.text o
            )
        Left  e ->
          return
            ( ctx
            , ACT.text e
            )

    MSG.List mfil ->
      EFF.llmCodeSeq mfil >>= \ ecfs ->
      case ecfs of
        Right fps ->
          return
            ( ctx
            , ACT.paths fps
            )
        Left errs ->
          return
            ( ctx
            , (ACT.text . unlines) errs
            )

    MSG.Path _ Nothing ->
      return
        ( ctx
        , ACT.text "Index is out of bounds"
        )

    MSG.Path _ (Just afp) ->
      EFF.llmCodeGet afp >>= \ ecf ->
      case ecf of
        Right f ->
          return
            ( ctx
            , ACT.file f
            )
        Left  e ->
          return
            ( ctx
            , ACT.text e
            )

    MSG.Repo ->
      EFF.llmCodeGit >>= \ ebs ->
      case ebs of
        Right o ->
          return
            ( ctx
            , ACT.text o
            )
        Left  e ->
          return
            ( ctx
            , ACT.text e
            )

    MSG.Root ->
      EFF.llmPathCWD >>= \ eroot ->
      case eroot of
        Just root ->
          return
            ( ctx
            , ACT.root root
            )
        Nothing   ->
          return
            ( ctx
            , ACT.none
            )

    MSG.Send xmls ->
      -- NOTE: 37.7 GB LLM (max context length 256k)
      payloadToLLM ctx xmls  262144 0.7 "mlx-community/Qwen3.6-35B-A3B-8bit"
      -- NOTE: 54.70 GB LLM (max context length 256k)
      --payloadToLLM ctx xmls  262144 0.7 "mlx-community/Qwen3.6-27B-bf16"

    MSG.Text txt ->
      EFF.llmCodeWeb (JSON.encode req) >>= \ eres ->
        case eres of
          Right json ->
            return
              ( nxt
              , case res of
                  Left  err ->
                    ACT.text $ error err
                  Right val ->
                    ACT.text $ concatMap (show . message) $ choices val
              )
            where
              nxt =
                ctx
                  { load = Just l
                  }
                where
                  l =
                    case com of
                      Communication ems ->
                        case res of
                          Left  err ->
                            Communication $ ems ++
                            [Left err]
                          Right val ->
                            Communication $ ems ++
                            (map (Right . message) $ choices val)
              res =
                case JSON.decode json of
                  Right a                -> Right a
                  Left  JSON.InvalidJSON -> Left $ Error $ "Invalid: " ++ json
                  Left  JSON.DiffSchema  ->
                    case JSON.decode json of
                      Right e                -> Left e
                      Left  JSON.InvalidJSON -> Left $ Error $ "Invalid: " ++ json
                      Left  JSON.DiffSchema  -> Left $ Error $ "Schema: "  ++ json
          Left err ->
            return (nxt, ACT.text err)
            where
              nxt =
                ctx
                  { load = Just l
                  }
                where
                  l =
                    case com of
                      Communication ems ->
                        Communication $ ems ++ [Left $ Error err]
      where
        req =
          -- NOTE: 20.4 GB LLM (max context length 256k)
          -- "mlx-community/Qwen3.6-35B-A3B-4bit"
          -- NOTE: 32.5 GB LLM (max context length 256k)
          -- "mlx-community/Qwen3-30B-A3B-8bit"
          Request
            { model       = "mlx-community/Qwen3.6-35B-A3B-4bit"
            , max_tokens  = 262144
            , temperature = 0.7
            , messages    = rs
            }
          where
            rs =
              case com of
                Communication ems -> rights ems
        com =
            case load ctx of
              Just (Communication ems) -> Communication $ ems ++ [Right erm]
              Nothing                  -> Communication          [Right erm]
        erm =
          Message
            { role      = "user"
            , content   = txt
            }

    MSG.Tmpl (Just afps) ->
           EFF.llmCodeIns      >>= \  cis ->
           EFF.llmCodeExa      >>= \  ces ->
      mapM EFF.llmCodeGet afps >>= \ ecfs ->
      case partitionEithers ecfs of
        ( [  ], cfs ) ->
          return
            ( ctx
            , ACT.template parseFiles cis ces cfs
            )
        ( errs, ___ ) ->
          return
            ( ctx
            , (ACT.text . unlines) errs
            )

    ________ ->
      return
        ( ctx
        , ACT.text
          $ "The following message '" ++ show msg ++ "', is not supported"
        )

--------------------------------------------------------------------------------

main :: IO ()
main =
  replWithMode Code eval

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- NOTE: Analogy to H(P)FI relays and fuse groups.

--------------------------------------------------------------------------------

instance EFF.LlmConf RIO where
  llmPathCWD = RIO.llmPathCWD

--------------------------------------------------------------------------------

instance EFF.LlmCodeRoot RIO where
  llmCodeDir = RIO.llmCodeDir

instance EFF.LlmCodeMask RIO where
  llmCodeMsk = RIO.llmCodeMsk

instance EFF.LlmCodeTmpl RIO where
  llmCodeIns = RIO.llmCodeIns
  llmCodeExa = RIO.llmCodeExa

--------------------------------------------------------------------------------

instance EFF.LlmCodeRead RIO where
  llmCodeSeq = RIO.llmCodeSeq
  llmCodeGet = RIO.llmCodeGet
  llmCodeGit = RIO.llmCodeGit

--------------------------------------------------------------------------------

instance EFF.LlmCodeSave RIO where
  llmCodePut = RIO.llmCodePut

--------------------------------------------------------------------------------

instance EFF.LlmCodeConf RIO where
  llmCodeAPI = RIO.llmCodeAPI
  llmCodeKey = RIO.llmCodeKey

instance EFF.LlmCodePost RIO where
  llmCodeWeb = RIO.llmCodeWeb

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- HELPERS

--------------------------------------------------------------------------------

payloadToLLM
  :: EFF.LlmCodePost rio
  => Context a
  -> [String]
  -> Int
  -> Double
  -> String
  -> rio (Context a, ACT.Action)
payloadToLLM ctx txts tok tem mod =
  EFF.llmCodeWeb (JSON.encode req) >>= \ eres ->
    case eres of
      Right json ->
        return
          ( ctx
          , case res of
              Left  err ->
                ACT.text $ error err
              Right val ->
                ACT.text $ concatMap (content . message) $ choices val
          )
        where
          res =
            case JSON.decode json of
              Right a                -> Right a
              Left  JSON.InvalidJSON -> Left $ Error $ "Invalid: " ++ json
              Left  JSON.DiffSchema  ->
                case JSON.decode json of
                  Right e                ->
                    Left e
                  Left  JSON.InvalidJSON ->
                    Left $ Error $ "Invalid: "      ++ json
                  Left  JSON.DiffSchema  ->
                    Left $ Error $ "Diff Schema: "  ++ json
      Left err ->
        return (ctx, ACT.text err)
  where
    req =
      Request
        { model              = mod
        , max_tokens         = tok
        , temperature        = tem
        , messages           =
            map
              ( \ txt ->
                  Message
                    { role      = "user"
                    , content   = txt
                    }
              )
              txts
        }

--------------------------------------------------------------------------------

parseFiles
  :: String
  -> [(FilePath, [String])]
parseFiles txt =
  (snd . partitionEithers) $ aux (fn sp np) np
  where
    sp = "```\n"
    np = txt ++ "\n"
    lp = length sp
    fn = (. tails) . findIndex . isPrefixOf
    aux midx cs =
      case midx of
        Nothing  -> []
        Just idx ->
          par (take idx cs) : aux (fn sp rs) rs
          where
            rs = drop (idx + lp) cs
    par cs =
      case JSON.decode (g cs) :: Either JSON.DecodeError [String] of
        Right fls ->
          Right
            ( f cs
            , fls
            )
        Left  err ->
          case err of
            JSON.InvalidJSON -> Left "Invalid JSON"
            JSON.DiffSchema  -> Left "Diff JSON Schema"
      where
        f = takeWhile (/= '\n') . drop 2 . dropWhile (/= '#')
        g = drop 1 . dropWhile (/= '\n') . dropWhile (/= '`')

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------