using System;
using System.Collections.Generic;
using System.Net.Http;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;
using MemScanEDR.Models;

namespace MemScanEDR.Services;

public class AIAnalyzer
{
    private readonly HttpClient _http = new();
    private AIConfig _config;

    public AIAnalyzer(AIConfig config)
    {
        _config = config;
        _http.Timeout = TimeSpan.FromSeconds(config.Timeout);
    }

    public void UpdateConfig(AIConfig config)
    {
        _config = config;
        _http.Timeout = TimeSpan.FromSeconds(config.Timeout);
    }

    public async Task<ScanResult> AnalyzeProcessAsync(ProcessInfo proc, List<SignatureMatch> matches, List<MemoryRegion> suspiciousRegions)
    {
        if (_config.Provider == AIProvider.Heuristic)
            return HeuristicAnalyze(proc, matches, suspiciousRegions);

        // Pre-filter: skip API call if nothing suspicious at all
        var preScore = QuickPreScore(proc, matches, suspiciousRegions);
        if (preScore < 0.05)
        {
            return new ScanResult
            {
                Pid          = proc.Pid,
                ProcessName  = proc.Name,
                RiskLevel    = RiskLevel.Safe,
                RiskScore    = 0.0,
                Indicators   = new List<string>(),
                Reasoning    = $"PID {proc.Pid} ({proc.Name}) | no suspicious indicators, skipped AI analysis",
                AIProvider   = "heuristic/pre-filter",
                ScanTime     = DateTime.Now
            };
        }

        try
        {
            var prompt = BuildPrompt(proc, matches, suspiciousRegions);
            var response = await CallAIAsync(prompt);
            return ParseResponse(proc, response);
        }
        catch (Exception ex)
        {
            var result = HeuristicAnalyze(proc, matches, suspiciousRegions);
            result.Reasoning += $" [AI API error: {ex.Message}]";
            return result;
        }
    }

    private static double QuickPreScore(ProcessInfo proc, List<SignatureMatch> matches, List<MemoryRegion> suspiciousRegions)
    {
        double s = 0;
        if (!proc.IsSigned && proc.IsTempPath) s += 0.30;
        var sysNames = new HashSet<string>(StringComparer.OrdinalIgnoreCase) { "svchost.exe", "csrss.exe", "lsass.exe", "winlogon.exe", "services.exe" };
        if (sysNames.Contains(proc.Name) && !proc.IsSystemPath) s += 0.25;
        var suspPorts = new HashSet<int> { 4444, 5555, 6666, 7777, 1337, 31337, 9999, 12345, 27015 };
        if (proc.Connections.Any(c => { var parts = c.RemoteAddr.Split(':'); return parts.Length >= 2 && int.TryParse(parts[^1], out var p) && suspPorts.Contains(p); })) s += 0.20;
        if (matches.Count > 0) s += 0.20;
        if (suspiciousRegions.Any(r => r.Protect.Contains("EXECUTE_READWRITE"))) s += 0.15;
        return Math.Min(s, 1.0);
    }

    private async Task<string> CallAIAsync(string userMessage)
    {
        if (_config.Provider == AIProvider.Claude)
            return await CallClaudeAsync(userMessage);

        // OpenAI-compatible / Local use the same OpenAI chat format
        var url = _config.Provider == AIProvider.Local
            ? $"{_config.BaseUrl}/api/chat"
            : $"{_config.BaseUrl.TrimEnd('/')}/chat/completions";

        object requestBody = _config.Provider == AIProvider.Local
            ? new { model = _config.Model, messages = new[] { new { role = "system", content = SystemPrompt }, new { role = "user", content = userMessage } }, stream = false }
            : new { model = _config.Model, messages = new[] { new { role = "system", content = SystemPrompt }, new { role = "user", content = userMessage } }, max_tokens = _config.MaxTokens, temperature = _config.Temperature };

        var json = JsonSerializer.Serialize(requestBody);
        var request = new HttpRequestMessage(HttpMethod.Post, url);
        request.Content = new StringContent(json, Encoding.UTF8, "application/json");

        if (!string.IsNullOrEmpty(_config.ApiKey))
            request.Headers.Add("Authorization", $"Bearer {_config.ApiKey}");

        var response = await _http.SendAsync(request);
        response.EnsureSuccessStatusCode();
        var responseJson = await response.Content.ReadAsStringAsync();

        using var doc = JsonDocument.Parse(responseJson);

        if (_config.Provider == AIProvider.Local)
            return doc.RootElement.GetProperty("message").GetProperty("content").GetString() ?? "";

        return doc.RootElement.GetProperty("choices")[0].GetProperty("message").GetProperty("content").GetString() ?? "";
    }

