// Copyright 2025 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
struct Server {
  tools : Map[String, Tool]
  prompts : Map[String, Prompt]
  in_channel : @channel.T[Message]
  out_channel : @channel.T[Message]
  mut running : Bool
  mut capabilities : Json
  mut clientInfo : Json
  callbacks : Map[Id, @channel.T[Message]]
  random : @random.Rand
  mut level : LogLevel
}

///|
pub fn new(seed : Bytes) -> Server {
  Server::{
    tools: {},
    prompts: {},
    in_channel: @channel.make(size=10),
    out_channel: @channel.make(size=10),
    running: false,
    capabilities: {},
    clientInfo: {},
    callbacks: {},
    random: @random.Rand::new(generator=@random.chacha8(seed~)),
    level: Debug,
  }
}

///|
pub fn serve(
  self : Server,
  stdin~ : &@stream.Reader,
  stdout~ : &@stream.Writer,
  stderr~ : &@stream.Writer
) -> Unit {
  let stdin = @stream.buffered_reader(stdin)
  self.running = true
  // read from stdin for input channel
  @promise.spawn(async fn(_defer) {
    while self.running {
      try {
        let message = stdin.read_line().await()
          |> @json.parse()
          |> @json.from_json()
        self.in_channel.push(message)
      } catch {
        @json.InvalidChar(_)
        | @json.InvalidEof
        | @json.InvalidNumber(_)
        | @json.InvalidIdentEscape(_)
        | @json.JsonDecodeError(_) =>
          self.out_channel.push(
            Response(
              id=None,
              payload=Err({ code: -32700, message: "Parse error", data: None }),
            ),
          )
        e => {
          stderr.write("\{e}\n").await()
          break
        }
      }
    }
  })
  |> ignore
  // write to stdout for output channel
  @promise.spawn(async fn(_defer) {
    while true {
      let message = self.out_channel.pop()
      stdout.write("\{message.to_json() |> Json::stringify}\n").await()
    }
  })
  |> ignore
  // handle messages
  @promise.spawn(async fn(defer_) {
    defer_(fn() { self.running = false })
    let request = self.in_channel.pop()
    guard request is Request(id~, method_="initialize", params~) else {
      fail("Invalid initialization request \{request}")
    }
    guard params
      is {
        "protocolVersion": "2024-11-05",
        "capabilities": capabilities,
        "clientInfo": clientInfo,
        ..
      } else {
      self.out_channel.push(
        @rpc.response(
          id~,
          payload=Err({
            code: -32602,
            message: "Unsupported protocol version",
            data: Some({ "supported": ["2024-11-05"] }),
          }),
        ),
      )
    }
    let response = @rpc.response(
      id~,
      payload=Ok({
        "protocolVersion": "2024-11-05",
        "capabilities": {
          "logging": {},
          "tools": { "listChanged": true },
          "prompts": { "listChanged": true },
        },
        "serverInfo": { "name": "mcp-cli", "version": "0.1.0" },
      }),
    )
    self.out_channel.push(response)
    let notification = self.in_channel.pop()
    guard notification
      is Notification(method_="notifications/initialized", params=_) else {
      fail("Invalid initializaiton notification \{notification}")
    }
    self.capabilities = capabilities
    self.clientInfo = clientInfo
    while true {
      let message = self.in_channel.pop()
      @promise.spawn(async fn(_defer) {
        try self.handle(message) catch {
          error => self.log(Json::string(error.to_string()), level=Error)
        } else {
          Some(response) => self.out_channel.push(response)
          None => ()
        }
      })
      |> ignore
    }
  })
  |> ignore
}

