#include "../parser.hpp"
#include <common/enumerate.hpp>
#include <format>
#include "codegen.hpp"
#include <iomanip>
#include <sstream>
#include <ranges>

namespace {

inline std::string_view getTypename(const TypeValue &type) {
  if (auto ptr = std::get_if<TypeStruct *>(&type.data)) { return (*ptr)->name; }
  if (auto ptr = std::get_if<EnumValue>(&type.data)) { return ptr->name; }
  if (auto ptype = std::get_if<PrimitiveType>(&type.data)) {
    switch (*ptype) {
    case PrimitiveType::Void:
      return "void";
    case PrimitiveType::Boolean:
      return "boolean";
    case PrimitiveType::Int:
    case PrimitiveType::UInt:
    case PrimitiveType::Double:
      return "number";
    case PrimitiveType::String:
      return "string";
    case PrimitiveType::Any:
      return "any";
    }
  }

  return "string";
}

inline std::string getTypeSignature(const TypeValue &type) {
  std::string str{getTypename(type)};
  if (type.isArray) str.append("[]");
  return str;
}

inline std::string generateType(const TypeStruct &s) {
  std::ostringstream oss;
  oss << "export type " << s.name << " = {\n";

  for (const auto &field : s.fields) {
    oss << "\t" << field.name;

    if (field.type.isOptional) oss << "?";

    oss << ": " << getTypeSignature(field.type) << ";\n";
  }

  oss << "}";

  return oss.str();
}

inline std::string generateEnum(const EnumValue &e) {
  std::ostringstream oss;
  oss << "export type " << e.name << " = ";

  for (const auto &[idx, value] : e.values | vicinae::enumerate) {
    if (idx > 0) oss << " | ";
    oss << "'" << value << "'";
  }

  oss << ";";

  return oss.str();
}

inline std::string generateEventHandlerTypeSignature(const Event &event) {
  std::ostringstream oss;
  oss << "(";
  for (const auto &[idx, param] : event.params | vicinae::enumerate) {
    if (idx > 0) oss << ", ";
    oss << param.name << ": " << getTypeSignature(param.type);
  }
  oss << ") => void";
  return oss.str();
}

inline std::string serializeParamList(std::span<const MethodParameter> params) {
  std::ostringstream oss;

  for (const auto &[idx, param] : params | vicinae::enumerate) {
    if (idx > 0) oss << ", ";
    oss << param.name << ": " << getTypename(param.type);
  }

  return oss.str();
}

inline std::string generateService(const Service &s) {
  std::ostringstream oss;
  oss << "class " << s.name << "Service {\n";

  oss << "\tconstructor(private readonly transport: RpcTransport) {}\n\n";

  for (const auto &method : s.methods) {
    oss << "\t" << method.name << "(";

    for (const auto &[idx, param] : method.params | vicinae::enumerate) {
      if (idx > 0) oss << ", ";
      oss << param.name;
      if (param.type.isOptional) oss << "?";
      oss << ": " << getTypeSignature(param.type);
    }

    oss << "): " << "Promise<" << getTypeSignature(method.returnType) << "> {\n";

    std::string methodName = std::format("{}/{}", s.name, method.name);

    oss << "\t\treturn this.transport.request(" << std::quoted(methodName) << ", { ";

    for (const auto &[idx, param] : method.params | vicinae::enumerate) {
      if (idx > 0) oss << ", ";
      oss << param.name;
    }

    oss << "});";
    oss << "\t" << "\n\t}\n\n";
  }

  if (!s.events.empty()) {
    for (const auto &event : s.events) {
      std::string eventName = std::format("{}/{}", s.name, event.name);

      oss << "\t" << event.name << "(handler: " << generateEventHandlerTypeSignature(event)
          << "): EventSubscription {\n";
      oss << "\t\treturn this.transport.subscribe(" << std::quoted(eventName) << ", (msg) => handler(";

      for (const auto &[idx, param] : event.params | vicinae::enumerate) {
        if (idx > 0) oss << ", ";
        oss << "msg." << param.name;
      }

      oss << "))";
      oss << "\n\t}\n";
    }
  }

  oss << "}";

  return oss.str();
}

}; // namespace

static constexpr const auto commonCode = R"(
// generated by vicinae figura codegen - do not edit manually

interface JsonRpcMessage {
	jsonrpc: "2.0";
	method?: string;
	id?: number;
	params?: Record<string, any>;
	result?: any;
	error?: any;
}

interface ITransport {
	send(data: string): void;
}
)";

static constexpr const auto serverBoilerplate = R"(
export class RpcTransport {
	constructor(private readonly transport: ITransport) { }

	emit(method: string, params: Record<string, any>): void {
		this.sendMessage({ jsonrpc: '2.0', method, params });
	}

	reply(id: number, result: any) {
		this.sendMessage({ jsonrpc: '2.0', id, result });
	}

	private sendMessage(msg: JsonRpcMessage) {
		this.transport.send(JSON.stringify(msg));
	}
};

)";

static constexpr const auto busCode = R"(
type EventSubscription = {
	unsubscribe: () => void;
};

export class RpcTransport {
	constructor(private readonly transport: ITransport) { }

	dispatchMessage(data: string) {
		const msg = JSON.parse(data) as JsonRpcMessage;

		if (msg.id !== undefined) {
			const handler = this.requestMap.get(msg.id);

			if (handler) {
				if (msg.error) handler.reject(msg.error);
				else handler.resolve(msg.result);
				this.requestMap.delete(msg.id);
			}
		}
		if (msg.id === undefined && msg.method) {
			for (const cb of this.handlers.get(msg.method) ?? []) {
				cb(msg.params);
			}
		}
	}