    private async Task<string> CallClaudeAsync(string userMessage)
    {
        var url = $"{_config.BaseUrl.TrimEnd('/')}/v1/messages";
        var requestBody = new
        {
            model = _config.Model,
            max_tokens = _config.MaxTokens,
            system = SystemPrompt,
            messages = new[] { new { role = "user", content = userMessage } }
        };

        var json = JsonSerializer.Serialize(requestBody);
        var request = new HttpRequestMessage(HttpMethod.Post, url);
        request.Content = new StringContent(json, Encoding.UTF8, "application/json");
        request.Headers.Add("x-api-key", _config.ApiKey);
        request.Headers.Add("anthropic-version", "2023-06-01");

        var response = await _http.SendAsync(request);
        response.EnsureSuccessStatusCode();
        var responseJson = await response.Content.ReadAsStringAsync();

        using var doc = JsonDocument.Parse(responseJson);
        return doc.RootElement.GetProperty("content")[0].GetProperty("text").GetString() ?? "";
    }

    private ScanResult ParseResponse(ProcessInfo proc, string content)
    {
        content = content.Trim();
        if (content.StartsWith("```"))
        {
            var lines = content.Split('\n');
            content = string.Join('\n', lines[1..^1]);
        }

        try
        {
            using var doc = JsonDocument.Parse(content);
            var root = doc.RootElement;

            var riskLevel = root.GetProperty("risk_level").GetString() switch
            {
                "high" => RiskLevel.High,
                "suspicious" => RiskLevel.Suspicious,
                _ => RiskLevel.Safe
            };

            var indicators = new List<string>();
            if (root.TryGetProperty("indicators", out var indArr))
                foreach (var ind in indArr.EnumerateArray())
                    indicators.Add(ind.GetString() ?? "");

            return new ScanResult
            {
                Pid = proc.Pid,
                ProcessName = proc.Name,
                RiskLevel = riskLevel,
                RiskScore = root.GetProperty("risk_score").GetDouble(),
                Indicators = indicators,
                Reasoning = root.GetProperty("reasoning").GetString() ?? "",
                AIProvider = $"{_config.Provider}/{_config.Model}",
                ScanTime = DateTime.Now
            };
        }
        catch
        {
            return new ScanResult
            {
                Pid = proc.Pid,
                ProcessName = proc.Name,
                RiskLevel = RiskLevel.Safe,
                RiskScore = 0,
                Reasoning = content,
                AIProvider = $"{_config.Provider}/{_config.Model}",
                ScanTime = DateTime.Now
            };
        }
    }