///|
async fn handle(self : Server, message : Message) -> Message? raise {
  // Check the used ids
  match message {
    Request(method_="ping", id~, ..) => Some(@rpc.response(id~, payload=Ok({})))
    Request(method_="logging/setLevel", id~, params~) => {
      guard params is { "level": String(level), .. } else {
        Some(
          @rpc.response(
            id~,
            payload=Err({ code: -32602, message: "Invalid params", data: None }),
          ),
        )
      }
      match level.to_lower() {
        "debug" => self.level = Debug
        "info" => self.level = Info
        "notice" => self.level = Notice
        "warning" => self.level = Warning
        "error" => self.level = Error
        "critical" => self.level = Critical
        "alert" => self.level = Alert
        "emergency" => self.level = Emergency
        _ =>
          return Some(
            @rpc.response(
              id~,
              payload=Err({
                code: -32602,
                message: "Invalid logging level",
                data: None,
              }),
            ),
          )
      }
      Some(@rpc.response(id~, payload=Ok({})))
    }
    Request(method_="tools/list", id~, ..) =>
      Some(
        @rpc.response(
          id~,
          payload=Ok({
            "tools": Json::array(
              self.tools
              .iter()
              .map(tool => {
                let (name, { description, inputSchema, .. }) = tool
                (
                  {
                    "name": Json::string(name),
                    "description": Json::string(description),
                    "inputSchema": inputSchema.to_json(),
                  } : Json)
              })
              .collect(),
            ),
          }),
        ),
      )
    Request(method_="tools/call", id~, params~) =>
      if params is { "name": String(name), "arguments": arguments, .. } &&
        self.tools.get(name) is Some(tool) {
        if not(tool.inputSchema.verify(arguments)) {
          return Some(
            @rpc.response(
              id~,
              payload=Err(RPCError::{
                code: -32602,
                message: "Failed to validate tool call argument",
                data: None,
              }),
            ),
          )
        }
        let { isError, payload } = (tool.cb)(arguments, self)
        return Some(
          @rpc.response(
            id~,
            payload=Ok({
              "isError": Json::boolean(isError),
              "content": Json::array(
                payload.map(data => match data {
                  Text(text~) => { "type": "text", "text": Json::string(text) }
                  Image(data~, mimeType~) =>
                    {
                      "type": "image",
                      "data": Json::string(data),
                      "mimeType": Json::string(mimeType),
                    }
                  Resource(uri~, mimeType~, text~) =>
                    {
                      "type": "resource",
                      "uri": Json::string(uri),
                      "mimeType": Json::string(mimeType),
                      "text": Json::string(text),
                    }
                }),
              ),
            }),
          ),
        )
      } else {
        return Some(
          @rpc.response(
            id~,
            payload=Err(RPCError::{
              code: -32602,
              message: "Failed to verify tool call payload",
              data: None,
            }),
          ),
        )
      }
    Request(method_="prompts/list", id~, ..) =>
      Some(
        @rpc.response(
          id~,
          payload=Ok({
            "prompts": Json::array(
              self.prompts
              .iter()
              .map(prompt => {
                let (name, { description, arguments, .. }) = prompt
                (
                  {
                    "name": Json::string(name),
                    "description": Json::string(description),
                    "arguments": Json::array(
                      arguments.map(argument => {
                        let { name, description, required } = argument
                        {
                          "name": Json::string(name),
                          "description": Json::string(description.or("")),
                          "required": Json::boolean(required is Some(true)),
                        }
                      }),
                    ),
                  } : Json)
              })
              .collect(),
            ),
          }),
        ),
      )
    Request(method_="prompts/get", id~, params~) =>
      if params is { "name": String(name), "arguments": Object(arguments), .. } &&
        self.prompts.get(name) is Some(prompt) {
        for argument in prompt.arguments {
          if argument.required is Some(true) &&
            not(arguments.contains(argument.name)) {
            return Some(
              @rpc.response(
                id~,
                payload=Err(RPCError::{
                  code: -32602,
                  message: "Failed to validate prompt get argument due to missing argument \{argument.name}",
                  data: None,
                }),
              ),
            )
          }
        }
        let response = (prompt.cb)(arguments)
        return Some(
          @rpc.response(
            id~,
            payload=Ok({
              "messages": Json::array(
                response.map(resp => match resp {
                  { role, message: Text(text~) } =>
                    {
                      "role": match role {
                        User => "user"
                        Assistant => "assistant"
                      },
                      "content": { "type": "text", "text": Json::string(text) },
                    }
                  { role, message: Image(data~, mimeType~) } =>
                    {
                      "role": match role {
                        User => "user"
                        Assistant => "assistant"
                      },
                      "content": {
                        "type": "image",
                        "data": Json::string(data),
                        "mimeType": Json::string(mimeType),
                      },
                    }
                  { role, message: Resource(uri~, mimeType~, text~) } =>
                    {
                      "role": match role {
                        User => "user"
                        Assistant => "assistant"
                      },
                      "content": {
                        "type": "resource",
                        "uri": Json::string(uri),
                        "mimeType": Json::string(mimeType),
                        "text": Json::string(text),
                      },
                    }
                }),
              ),
            }),
          ),
        )
      } else {
        return Some(
          @rpc.response(
            id~,
            payload=Err(RPCError::{
              code: -32602,
              message: "Failed to validate prompt get payload",
              data: None,
            }),
          ),
        )
      }
    Request(method_="resources/list", id~, ..) =>
      Some(@rpc.response(id~, payload=Ok({ "resources": [] })))
    Request(method_="resources/templates/list", id~, ..) =>
      Some(@rpc.response(id~, payload=Ok({ "resourceTemplates": [] })))
    Response(id=Some(id), ..) as response => {
      if self.callbacks.get(id) is Some(channel) {
        @promise.spawn(fn(_dfer) { channel.push(response) }) |> ignore
      } else {
        fail("Unexpected response for request \{id}: \{response}")
      }
      None
    }
    Response(id=None, ..) => fail("Unexpected response for request \{message}")
    other => fail("Unhandled \{other}")
  }
}

