using System.Collections.Concurrent;
using System.ComponentModel;
using System.Diagnostics;
using System.Text;
using System.Text.RegularExpressions;
using ModelContextProtocol.Server;

namespace Alchegos.MCP.Tools;
[McpServerToolType, Description("""
                                Manages multiple persistent terminal sessions. 
                                **Sessions are identified and operated on using their unique `sessionId`. The `label` is for human reference only.** 
                                Enables creation, closure, command execution, and output retrieval.
                                """)]
public static class TerminalTool
{
    private static readonly ConcurrentDictionary<string, TerminalSession> _sessions = new();
    public record SessionInfo(
        [property: Description("Unique identifier (`sessionId`) for the session. **This ID is required for all operations targeting this specific session (e.g., sending input, exiting).**")] string SessionId,
        [property: Description("Human-readable label for the session. **This is for display and organizational purposes only and CANNOT be used to interact with the session via subsequent tool calls.**")] string Label
    );
    public record ListSessionsResult(
        [property: Description("List of all active sessions, including their `sessionId` and `label`. **Use the `sessionId` from this list for interacting with a specific session in other tool calls.**")] List<SessionInfo> Sessions
    );
    public record CreateSessionResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status,
        [property: Description("Details of the created session if successful; otherwise, with empty fields.")] SessionInfo Session
    );
    public record ExitSessionResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );
    public record SendSignalResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );
    public record SendInputWithExceptResult(
        [property: Description("Indicates how the response was triggered: 'except' when a regex matched, or 'timeout' when the timeout was reached.")] string OutputBy,
        [property: Description("The regex pattern that matched, or null if none matched before timeout.")] string? MatchedExcept,
        [property: Description("Accumulated standard output from the session after sending input.")] string Stdout,
        [property: Description("Accumulated standard error from the session after sending input.")] string Stderr
    );
    public record SendInputAndForgetResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );
    public record OneOffCommandResult(
        [property: Description("How the command completed: 'finish' when the process exited normally, or 'timeout' when the timeout was reached. 'error' otherwise")] string ReturnBy,
        [property: Description("Standard output produced by the command.")] string Stdout,
        [property: Description("Standard error produced by the command.")] string Stderr
    );
    public record GetCachedResultResult(
        [property: Description("Cached standard output of the session.")] string Stdout,
        [property: Description("Cached standard error of the session.")] string Stderr
    );
    public record ClearResult(
        [property: Description("Status of the operation: 'success' or 'fail'.")] string Status
    );
    private class TerminalSession : IDisposable
    {
        public string Id { get; }
        public string Label { get; }
        public Process Process { get; }
        
        private readonly StringBuilder _stdoutBuffer = new();
        private readonly StringBuilder _stderrBuffer = new();
        private readonly Lock _lockObject = new();
        private bool _isDisposed;

        public TerminalSession(string id, string label)
        {
            Id = id;
            Label = label;
            Process = new Process
            {
                StartInfo = new ProcessStartInfo
                {
                    FileName = GetShellPath(),
                    Arguments = GetShellArguments(),
                    RedirectStandardInput = true,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true
                },
                EnableRaisingEvents = true
            };
            
            Process.OutputDataReceived += (sender, e) =>
            {
                if (e.Data != null)
                    lock (_lockObject)
                    {
                        _stdoutBuffer.AppendLine(e.Data);
                    }
            };
            
            Process.ErrorDataReceived += (sender, e) =>
            {
                if (e.Data != null)
                    lock (_lockObject)
                    {
                        _stderrBuffer.AppendLine(e.Data);
                    }
            };
            
            Process.Start();
            Process.BeginOutputReadLine();
            Process.BeginErrorReadLine();
        }

        private static string GetShellPath()
        {
            if (OperatingSystem.IsWindows())
                return "cmd.exe";
            return "/bin/bash";
        }

        private static string GetShellArguments()
        {
            if (OperatingSystem.IsWindows())
                return "/q";
            return "-i";
        }

        public void SendInput(string input)
        {
            if (!Process.HasExited)
            {
                Process.StandardInput.WriteLine(input);
                Process.StandardInput.Flush();
            }
        }

        public bool SendSignal(int signal)
        {
            try
            {
                if (OperatingSystem.IsWindows())
                {
                    if (signal == 2)
                        return NativeMethods.SendCtrlC(Process);
                    return false;
                }
                return NativeMethods.SendSignal(Process.Id, signal);
            }
            catch
            {
                return false;
            }
        }

        public (string stdout, string stderr) GetAndClearBuffers()
        {
            lock (_lockObject)
            {
                string stdout = _stdoutBuffer.ToString();
                string stderr = _stderrBuffer.ToString();
                _stdoutBuffer.Clear();
                _stderrBuffer.Clear();
                return (stdout, stderr);
            }
        }

        public (string stdout, string stderr) GetBuffers()
        {
            lock (_lockObject)
            {
                return (_stdoutBuffer.ToString(), _stderrBuffer.ToString());
            }
        }

        public void ClearBuffers()
        {
            lock (_lockObject)
            {
                _stdoutBuffer.Clear();
                _stderrBuffer.Clear();
            }
        }

        public void Dispose()
        {
            if (_isDisposed)
                return;
            _isDisposed = true;
            try
            {
                if (!Process.HasExited)
                {
                    Process.StandardInput.WriteLine("exit");
                    Process.StandardInput.Flush();
                    if (!Process.WaitForExit(1000))
                        Process.Kill();
                }
            }
            catch
            {
                // Best effort only
            }
            finally
            {
                Process.Dispose();
            }
        }
    }

    private static class NativeMethods
    {
        public static bool SendCtrlC(Process process)
        {
            try
            {
                if (process.HasExited)
                    return false;
                if (!AttachConsole((uint)process.Id))
                    return false;
                SetConsoleCtrlHandler(null, true);
                try
                {
                    if (!GenerateConsoleCtrlEvent(CtrlTypes.CTRL_C_EVENT, 0))
                        return false;
                    return true;
                }
                finally
                {
                    SetConsoleCtrlHandler(null, false);
                    FreeConsole();
                }
            }
            catch
            {
                return false;
            }
        }

        public static bool SendSignal(int pid, int signal)
        {
            try
            {
                if (OperatingSystem.IsLinux() || OperatingSystem.IsMacOS())
                {
                    var killProcess = Process.Start(new ProcessStartInfo
                    {
                        FileName = "kill",
                        Arguments = $"-{signal} {pid}",
                        UseShellExecute = false,
                        CreateNoWindow = true
                    });
                    killProcess?.WaitForExit();
                    return killProcess?.ExitCode == 0;
                }
                return false;
            }
            catch
            {
                return false;
            }
        }

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool AttachConsole(uint dwProcessId);

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool FreeConsole();

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool SetConsoleCtrlHandler(ConsoleCtrlHandler? handlerRoutine, bool add);

        [System.Runtime.InteropServices.DllImport("kernel32.dll")]
        public static extern bool GenerateConsoleCtrlEvent(CtrlTypes dwCtrlEvent, uint dwProcessGroupId);

        public enum CtrlTypes
        {
            CTRL_C_EVENT = 0,
            CTRL_BREAK_EVENT = 1,
            CTRL_CLOSE_EVENT = 2,
            CTRL_LOGOFF_EVENT = 5,
            CTRL_SHUTDOWN_EVENT = 6
        }
        public delegate bool ConsoleCtrlHandler(CtrlTypes ctrlType);
    }

    [McpServerTool, Description(
        """
        Lists all currently active persistent terminal sessions being managed.
        Provides the essential `sessionId` (required for all other operations on a session) and the human-readable `label` for each.
        Use this to find the `sessionId` of an existing session you want to interact with.
        Return Format: { sessions: [ { session_id: string, label: string }, ... ] }
        """)]
    public static ListSessionsResult ListTerminalSessions()
    {
        List<SessionInfo> sessions = _sessions.Values.Select(s => new SessionInfo(s.Id, s.Label)).ToList();
        return new ListSessionsResult(sessions);
    }

    [McpServerTool, Description(
        """
        Creates a new persistent terminal session (a background shell like bash or cmd) and assigns it a unique `sessionId`.
        This `sessionId` is **essential** and **must be used** in subsequent calls to other TerminalTool functions (like SendInput..., OneOffCommand, GetCachedResult, ExitSession) to target this specific session.
        The `label` is purely for human reference (e.g., in ListTerminalSessions).
        Return Format: { status: 'success' | 'fail', session: { session_id: string, label: string } }
        """)]
    public static CreateSessionResult CreateTerminalSession(
        [Description("Memorable label for the new session.")] string label
    )
    {
        try
        {
            string sessionId = Guid.NewGuid().ToString("N");
            var session = new TerminalSession(sessionId, label);
            if (_sessions.TryAdd(sessionId, session))
                return new CreateSessionResult("success", new SessionInfo(sessionId, label));
            session.Dispose();
            return new CreateSessionResult("fail", new SessionInfo(string.Empty, string.Empty));
        }
        catch (Exception)
        {
            return new CreateSessionResult("fail", new SessionInfo(string.Empty, string.Empty));
        }
    }
    [McpServerTool, Description(
        """
        Terminates and cleans up a specific persistent terminal session identified by its `sessionId`.
        This effectively closes the background shell process associated with that session.
        Use this when you are finished working with a specific session. Fails if the `sessionId` is invalid.
        Return Format: { status: 'success' | 'fail' }
        """)]
    public static ExitSessionResult ExitTerminalSession(
        [Description("ID of the session to exit.")] string sessionId
    )
    {
        if (_sessions.TryRemove(sessionId, out var session))
        {
            session.Dispose();
            return new ExitSessionResult("success");
        }
        return new ExitSessionResult("fail");
    }
    [McpServerTool, Description(
        """
        Terminates and cleans up a specific persistent terminal session identified by its `sessionId`.
        This effectively closes the background shell process associated with that session.
        Use this when you are finished working with a specific session. Fails if the `sessionId` is invalid.
        Return Format: { status: 'success' | 'fail' }
        """)]
    public static SendSignalResult SendSignalToTerminalSession(
        [Description("Session ID to send the signal to.")] string sessionId,
        [Description("Signal number to send (e.g., 2 for SIGINT)." )] int signal
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            bool success = session.SendSignal(signal);
            return new SendSignalResult(success ? "success" : "fail");
        }
        return new SendSignalResult("fail");
    }

    [McpServerTool, Description(
        """
        Sends input text to a specific terminal session (identified by `sessionId`) and WAITS for output.
        It returns ONLY when specific regex patterns (`excepts`) are matched in the session's stdout/stderr OR the `timeout` is reached.
        **Use Case:** Handling interactive prompts within the shell (e.g., password prompts, 'yes/no' confirmations) where you need to wait for a specific response pattern before proceeding.
        Clears the session's output buffer upon returning. Compare with `SendInputAndForget` and `OneOffCommand` for different interaction patterns.
        Return Format: { output_by: 'except' | 'timeout', matched_except: string | null, stdout: string, stderr: string }
        """)]
    public static SendInputWithExceptResult SendCommandToTerminalSessionWithExcept(
        [Description("Session ID to send input to.")] string sessionId,
        [Description("Input text to send (without newline)." )] string input,
        [Description("Timeout in milliseconds.")] int timeout,
        [Description("Regex patterns to match against output.")] params string[] excepts
    )
    {
        if (!_sessions.TryGetValue(sessionId, out var session))
            return new SendInputWithExceptResult("timeout", null, string.Empty, string.Empty);
        session.SendInput(input);
        var regexPatterns = excepts.Select(pattern => new Regex(pattern, RegexOptions.Compiled)).ToArray();
        string matchedPattern = null;
        bool matched = false;
        var watch = Stopwatch.StartNew();
        while (watch.ElapsedMilliseconds < timeout && !matched)
        {
            var (stdout, stderr) = session.GetBuffers();
            string combinedOutput = stdout + stderr;
            for (int i = 0; i < regexPatterns.Length; i++)
            {
                if (regexPatterns[i].IsMatch(combinedOutput))
                {
                    matchedPattern = excepts[i];
                    matched = true;
                    break;
                }
            }
            if (!matched)
                Thread.Sleep(10);
        }
        watch.Stop();
        var (finalStdout, finalStderr) = session.GetAndClearBuffers();
        return new SendInputWithExceptResult(
            matched ? "except" : "timeout",
            matchedPattern,
            finalStdout,
            finalStderr
        );
    }

    [McpServerTool, Description(
        """
        Sends a POSIX-style signal to the process running within the specified terminal session (identified by `sessionId`).
        Common use is sending signal 2 (SIGINT) to simulate Ctrl+C, interrupting a running command within that specific session's shell.
        Requires the `sessionId` of an active session. Fails if the `sessionId` is invalid or the signal cannot be delivered.
        Return Format: { status: 'success' | 'fail' }
        """)]
    public static SendInputAndForgetResult SendCommandToTerminalSessionAndForget(
        [Description("Session ID to send input to.")] string sessionId,
        [Description("Input text to send (without newline)." )] string command
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            try
            {
                session.SendInput(command);
                return new SendInputAndForgetResult("success");
            }
            catch
            {
                return new SendInputAndForgetResult("fail");
            }
        }
        
        return new SendInputAndForgetResult("fail");
    }

    [McpServerTool, Description(
        """
        Executes a single command line within a specific terminal session (identified by `sessionId`), waits for it to finish, and returns the collected stdout/stderr.
        This tool manages sending the command and detecting its completion automatically.
        **Use Case:** Running relatively quick, synchronous commands where you need the output right away (e.g., `ls -l`, `pwd`, `git status`, `echo $VAR`).
        Compare with `SendInputAndForget` (for background tasks) and `SendInputWithExcept` (for interactive prompts).
        Return Format: { return_by: 'finish' | 'timeout' | 'error', stdout: string, stderr: string }
        """)]
    public static OneOffCommandResult SendOneOffCommandToTerminalSession(
        [Description("Session ID to execute the command in.")] string sessionId,
        [Description("Command to execute, including arguments.")] string command,
        [Description("Timeout in milliseconds.")] int timeout
    )
    {
        if (!_sessions.TryGetValue(sessionId, out var session))
            return new OneOffCommandResult("error", string.Empty, string.Empty);
        session.ClearBuffers();
        string marker = $"CMD_COMPLETE_{Guid.NewGuid():N}";
        if (OperatingSystem.IsWindows())
            session.SendInput($"{command} & echo {marker}");
        else
            session.SendInput($"{command}; echo {marker}");
        bool completed = false;
        var watch = Stopwatch.StartNew();
        while (watch.ElapsedMilliseconds < timeout && !completed)
        {
            var (stdout, _) = session.GetBuffers();
            if (stdout.Contains(marker))
                completed = true;
            else
                Thread.Sleep(10);
        }
        watch.Stop();
        var (finalStdout, finalStderr) = session.GetAndClearBuffers();
        if (completed)
        {
            finalStdout = Regex.Replace(finalStdout, $"{marker}\r?\n?$", "");
            finalStderr = Regex.Replace(finalStderr, $"{marker}\r?\n?$", "");
        }
        return new OneOffCommandResult(
            completed ? "finish" : "timeout",
            finalStdout,
            finalStderr
        );
    }

    [McpServerTool, Description(
        """
        Retrieves AND clears all accumulated standard output (stdout) and standard error (stderr) that have been buffered within the specified terminal session (identified by `sessionId`) since the last read or clear operation.
        **Use Case:** Primarily used to collect the output from commands previously started with `SendInputAndForget`. Call this periodically or when you expect a background task might have produced output.
        Return Format: { stdout: string, stderr: string }
        """)]
    public static GetCachedResultResult GetTerminalSessionCachedResult(
        [Description("Session ID to retrieve output from.")] string sessionId
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            var (stdout, stderr) = session.GetAndClearBuffers();
            return new GetCachedResultResult(stdout, stderr);
        }
        return new GetCachedResultResult(string.Empty, string.Empty);
    }

    [McpServerTool, Description(
        """
        Discards (clears) any accumulated stdout and stderr buffered within the specified terminal session (identified by `sessionId`) without returning the content.
        **Use Case:** Useful if you want to ignore any previous output before running a new command or checking for fresh output, particularly if using `SendInputAndForget`.
        Return Format: { status: 'success' | 'fail' }
        """)]
    public static ClearResult Clear(
        [Description("Session ID to clear cache for.")] string sessionId
    )
    {
        if (_sessions.TryGetValue(sessionId, out var session))
        {
            session.ClearBuffers();
            return new ClearResult("success");
        }
        return new ClearResult("fail");
    }
    
}