    private string BuildPrompt(ProcessInfo proc, List<SignatureMatch> matches, List<MemoryRegion> suspiciousRegions)
    {
        var sb = new StringBuilder();
        sb.AppendLine("=== PROCESS DATA ===");
        sb.AppendLine($"PID: {proc.Pid}");
        sb.AppendLine($"Name: {proc.Name}");
        sb.AppendLine($"Path: {(string.IsNullOrEmpty(proc.ExePath) ? "unknown" : proc.ExePath)}");
        sb.AppendLine($"Parent: {proc.ParentName} (PPID: {proc.Ppid})");
        sb.AppendLine($"CommandLine: {proc.CmdLine}");
        sb.AppendLine($"Signed: {proc.IsSigned}");
        sb.AppendLine($"Signer: {(string.IsNullOrEmpty(proc.Signer) ? "-" : proc.Signer)}");
        sb.AppendLine($"SystemPath: {proc.IsSystemPath}");
        sb.AppendLine($"TempPath: {proc.IsTempPath}");
        sb.AppendLine($"ThreadCount: {proc.ThreadCount}");
        sb.AppendLine($"WorkingSetMB: {proc.WorkingSetMB}");

        if (proc.Connections.Count > 0)
        {
            sb.AppendLine($"\nNetworkConnections ({proc.Connections.Count}):");
            foreach (var c in proc.Connections.Take(8))
                sb.AppendLine($"  {c.Protocol} {c.LocalAddr} -> {c.RemoteAddr} [{c.State}]");
        }
        else
        {
            sb.AppendLine("\nNetworkConnections: none");
        }

        if (matches.Count > 0)
        {
            sb.AppendLine($"\nMemorySignatures ({matches.Count}):");
            foreach (var m in matches.Take(10))
                sb.AppendLine($"  [{m.Severity}] [{m.Category}] {m.SigId} @ {m.Address}: {m.Description}");
        }
        else
        {
            sb.AppendLine("\nMemorySignatures: none");
        }

        if (suspiciousRegions.Count > 0)
        {
            sb.AppendLine($"\nSuspiciousMemoryRegions ({suspiciousRegions.Count}):");
            foreach (var r in suspiciousRegions.Take(5))
                sb.AppendLine($"  0x{r.BaseAddress:X} size={r.Size} protect={r.Protect} type={r.Type} state={r.State}");
        }
        else
        {
            sb.AppendLine("\nSuspiciousMemoryRegions: none");
        }

        sb.AppendLine("\n=== ANALYZE THIS PROCESS ===");
        return sb.ToString();
    }

    public ScanResult HeuristicAnalyze(ProcessInfo proc, List<SignatureMatch> matches, List<MemoryRegion> suspiciousRegions)
    {
        double score = 0;
        var indicators = new List<string>();
        var reasons = new List<string>();

        // Rule 1: unsigned + temp path
        if (!proc.IsSigned && proc.IsTempPath)
        {
            score += 0.30;
            indicators.Add("unsigned_temp_path");
            reasons.Add("unsigned executable in temp directory");
        }

        // Rule 2: process masquerading
        var sysNames = new HashSet<string>(StringComparer.OrdinalIgnoreCase) {
            "svchost.exe", "csrss.exe", "lsass.exe", "winlogon.exe", "services.exe"
        };
        if (sysNames.Contains(proc.Name) && !proc.IsSystemPath)
        {
            score += 0.25;
            indicators.Add("process_masquerading");
            reasons.Add($"system process name in non-system path");
        }

        // Rule 3: suspicious ports
        var suspPorts = new HashSet<int> { 4444, 5555, 6666, 7777, 1337, 31337 };
        foreach (var conn in proc.Connections)
        {
            var parts = conn.RemoteAddr.Split(':');
            if (parts.Length == 2 && int.TryParse(parts[1], out var port) && suspPorts.Contains(port))
            {
                score += 0.20;
                indicators.Add("suspicious_port");
                reasons.Add($"suspicious remote port {port}");
                break;
            }
        }

        // Rule 4: signature matches
        if (matches.Count > 0)
        {
            var high = matches.Count(m => m.Severity == "high");
            var med = matches.Count(m => m.Severity == "medium");
            score += Math.Min(high * 0.20, 0.4);
            score += Math.Min(med * 0.08, 0.2);
            if (high > 0) { indicators.Add("high_severity_match"); reasons.Add($"{high} high-severity signatures"); }
            if (matches.Any(m => m.Category == "shellcode")) { score += 0.15; indicators.Add("shellcode_detected"); }
            if (matches.Any(m => m.Category == "c2")) { score += 0.20; indicators.Add("c2_pattern"); }
        }

        // Rule 5: suspicious memory regions
        var rwxCount = suspiciousRegions.Count(r => r.Protect.Contains("EXECUTE_READWRITE") && r.Type == "PRIVATE");
        if (rwxCount > 0) { score += Math.Min(rwxCount * 0.12, 0.25); indicators.Add("rwx_private_memory"); }

        score = Math.Min(score, 1.0);
        var level = score >= 0.75 ? RiskLevel.High : score >= 0.50 ? RiskLevel.Suspicious : RiskLevel.Safe;

        return new ScanResult
        {
            Pid = proc.Pid,
            ProcessName = proc.Name,
            RiskLevel = level,
            RiskScore = score,
            Indicators = indicators,
            Reasoning = $"PID {proc.Pid} ({proc.Name}) | {string.Join("; ", reasons)} | score={score:F2}",
            AIProvider = "heuristic/rule-engine-v1",
            ScanTime = DateTime.Now
        };
    }