///|
pub async fn log(
  self : Server,
  message : Json,
  level~ : LogLevel = Info
) -> Unit {
  if level >= self.level {
    let notification = @rpc.Notification(method_="notifications/message", params={
      "level": Json::string(level.to_string().to_lower()),
      "data": message,
    })
    self.out_channel.push(notification)
  }
}

///|
pub fn add_tool(
  self : Server,
  name~ : String,
  description~ : String,
  inputSchema~ : Schema,
  cb~ : async (Json, Server) -> Response raise
) -> Unit {
  self.tools[name] = Tool::{ description, inputSchema, cb }
  if self.running {
    @promise.spawn(async fn(_defer) {
      self.out_channel.push(
        Notification(method_="notifications/tools/list_changed", params={}),
      )
    })
    |> ignore
  }
}

///|
pub fn remove_tool(self : Server, name : String) -> Unit {
  self.tools.remove(name)
  if self.running {
    @promise.spawn(async fn(_defer) {
      self.out_channel.push(
        Notification(method_="notifications/tools/list_changed", params={}),
      )
    })
    |> ignore
  }
}

///|
pub fn add_prompt(
  self : Server,
  name~ : String,
  description~ : String,
  arguments~ : Array[PromptArgument] = [],
  cb~ : (Map[String, Json]) -> Array[PromptMessage]
) -> Unit {
  self.prompts[name] = Prompt::{ description, arguments, cb }
  if self.running {
    @promise.spawn(async fn(_defer) {
      self.out_channel.push(
        Notification(method_="notifications/prompts/list_changed", params={}),
      )
    })
    |> ignore
  }
}

///|
pub fn remove_prompt(self : Server, name : String) -> Unit {
  self.prompts.remove(name)
  if self.running {
    @promise.spawn(async fn(_defer) {
      self.out_channel.push(
        Notification(method_="notifications/prompts/list_changed", params={}),
      )
    })
    |> ignore
  }
}

///|
pub async fn list_roots(self : Server) -> Array[Root] raise {
  if self.capabilities is { "roots": _, .. } && self.running {
    let channel = @channel.make(size=1)
    let lower = self.random.uint64().to_le_bytes()
    let higher = self.random.uint64().to_le_bytes()
    let random = FixedArray::make(16, b'\x00')
    random.blit_from_bytes(0, lower, 0, 8)
    random.blit_from_bytes(8, higher, 0, 8)
    let uuid = @uuid.from_bytes(Bytes::from_fixedarray(random)).as_version(V4)
    let id = @rpc.String(uuid.to_string())
    self.callbacks[id] = channel
    self.out_channel.push(Request(id~, method_="roots/list", params={}))
    let response = channel.pop()
    self.callbacks.remove(id)
    guard response is Response(payload~, ..) else {
      fail("Unexpected message for request roots/list: \{response}")
    }
    guard payload is Ok({ "roots": Array(roots), .. }) else {
      fail("Unexpected payload for request roots/list")
    }
    roots.map(payload => {
      guard payload is { "uri": String(uri), "name"? : name, .. }
      { uri, name: name.bind(Json::as_string) }
    })
  } else {
    []
  }
}