{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

{- |
Module      : Data.Ollama.Common.SchemaBuilder
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : DSL for constructing structured JSON Schemas for Ollama's structured output API.

== Overview

This module defines a simple schema builder DSL for programmatically constructing
JSON Schemas compatible with the structured output features in the Ollama API.

It supports nested objects, arrays, required fields, and custom types, and
provides infix operators for a fluent and expressive syntax.

== Example

@
import Data.Ollama.Common.SchemaBuilder

let schema =
      emptyObject
        |+ ("name", JString)
        |+ ("age", JInteger)
        |++ ("address", buildSchema $
              emptyObject
                |+ ("city", JString)
                |+ ("zip", JInteger)
                |! "city"
            )
        |!! ["name", "age"]
        & buildSchema

printSchema schema
@
-}
module Data.Ollama.Common.SchemaBuilder
  ( -- * Core Types
    JsonType (..)
  , Property (..)
  , Schema (..)

    -- * Schema Construction
  , emptyObject
  , addProperty
  , addObjectProperty
  , requireField
  , requireFields
  , buildSchema

    -- * Schema Utilities
  , objectOf
  , arrayOf
  , toOllamaFormat
  , printSchema

    -- * Infix Schema DSL
  , (|+)
  , (|++)
  , (|!)
  , (|!!)
  ) where

import Data.Aeson
import Data.Map.Strict qualified as HM
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Lazy qualified as TL
import Data.Text.Lazy.Encoding qualified as T
import GHC.Generics

-- | Supported JSON types for schema generation.
data JsonType
  = JString
  | JNumber
  | JInteger
  | JBoolean
  | JNull
  | -- | Array of a specific type
    JArray JsonType
  | -- | Nested object schema
    JObject Schema
  deriving (Int -> JsonType -> ShowS
[JsonType] -> ShowS
JsonType -> String
(Int -> JsonType -> ShowS)
-> (JsonType -> String) -> ([JsonType] -> ShowS) -> Show JsonType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> JsonType -> ShowS
showsPrec :: Int -> JsonType -> ShowS
$cshow :: JsonType -> String
show :: JsonType -> String
$cshowList :: [JsonType] -> ShowS
showList :: [JsonType] -> ShowS
Show, JsonType -> JsonType -> Bool
(JsonType -> JsonType -> Bool)
-> (JsonType -> JsonType -> Bool) -> Eq JsonType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: JsonType -> JsonType -> Bool
== :: JsonType -> JsonType -> Bool
$c/= :: JsonType -> JsonType -> Bool
/= :: JsonType -> JsonType -> Bool
Eq, (forall x. JsonType -> Rep JsonType x)
-> (forall x. Rep JsonType x -> JsonType) -> Generic JsonType
forall x. Rep JsonType x -> JsonType
forall x. JsonType -> Rep JsonType x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. JsonType -> Rep JsonType x
from :: forall x. JsonType -> Rep JsonType x
$cto :: forall x. Rep JsonType x -> JsonType
to :: forall x. Rep JsonType x -> JsonType
Generic)

instance ToJSON JsonType where
  toJSON :: JsonType -> Value
toJSON JsonType
JString = Value
"string"
  toJSON JsonType
JNumber = Value
"number"
  toJSON JsonType
JInteger = Value
"integer"
  toJSON JsonType
JBoolean = Value
"boolean"
  toJSON JsonType
JNull = Value
"null"
  toJSON (JArray JsonType
_) = Value
"array"
  toJSON (JObject Schema
_) = Value
"object"

-- | A named property with a given type (supports nested values).
newtype Property = Property JsonType
  deriving (Int -> Property -> ShowS
[Property] -> ShowS
Property -> String
(Int -> Property -> ShowS)
-> (Property -> String) -> ([Property] -> ShowS) -> Show Property
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Property -> ShowS
showsPrec :: Int -> Property -> ShowS
$cshow :: Property -> String
show :: Property -> String
$cshowList :: [Property] -> ShowS
showList :: [Property] -> ShowS
Show, Property -> Property -> Bool
(Property -> Property -> Bool)
-> (Property -> Property -> Bool) -> Eq Property
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Property -> Property -> Bool
== :: Property -> Property -> Bool
$c/= :: Property -> Property -> Bool
/= :: Property -> Property -> Bool
Eq, (forall x. Property -> Rep Property x)
-> (forall x. Rep Property x -> Property) -> Generic Property
forall x. Rep Property x -> Property
forall x. Property -> Rep Property x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Property -> Rep Property x
from :: forall x. Property -> Rep Property x
$cto :: forall x. Rep Property x -> Property
to :: forall x. Rep Property x -> Property
Generic)

instance ToJSON Property where
  toJSON :: Property -> Value
toJSON (Property (JArray JsonType
itemType)) =
    [Pair] -> Value
object [Key
"type" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= (Text
"array" :: Text), Key
"items" Key -> Property -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JsonType -> Property
Property JsonType
itemType]
  toJSON (Property (JObject Schema
schema)) = Schema -> Value
forall a. ToJSON a => a -> Value
toJSON Schema
schema
  toJSON (Property JsonType
typ) = [Pair] -> Value
object [Key
"type" Key -> JsonType -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JsonType
typ]

-- | Complete schema representation.
--
-- @since 0.2.0.0
data Schema = Schema
  { Schema -> Map Text Property
schemaProperties :: HM.Map Text Property
  , Schema -> [Text]
schemaRequired :: [Text]
  }
  deriving (Int -> Schema -> ShowS
[Schema] -> ShowS
Schema -> String
(Int -> Schema -> ShowS)
-> (Schema -> String) -> ([Schema] -> ShowS) -> Show Schema
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Schema -> ShowS
showsPrec :: Int -> Schema -> ShowS
$cshow :: Schema -> String
show :: Schema -> String
$cshowList :: [Schema] -> ShowS
showList :: [Schema] -> ShowS
Show, Schema -> Schema -> Bool
(Schema -> Schema -> Bool)
-> (Schema -> Schema -> Bool) -> Eq Schema
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Schema -> Schema -> Bool
== :: Schema -> Schema -> Bool
$c/= :: Schema -> Schema -> Bool
/= :: Schema -> Schema -> Bool
Eq, (forall x. Schema -> Rep Schema x)
-> (forall x. Rep Schema x -> Schema) -> Generic Schema
forall x. Rep Schema x -> Schema
forall x. Schema -> Rep Schema x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Schema -> Rep Schema x
from :: forall x. Schema -> Rep Schema x
$cto :: forall x. Rep Schema x -> Schema
to :: forall x. Rep Schema x -> Schema
Generic)

instance ToJSON Schema where
  toJSON :: Schema -> Value
toJSON (Schema Map Text Property
props [Text]
req) =
    [Pair] -> Value
object
      [ Key
"type" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= (Text
"object" :: Text)
      , Key
"properties" Key -> Map Text Property -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Map Text Property
props
      , Key
"required" Key -> [Text] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= [Text]
req
      ]

-- | Internal builder for schema DSL.
newtype SchemaBuilder = SchemaBuilder Schema
  deriving (Int -> SchemaBuilder -> ShowS
[SchemaBuilder] -> ShowS
SchemaBuilder -> String
(Int -> SchemaBuilder -> ShowS)
-> (SchemaBuilder -> String)
-> ([SchemaBuilder] -> ShowS)
-> Show SchemaBuilder
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SchemaBuilder -> ShowS
showsPrec :: Int -> SchemaBuilder -> ShowS
$cshow :: SchemaBuilder -> String
show :: SchemaBuilder -> String
$cshowList :: [SchemaBuilder] -> ShowS
showList :: [SchemaBuilder] -> ShowS
Show, SchemaBuilder -> SchemaBuilder -> Bool
(SchemaBuilder -> SchemaBuilder -> Bool)
-> (SchemaBuilder -> SchemaBuilder -> Bool) -> Eq SchemaBuilder
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SchemaBuilder -> SchemaBuilder -> Bool
== :: SchemaBuilder -> SchemaBuilder -> Bool
$c/= :: SchemaBuilder -> SchemaBuilder -> Bool
/= :: SchemaBuilder -> SchemaBuilder -> Bool
Eq)

-- | Create an empty schema object.
emptyObject :: SchemaBuilder
emptyObject :: SchemaBuilder
emptyObject = Schema -> SchemaBuilder
SchemaBuilder (Schema -> SchemaBuilder) -> Schema -> SchemaBuilder
forall a b. (a -> b) -> a -> b
$ Map Text Property -> [Text] -> Schema
Schema Map Text Property
forall k a. Map k a
HM.empty []

-- | Add a simple field with a given name and type.
addProperty :: Text -> JsonType -> SchemaBuilder -> SchemaBuilder
addProperty :: Text -> JsonType -> SchemaBuilder -> SchemaBuilder
addProperty Text
name JsonType
typ (SchemaBuilder Schema
s) =
  Schema -> SchemaBuilder
SchemaBuilder (Schema -> SchemaBuilder) -> Schema -> SchemaBuilder
forall a b. (a -> b) -> a -> b
$ Schema
s {schemaProperties = HM.insert name (Property typ) (schemaProperties s)}

-- | Add a nested object field with its own schema.
addObjectProperty :: Text -> Schema -> SchemaBuilder -> SchemaBuilder
addObjectProperty :: Text -> Schema -> SchemaBuilder -> SchemaBuilder
addObjectProperty Text
name Schema
nestedSchema (SchemaBuilder Schema
s) =
  Schema -> SchemaBuilder
SchemaBuilder (Schema -> SchemaBuilder) -> Schema -> SchemaBuilder
forall a b. (a -> b) -> a -> b
$
    Schema
s {schemaProperties = HM.insert name (Property (JObject nestedSchema)) (schemaProperties s)}

-- | Mark a field as required.
requireField :: Text -> SchemaBuilder -> SchemaBuilder
requireField :: Text -> SchemaBuilder -> SchemaBuilder
requireField Text
name (SchemaBuilder Schema
s) =
  Schema -> SchemaBuilder
SchemaBuilder (Schema -> SchemaBuilder) -> Schema -> SchemaBuilder
forall a b. (a -> b) -> a -> b
$ Schema
s {schemaRequired = name : schemaRequired s}

-- | Mark multiple fields as required.
requireFields :: [Text] -> SchemaBuilder -> SchemaBuilder
requireFields :: [Text] -> SchemaBuilder -> SchemaBuilder
requireFields [Text]
names SchemaBuilder
builder = (Text -> SchemaBuilder -> SchemaBuilder)
-> SchemaBuilder -> [Text] -> SchemaBuilder
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Text -> SchemaBuilder -> SchemaBuilder
requireField SchemaBuilder
builder [Text]
names

-- | Finalize the schema from a builder.
buildSchema :: SchemaBuilder -> Schema
buildSchema :: SchemaBuilder -> Schema
buildSchema (SchemaBuilder Schema
s) = Schema
s

-- | Wrap a 'SchemaBuilder' as a nested object type.
objectOf :: SchemaBuilder -> JsonType
objectOf :: SchemaBuilder -> JsonType
objectOf SchemaBuilder
builder = Schema -> JsonType
JObject (SchemaBuilder -> Schema
buildSchema SchemaBuilder
builder)

-- | Create an array of a given JSON type.
arrayOf :: JsonType -> JsonType
arrayOf :: JsonType -> JsonType
arrayOf = JsonType -> JsonType
JArray

-- | Convert schema into a JSON 'Value' suitable for API submission.
toOllamaFormat :: Schema -> Value
toOllamaFormat :: Schema -> Value
toOllamaFormat = Schema -> Value
forall a. ToJSON a => a -> Value
toJSON

-- | Pretty print a schema as formatted JSON.
printSchema :: Schema -> IO ()
printSchema :: Schema -> IO ()
printSchema = String -> IO ()
putStrLn (String -> IO ()) -> (Schema -> String) -> Schema -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack (Text -> String) -> (Schema -> Text) -> Schema -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LazyText -> Text
TL.toStrict (LazyText -> Text) -> (Schema -> LazyText) -> Schema -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> LazyText
T.decodeUtf8 (ByteString -> LazyText)
-> (Schema -> ByteString) -> Schema -> LazyText
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schema -> ByteString
forall a. ToJSON a => a -> ByteString
encode

-- | Infix alias for 'addProperty'.
(|+) :: SchemaBuilder -> (Text, JsonType) -> SchemaBuilder
SchemaBuilder
builder |+ :: SchemaBuilder -> (Text, JsonType) -> SchemaBuilder
|+ (Text
name, JsonType
typ) = Text -> JsonType -> SchemaBuilder -> SchemaBuilder
addProperty Text
name JsonType
typ SchemaBuilder
builder

-- | Infix alias for 'addObjectProperty'.
(|++) :: SchemaBuilder -> (Text, Schema) -> SchemaBuilder
SchemaBuilder
builder |++ :: SchemaBuilder -> (Text, Schema) -> SchemaBuilder
|++ (Text
name, Schema
schema) = Text -> Schema -> SchemaBuilder -> SchemaBuilder
addObjectProperty Text
name Schema
schema SchemaBuilder
builder

-- | Infix alias for 'requireField'.
(|!) :: SchemaBuilder -> Text -> SchemaBuilder
SchemaBuilder
builder |! :: SchemaBuilder -> Text -> SchemaBuilder
|! Text
name = Text -> SchemaBuilder -> SchemaBuilder
requireField Text
name SchemaBuilder
builder

-- | Infix alias for 'requireFields'.
(|!!) :: SchemaBuilder -> [Text] -> SchemaBuilder
SchemaBuilder
builder |!! :: SchemaBuilder -> [Text] -> SchemaBuilder
|!! [Text]
names = [Text] -> SchemaBuilder -> SchemaBuilder
requireFields [Text]
names SchemaBuilder
builder

infixl 7 |+, |++
infixl 6 |!, |!!