	request<T>(method: string, params: Record<string, any>): Promise<T> {
		const id = this.id++;
		const promise = new Promise<T>((resolve, reject) => {
			this.requestMap.set(id, { resolve: (msg) => resolve(msg as T), reject });
		});

		this.sendMessage({ jsonrpc: '2.0', id, method, params });

		return promise;
	}

	subscribe(method: string, cb: (result: any) => void): EventSubscription {
		const handlers = this.handlers.get(method);

		if (handlers) handlers.push(cb);
		else this.handlers.set(method, [cb]);

		return {
			unsubscribe: () => {
				const handlers = this.handlers.get(method);

				if (handlers) {
					handlers.splice(handlers.indexOf(cb), 1);
				}
			},
		};
	}

	private sendMessage(msg: JsonRpcMessage) {
		this.transport.send(JSON.stringify(msg));
	}


	private id = 1;
	private requestMap = new Map<number, { resolve: (value: any) => void, reject: (error: any) => void }>;
	private handlers = new Map<string, Array<(result: any) => void>>;
};
)";

static std::string tab(int n) {
  std::string str;
  str.reserve(n);
  for (int i = 0; i != n; ++i)
    str += '\t';
  return str;
}

class TypeScriptCodeGenerator : public AbstractCodeGenerator {
  std::string name() const override { return "typescript"; }

  static bool isVoid(const TypeValue &value) {
    auto const visitor = overloads{[](PrimitiveType type) { return type == PrimitiveType::Void; },
                                   [](auto &&other) { return false; }};

    return std::visit(visitor, value.data);
  }

  static void generateTypes(std::ostringstream &oss, const Tree &tree) {
    for (auto const &s : tree.enums) {
      oss << generateEnum(s) << "\n\n";
    }

    for (auto const &s : tree.structs) {
      oss << generateType(*s) << "\n\n";
    }
  }

  std::string generateClient(const Tree &tree, const CodegenOptions &opts) override {
    std::ostringstream oss;

    oss << commonCode << busCode << "\n";

    generateTypes(oss, tree);

    for (auto const &s : tree.services) {
      oss << generateService(*s) << "\n\n";
    }

    oss << "export class Client {\n";
    oss << "\tconstructor(private readonly transport: RpcTransport) {\n";

    for (auto const &s : tree.services) {
      oss << "\t\tthis." << s->name << " = new " << s->name << "Service(this.transport);\n";
    }

    oss << "\t}\n";

    oss << R"(
  	route(msg: string): void { this.transport.dispatchMessage(msg); }
  )";

    for (auto const &s : tree.services) {
      oss << "\t" << s->name << ": " << s->name << "Service;\n";
    }

    oss << "\n}\n";

    return oss.str();
  }

  std::string serializeCases(const Tree &ast) {
    std::ostringstream oss;

    for (const auto &s : ast.services) {
      for (const auto &m : s->methods) {
        std::string methodId = std::format("{}/{}", s->name, m.name);
        oss << "case " << std::quoted(methodId) << ":";
        oss << "\t" << "return this." << s->name << "." << m.name << "(";

        for (const auto &[idx, p] : m.params | vicinae::enumerate) {
          if (idx > 0) oss << ", ";
          oss << "msg.params." << p.name;
        }

        oss << ")";

        if (isVoid(m.returnType)) {
          oss << ".then(() => this.rpc.reply(id, null))\n";
        } else {
          oss << ".then((r) => this.rpc.reply(id, r))\n";
        }
      }
    }

    oss << "default: throw new Error(`No handler for method ${msg.method}`)";

    return oss.str();
  }

  static std::string eventEmitterName(const Service &service) {
    return std::format("{}EventEmitter", service.name);
  }

  std::string generateServer(const Tree &ast, const CodegenOptions &opts) override {
    std::ostringstream oss;

    oss << commonCode << serverBoilerplate;

    generateTypes(oss, ast);

    // generate service interfaces, to be implemented by the user
    for (const auto &service : ast.services) {
      if (!service->methods.empty()) {
        oss << "export abstract class " << service->name << "Service {\n";

        oss << "\tconstructor(private readonly rpc: RpcTransport) {}\n";

        for (const auto &m : service->methods) {
          oss << "\tabstract " << m.name << "(" << serializeParamList(m.params);
          oss << "): " << "Promise <" << getTypename(m.returnType) << ">;\n";
        }

        for (const auto &ev : service->events) {
          std::string eventName = std::format("{}/{}", service->name, ev.name);
          oss << "emit_" << ev.name << "(" << serializeParamList(ev.params) << ") {\n";
          oss << "\tthis.rpc.emit(" << std::quoted(eventName) << ", { ";
          for (const auto &[idx, param] : ev.params | vicinae::enumerate) {
            if (idx > 0) oss << ", ";
            oss << param.name;
          }
          oss << "});";
          oss << "\n}";
        }

        oss << "\n}\n";
      }
    }

    oss << "export class Server {\n";
    oss << "\tconstructor(private readonly rpc: RpcTransport";

    for (const auto &[idx, service] : ast.services | vicinae::enumerate) {
      if (!service->methods.empty()) {
        oss << ", readonly " << service->name << ": " << service->name << "Service";
      }
    }

    oss << ") {\n";

    for (const auto &service : ast.services) {}

    oss << "}\n\n";

    oss << R"(
	route(raw: string) {
		const msg = JSON.parse(raw) as JsonRpcMessage;

		// request
		if (msg.id !== undefined && msg.method && msg.params) {
			const id = msg.id;

			switch (msg.method) {
	)";

    oss << serializeCases(ast);
    oss << "\n}";
    oss << "\n}";
    oss << "\n}";
    oss << "\n}";

    return oss.str();
  }
};
