using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

public class Script : ScriptBase
{
    private const int StartTimeoutSeconds = 100;
    private const int MaxRequestBytes = 131072;
    private const int MaxQuestionCharacters = 32768;
    private const int MaxLineBytes = 262144;
    private const int MaxStreamBytes = 1048576;
    private const string IdPattern = "[A-Za-z0-9_-]{1,128}";

    public override async Task<HttpResponseMessage> ExecuteAsync()
    {
        string operation = ResolveOperation(this.Context.OperationId);
        if (operation == null)
            return Error(HttpStatusCode.BadRequest, "UNKNOWN_OPERATION", "Unknown connector operation.");

        HttpResponseMessage invalid = ValidateRoute(operation);
        if (invalid != null)
            return invalid;

        if (operation == "StartAgentMode")
        {
            using (var deadline = CancellationTokenSource.CreateLinkedTokenSource(this.CancellationToken))
            {
                deadline.CancelAfter(TimeSpan.FromSeconds(StartTimeoutSeconds));
                return await StartAsync(deadline.Token).ConfigureAwait(false);
            }
        }

        try
        {
            return await this.Context.SendAsync(this.Context.Request, this.CancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            return Error(HttpStatusCode.GatewayTimeout, "REQUEST_TIMED_OUT",
                "The request timed out. Consult the conversation before deciding whether to retry or cancel.");
        }
        catch (HttpRequestException)
        {
            return Error(HttpStatusCode.BadGateway, "DATABRICKS_UNREACHABLE", "The Databricks request failed.");
        }
    }

    internal async Task<HttpResponseMessage> StartAsync(CancellationToken token)
    {
        HttpResponseMessage invalid = ValidateRoute("StartAgentMode");
        if (invalid != null)
            return invalid;

        bool sent = false;
        try
        {
            if (this.Context.Request.Content == null)
                return Error(HttpStatusCode.BadRequest, "INVALID_REQUEST", "The request body is required.");

            JObject input;
            try
            {
                input = JObject.Parse(await ReadTextAsync(this.Context.Request.Content, MaxRequestBytes, token).ConfigureAwait(false));
            }
            catch (JsonReaderException)
            {
                return Error(HttpStatusCode.BadRequest, "INVALID_REQUEST", "The request body must be a JSON object.");
            }
            catch (InvalidDataException)
            {
                return Error(HttpStatusCode.BadRequest, "INVALID_REQUEST", "The request body exceeds 128 KiB.");
            }

            string question = ReadString(input["question"]);
            if (string.IsNullOrWhiteSpace(question) || question.Length > MaxQuestionCharacters)
                return Error(HttpStatusCode.BadRequest, "INVALID_QUESTION",
                    "question must contain between 1 and 32768 characters and cannot be whitespace.");
            if (input["enable_viz"] != null && input["enable_viz"].Type != JTokenType.Boolean)
                return Error(HttpStatusCode.BadRequest, "INVALID_REQUEST", "enable_viz must be a boolean.");
            foreach (JProperty property in input.Properties())
                if (property.Name != "question" && property.Name != "enable_viz")
                    return Error(HttpStatusCode.BadRequest, "INVALID_REQUEST",
                        "Only question and enable_viz are accepted. Each start creates a new conversation.");

            var payload = new JObject
            {
                ["input"] = new JArray(new JObject
                {
                    ["type"] = "message",
                    ["role"] = "user",
                    ["content"] = new JArray(new JObject
                    {
                        ["type"] = "input_text",
                        ["text"] = question
                    })
                }),
                ["execution_mode"] = "STANDARD",
                ["enable_viz"] = input["enable_viz"] == null ? false : (bool)input["enable_viz"]
            };

            HttpRequestMessage request = this.Context.Request;
            request.Headers.Accept.Clear();
            request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
            request.Content = new StringContent(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
            sent = true;
            // The platform owns credentials and transport. Early SSE delivery must be validated there.
            using (HttpResponseMessage upstream = await this.Context.SendAsync(request, token).ConfigureAwait(false))
            {
                int status = (int)upstream.StatusCode;
                if (status >= 300 && status < 400)
                    return Error(HttpStatusCode.BadGateway, "UNEXPECTED_REDIRECT",
                        "The managed transport returned a redirect. Verify the configured workspace host.");
                if (!upstream.IsSuccessStatusCode)
                {
                    var failure = Error(upstream.StatusCode, "DATABRICKS_HTTP_ERROR",
                        "Databricks rejected the start request with HTTP " + status + ". Check permissions, quota and preview availability.");
                    if (upstream.Headers.RetryAfter != null)
                        failure.Headers.RetryAfter = upstream.Headers.RetryAfter;
                    return failure;
                }
                string mediaType = upstream.Content == null || upstream.Content.Headers.ContentType == null
                    ? null : upstream.Content.Headers.ContentType.MediaType;
                if (!string.Equals(mediaType, "text/event-stream", StringComparison.OrdinalIgnoreCase))
                    return Error(HttpStatusCode.BadGateway, "UNEXPECTED_CONTENT_TYPE",
                        "The start endpoint did not return text/event-stream. Do not retry automatically.");

                using (Stream stream = await upstream.Content.ReadAsStreamAsync().ConfigureAwait(false))
                {
                    JObject response = await ReadCreatedEventAsync(stream, token).ConfigureAwait(false);
                    return Json(HttpStatusCode.OK, new JObject
                    {
                        ["conversation_id"] = response["conversation_id"],
                        ["response_id"] = response["id"],
                        ["status"] = response["status"]
                    });
                }
            }
        }
        catch (OperationCanceledException)
        {
            return Error(HttpStatusCode.GatewayTimeout, sent ? "START_RESULT_UNKNOWN" : "START_NOT_SENT",
                sent
                    ? "The first event was not obtained before cancellation. The managed transport may have buffered SSE, and the analysis may still be running. Do not retry automatically."
                    : "The operation was cancelled before the Databricks request was sent.");
        }
        catch (HttpRequestException)
        {
            return Error(HttpStatusCode.BadGateway, sent ? "START_RESULT_UNKNOWN" : "DATABRICKS_UNREACHABLE",
                "The connection failed. A submitted analysis may still be running. Do not retry automatically.");
        }
        catch (InvalidDataException ex)
        {
            return Error(HttpStatusCode.BadGateway, "INVALID_START_RESPONSE",
                ex.Message + " Do not retry automatically.");
        }
        catch (IOException)
        {
            return Error(HttpStatusCode.BadGateway, "INVALID_START_RESPONSE",
                "The SSE stream ended, exceeded its size limit, or failed before a valid response.created event. Do not retry automatically.");
        }
        catch (JsonReaderException)
        {
            return Error(HttpStatusCode.BadGateway, "INVALID_START_RESPONSE",
                "An SSE event contained invalid JSON. Do not retry automatically.");
        }
        catch (DecoderFallbackException)
        {
            return Error(HttpStatusCode.BadGateway, "INVALID_START_RESPONSE",
                "The SSE stream was not valid UTF-8. Do not retry automatically.");
        }
    }

    private HttpResponseMessage ValidateRoute(string operation)
    {
        Uri uri = this.Context.Request.RequestUri;
        if (uri == null || !uri.IsAbsoluteUri || uri.Scheme != "https" || uri.Port != 443 ||
            !string.IsNullOrEmpty(uri.UserInfo) || uri.Host.IndexOf("REPLACE", StringComparison.OrdinalIgnoreCase) >= 0 ||
            !(uri.Host.EndsWith(".azuredatabricks.net", StringComparison.OrdinalIgnoreCase) ||
              uri.Host.EndsWith(".cloud.databricks.com", StringComparison.OrdinalIgnoreCase) ||
              uri.Host.EndsWith(".gcp.databricks.com", StringComparison.OrdinalIgnoreCase)))
            return Error(HttpStatusCode.BadRequest, "INVALID_WORKSPACE",
                "Configure an HTTPS Databricks workspace host. Arbitrary hosts, userinfo and nonstandard ports are not allowed.");

        string prefix = "/api/2\\.0/genie/agents/" + IdPattern;
        string pattern;
        string method;
        switch (operation)
        {
            case "StartAgentMode":
                pattern = prefix + "/responses";
                method = "POST";
                break;
            case "GetAgentModeItems":
                pattern = prefix + "/conversations/" + IdPattern + "/items";
                method = "GET";
                break;
            case "CancelAgentModeResponse":
                pattern = prefix + "/conversations/" + IdPattern + "/responses/" + IdPattern + "/cancel";
                method = "POST";
                break;
            default:
                return Error(HttpStatusCode.BadRequest, "UNKNOWN_OPERATION", "Unknown connector operation.");
        }
        if (!Regex.IsMatch(uri.AbsolutePath, "^" + pattern + "$") || this.Context.Request.Method.Method != method)
            return Error(HttpStatusCode.BadRequest, "INVALID_ROUTE", "The operation does not match the HTTP method and Databricks path.");
        return null;
    }

    private static async Task<JObject> ReadCreatedEventAsync(Stream stream, CancellationToken token)
    {
        var buffer = new byte[4096];
        var lineBytes = new List<byte>();
        var data = new StringBuilder();
        var encoding = new UTF8Encoding(false, true);
        string eventName = "";
        int total = 0;
        bool firstLine = true;
        while (true)
        {
            int count = await stream.ReadAsync(buffer, 0, buffer.Length, token).ConfigureAwait(false);
            if (count == 0)
                throw new InvalidDataException("SSE closed before response.created.");
            total += count;
            if (total > MaxStreamBytes)
                throw new InvalidDataException("SSE exceeded the startup size limit.");
            for (int i = 0; i < count; i++)
            {
                if (buffer[i] != 10)
                {
                    lineBytes.Add(buffer[i]);
                    if (lineBytes.Count > MaxLineBytes)
                        throw new InvalidDataException("SSE line exceeded its size limit.");
                    continue;
                }
                string line = encoding.GetString(lineBytes.ToArray()).TrimEnd('\r');
                lineBytes.Clear();
                if (firstLine)
                {
                    line = line.TrimStart('\uFEFF');
                    firstLine = false;
                }
                if (line.Length == 0)
                {
                    if (data.Length > 0)
                    {
                        JObject envelope = JObject.Parse(data.ToString());
                        string type = eventName.Length == 0 ? ReadString(envelope["type"]) : eventName;
                        if (type == "response.failed" || type == "error")
                            throw new InvalidDataException("Databricks emitted a failure event.");
                        if (type == "response.created")
                        {
                            JObject nested = envelope["data"] as JObject;
                            JObject response = envelope["response"] as JObject;
                            if (response == null && nested != null)
                                response = nested["response"] as JObject;
                            if (response == null ||
                                string.IsNullOrWhiteSpace(ReadString(response["id"])) ||
                                string.IsNullOrWhiteSpace(ReadString(response["conversation_id"])) ||
                                ReadString(response["status"]) != "in_progress")
                                throw new InvalidDataException("response.created is missing required fields.");
                            return response;
                        }
                    }
                    data.Clear();
                    eventName = "";
                }
                else if (line.StartsWith("event:", StringComparison.Ordinal))
                    eventName = line.Substring(6).Trim();
                else if (line.StartsWith("data:", StringComparison.Ordinal))
                {
                    string value = line.Substring(5);
                    if (value.StartsWith(" ", StringComparison.Ordinal))
                        value = value.Substring(1);
                    data.Append(value).Append('\n');
                }
            }
        }
    }

    private static async Task<string> ReadTextAsync(HttpContent content, int maximumBytes, CancellationToken token)
    {
        using (Stream stream = await content.ReadAsStreamAsync().ConfigureAwait(false))
        using (var output = new MemoryStream())
        {
            var buffer = new byte[4096];
            int count;
            while ((count = await stream.ReadAsync(buffer, 0, buffer.Length, token).ConfigureAwait(false)) != 0)
            {
                if (output.Length + count > maximumBytes)
                    throw new InvalidDataException("Request body exceeds its size limit.");
                output.Write(buffer, 0, count);
            }
            return new UTF8Encoding(false, true).GetString(output.ToArray());
        }
    }

    private static string ReadString(JToken token)
    {
        return token != null && token.Type == JTokenType.String ? (string)token : null;
    }

    private static bool IsOperation(string operation)
    {
        return operation == "StartAgentMode" || operation == "GetAgentModeItems" || operation == "CancelAgentModeResponse";
    }

    private static string ResolveOperation(string operation)
    {
        if (IsOperation(operation))
            return operation;
        if (string.IsNullOrEmpty(operation))
            return null;
        try
        {
            string decoded = Encoding.UTF8.GetString(Convert.FromBase64String(operation));
            return IsOperation(decoded) ? decoded : null;
        }
        catch (FormatException)
        {
            return null;
        }
    }

    private HttpResponseMessage Error(HttpStatusCode status, string code, string message)
    {
        return Json(status, new JObject
        {
            ["error"] = new JObject
            {
                ["code"] = code,
                ["message"] = message,
                ["correlation_id"] = this.Context.CorrelationId
            }
        });
    }

    private static HttpResponseMessage Json(HttpStatusCode status, JObject body)
    {
        var response = new HttpResponseMessage(status) { Content = CreateJsonContent(body.ToString(Newtonsoft.Json.Formatting.None)) };
        response.Headers.CacheControl = new CacheControlHeaderValue { NoStore = true };
        return response;
    }
}