    public async Task<(bool ok, string message)> TestConnectionAsync()
    {
        if (_config.Provider == AIProvider.Heuristic) return (true, "Heuristic mode (no API)");
        try
        {
            await CallAIAsync("ping");
            return (true, $"OK - {_config.Provider}/{_config.Model}");
        }
        catch (Exception ex) { return (false, ex.Message); }
    }

    private const string SystemPrompt = @"You are a Windows memory security expert analyzing a single process for threats.

=== YOUR TASK ===
Given process metadata, memory scan results, and memory region data, determine if the process is safe, suspicious, or high risk.

=== RISK LEVEL DEFINITIONS ===

""safe"" — The process exhibits NORMAL or benign behavior:
- Signed binaries from trusted publishers (Microsoft, Adobe, major vendors)
- DLL injection by standard antivirus/EDR components
- Internet connections to normal CDNs/cloud services (Azure, AWS, Google)
- Small private RWX regions used by .NET JIT, JavaScript engines (node.js), or Java
- Hooking by regular security products
- Processes in System32 that are known Windows components
**IMPORTANT: The vast majority (>95%) of processes are SAFE. Default to ""safe"" unless you have clear threat evidence.**

""suspicious"" — ONE or TWO of the following are present:
- Unsigned executable in temporary folder (%TEMP%, Downloads)
- Process with system-critical name (svchost.exe, lsass.exe, csrss.exe) running from non-system path
- Known C2 ports (4444, 5555, 6666, 1337, etc.)
- Multiple large PAGE_EXECUTE_READWRITE private regions
- Signature matches of medium severity
- Unsigned DLL loaded from suspicious path
**A process needs at least 2-3 weak signals to be suspicious, or 1 strong signal.**

""high"" — CLEAR evidence of malware, requiring MULTIPLE concurring signals:
- Shellcode signature match + RWX private memory + unsigned temp path (all 3 together)
- C2 pattern match + network to suspicious remote address
- Reflective DLL injection signature + PE header in private memory
- Process hollowing or process injection signature matches
**Single indicators are NOT enough for ""high"". Requires 3+ concurring indicators.**

=== JSON RESPONSE FORMAT ===
Respond with ONLY this JSON (no markdown, no explanation):
{""risk_level"":""safe"",""risk_score"":0.0,""indicators"":[],""reasoning"":""""}

Fields:
- risk_level: ""safe"" | ""suspicious"" | ""high""
- risk_score: 0.0–1.0 (0.0–0.30=safe, 0.31–0.69=suspicious, 0.70–1.0=high)
- indicators: list of 0–5 key indicator strings (e.g. ""unsigned_temp_binary"", ""rwx_shellcode_region"", ""c2_network"", ""process_masquerading"", ""dll_injection_sig"")
- reasoning: 1–2 sentence summary

=== IMPORTANT RULES ===
1. Error on the side of ""safe"". Only elevate when you see CONCRETE threat indicators.
2. A memory signature match alone does NOT make a process malicious — many legitimate tools use similar techniques.
3. .NET/Java/Node.js processes legitimately use RWX memory — this alone is NOT suspicious.
4. Consider the process name and path context — calc.exe in System32 is safe; calc.exe in %TEMP% is not.
5. Process injection indicators combined with unsigned temp binaries is dangerous.
6. Only list indicators that are actually present — never fabricate evidence.";
}