using System.Data;
using System.Text.Json;
using Tango.Portal.Chat.Web.Models;
using Tango.Portal.Chat.Web.Services;
using Kusto.Data.Data;
using Microsoft.AspNetCore.Mvc;

namespace Tango.Portal.Chat.Web.Controllers
{
    [ApiController]
    [Route("api/[controller]")]
    public sealed class ChatController : ControllerBase
    {
        private readonly SchemaRegistry _schema;
        private readonly KqlGuard _guard;
        private readonly KustoQueryService _adx;
        private readonly LlmClient _llm;
        private static readonly string[] AllowTables = new[] { "JobRunsTable", "JobStatusTable", "TelemetryTable", "MachinesTable" };

        public ChatController(SchemaRegistry schema, KqlGuard guard, KustoQueryService adx, LlmClient llm)
        {
            _schema = schema;
            _guard = guard;
            _adx = adx;
            _llm = llm;
        }

        [HttpPost("ask")]
        public async Task<ActionResult<ChatResponse>> Ask([FromBody] ChatRequest req, CancellationToken ct)
        {
            try
            {
                var schemaJson = _schema.GetSchemaJson();
                var plannerPrompt = _schema.GetPlannerPrompt();

                // 1) Ask the model for KQL
                var plan = await _llm.ProposeKqlAsync(plannerPrompt, req.Question, schemaJson, req.History, ct);

                if (plan.Assistant == "data")
                {
                    // 2) Guardrail validation
                    var val = _guard.Validate(plan.Kql);
                    if (!val.IsOk) return BadRequest(new { error = "Invalid KQL", details = val.Error, plan });

                    // 4) Execute in ADX
                    DataTable table;
                    try
                    {
                        table = await _adx.QueryAsync(plan.Kql, plan.Parameters, ct);
                    }
                    catch (Exception ex)
                    {
                        // Return error to the client so they can iterate
                        return new ChatResponse
                        {
                            Answer = $"Seems like my kusto query ran into some issue..\n{ex.Message}",
                            ThreadId = req.ThreadId,
                            UsedKql = plan.Kql
                        };
                    }

                    // 5) Build compact facts (limit rows/cols)
                    var preview = ToPreview(table, 200);
                    var facts = JsonSerializer.Serialize(preview);

                    // 6) Ask model for final answer
                    //var answer = await _llm.AnswerFromFactsAsync(req.Question, facts, plan.Kql, ct);

                    var run = await _llm.AnswerWithAssistantAsync(
                        LlmClient.AssistantType.Data,
                        req.Question,
                        facts,
                        plan.Kql,
                        req.ThreadId,   // <-- reuse if provided
                        ct);

                    return new ChatResponse
                    {
                        Answer = run.Answer,
                        UsedKql = plan.Kql,
                        Preview = preview,
                        ThreadId = run.ThreadId   // <-- echo back the thread id used/created
                    };
                }
                else if (plan.Assistant == "docs")
                {
                    // AFTER
                    var run = await _llm.AnswerWithAssistantAsync(
                        LlmClient.AssistantType.Docs,
                        req.Question,
                        string.Empty,
                        plan.Kql,
                        req.ThreadId,   // <-- reuse if provided
                        ct);

                    return new ChatResponse
                    {
                        Answer = run.Answer,
                        ThreadId = run.ThreadId
                    };
                }
                else
                {
                    return new ChatResponse
                    {
                        Answer = plan.ConversationAnswer,
                        ThreadId = req.ThreadId
                    };
                }
            }
            catch (Exception ex)
            {
                return new ChatResponse
                {
                    Answer = $"Ooops something went wrong...\n{ex.Message}",
                    ThreadId = req.ThreadId
                };
            }
        }

        private static object ToPreview(DataTable dt, int maxRows)
        {
            var cols = dt.Columns.Cast<DataColumn>().Select(c => c.ColumnName).ToArray();
            var rows = new List<Dictionary<string, object?>>();
            int count = 0;
            foreach (DataRow r in dt.Rows)
            {
                if (count++ >= maxRows) break;
                var d = new Dictionary<string, object?>();
                foreach (var c in cols) d[c] = r[c];
                rows.Add(d);
            }
            return new { columns = cols, rows };
        }
    }
}