import csv
import io
import json
import os
import re
import socket
import unittest
from contextlib import redirect_stderr, redirect_stdout
from datetime import datetime

from helpers import EXAMPLE_CSV, EXAMPLE_JSONL, ROOT, TempDirTest

import shadow_ai_report
from shadow_ai_report import FULL_EDITION_URL, FULL_REGISTER_TEXT, cli
from shadow_ai_report.parse import TOOLS_FILE
from shadow_ai_report.report import Analysis, build_summary, render_html


def fmt(n):
    return "{:,}".format(n)


def tool_list():
    with open(TOOLS_FILE, encoding="utf-8-sig", newline="") as f:
        return {r["domain"].lower(): r for r in csv.DictReader(f)}


def naive_host(v):
    v = v.strip().lower().rstrip(".")
    return v[4:] if v.startswith("www.") else v


def naive_match(host, tools):
    for d in tools:
        if host == d or host.endswith("." + d):
            return d
    return None


def expected_from_rows(rows, tools):
    """Independent, deliberately simple re-count used to check the report numbers."""
    total, ai, users_all, ai_users, found, domains, per_tool, blocked = 0, 0, set(), set(), set(), set(), {}, 0
    for host, user, is_blocked in rows:
        total += 1
        domains.add(host)
        if user:
            users_all.add(user)
        d = naive_match(host, tools)
        if d:
            ai += 1
            found.add(d)
            per_tool[d] = per_tool.get(d, 0) + 1
            blocked += bool(is_blocked)
            if user:
                ai_users.add(user)
    high = sum(1 for d in found if tools[d]["risk_level"] == "high" or tools[d]["data_sovereignty"] == "high")
    flagged = sum(1 for d in found if tools[d]["risk_level"] in ("high", "medium")
                  or tools[d]["data_sovereignty"] in ("high", "medium"))
    return {"requests": total, "ai_requests": ai, "ai_tools": len(found), "users_seen": len(users_all),
            "ai_users": len(ai_users), "distinct_domains": len(domains), "per_tool": per_tool,
            "high": high, "flagged": flagged, "ai_blocked": blocked}


def example_csv_expected():
    tools = tool_list()
    rows, lines = [], 0
    with open(EXAMPLE_CSV, encoding="utf-8", newline="") as f:
        for r in csv.DictReader(f):
            lines += 1
            host = naive_host(r.get("query_name") or "")
            if not re.match(r"^[a-z0-9.-]+\.[a-z]+$", host):
                continue
            rows.append((host, r["user"] or r["src_ip"], r["action"] == "blocked"))
    exp = expected_from_rows(rows, tools)
    exp["lines"] = lines
    return exp


def example_jsonl_expected():
    tools = tool_list()
    rows, lines = [], 0
    with open(EXAMPLE_JSONL, encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            lines += 1
            try:
                o = json.loads(line)
            except ValueError:
                continue
            host = naive_host(o.get("QueryName") or "")
            if not host:
                continue
            email = o.get("Email") or ""
            user = email if email and not email.startswith("non_identity@") else o.get("SrcIP")
            rows.append((host, user, str(o.get("ResolverDecision", "")).startswith("blocked")))
    exp = expected_from_rows(rows, tools)
    exp["lines"] = lines
    return exp


def run_cli(*args):
    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        code = cli.main(list(args))
    return code, out.getvalue(), err.getvalue()


class TestExampleReports(TempDirTest):
    def tile(self, html, label):
        m = re.search(r'<div class="label">%s</div><div class="value">([^<]*)</div>' % re.escape(label), html)
        self.assertIsNotNone(m, label)
        return m.group(1)

    def fact(self, html, label):
        m = re.search(r'<div class="label">%s</div><div class="v">([^<]*)</div>' % re.escape(label), html)
        self.assertIsNotNone(m, label)
        return m.group(1)

    def check_html(self, html, exp):
        self.assertEqual(self.tile(html, "AI tools detected"), fmt(exp["ai_tools"]))
        self.assertEqual(self.tile(html, "Users on AI tools"), fmt(exp["ai_users"]))
        self.assertEqual(self.tile(html, "AI-tool requests"), fmt(exp["ai_requests"]))
        self.assertEqual(self.tile(html, "Share of requests"),
                         "{:.1f}%".format(100.0 * exp["ai_requests"] / exp["requests"]))
        self.assertEqual(self.tile(html, "High-risk tools"), fmt(exp["high"]))
        self.assertIn("of %s users or devices in the log" % fmt(exp["users_seen"]), html)
        self.assertIn("%s blocked by your filter" % fmt(exp["ai_blocked"]), html)
        self.assertEqual(self.fact(html, "Log lines read"), fmt(exp["lines"]))
        self.assertEqual(self.fact(html, "Requests analysed"), fmt(exp["requests"]))
        self.assertEqual(self.fact(html, "Lines skipped"), fmt(exp["lines"] - exp["requests"]))
        self.assertEqual(self.fact(html, "Distinct domains seen"), fmt(exp["distinct_domains"]))
        self.assertIn("%d of the AI tools found carry a high or medium rating" % exp["flagged"], html)
        # every detected tool appears in the table with its request count
        tools = tool_list()
        for d, n in exp["per_tool"].items():
            pat = (r'<div class="dom">%s</div></td>.*?data-label="Requests">([\d,]+)</td>' % re.escape(d))
            m = re.search(pat, html, re.S)
            self.assertIsNotNone(m, d)
            self.assertEqual(m.group(1), fmt(n), d)
            self.assertIn(tools[d]["name"].replace("&", "&amp;"), html)

    def test_example_csv(self):
        exp = example_csv_expected()
        # sanity-check the example itself: ~3,000 rows, 7 days, ~25 tools, some flagged
        self.assertEqual(exp["requests"], 3000)
        self.assertEqual(exp["lines"], 3003)
        self.assertEqual(exp["ai_tools"], 25)
        self.assertGreaterEqual(exp["high"], 2)
        for popular in ("chatgpt.com", "perplexity.ai", "character.ai", "midjourney.com", "openai.com"):
            self.assertIn(popular, exp["per_tool"])
        out = self.path("r.html")
        code, stdout, err = run_cli(EXAMPLE_CSV, "-o", out, "--org", "Example Industries Inc.")
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        self.check_html(html, exp)
        self.assertIn("Example Industries Inc.", html)
        self.assertIn("21&ndash;27 Sep 2026", html)
        self.assertIn("(7 days)", html)
        self.assertRegex(stdout, r"AI tools detected\s+%d\n" % exp["ai_tools"])
        self.assertRegex(stdout, r"Lines skipped\s+3\n")
        self.assertIn("Report written: %s" % os.path.abspath(out), stdout)

    def test_example_jsonl(self):
        exp = example_jsonl_expected()
        self.assertEqual(exp["requests"], 1000)
        out = self.path("r.html")
        code, _, err = run_cli(EXAMPLE_JSONL, "-o", out)
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        self.check_html(html, exp)
        self.assertIn("Your organisation", html)
        self.assertIn("cloudflare-dns", html)

    def test_both_files_together(self):
        e1, e2 = example_csv_expected(), example_jsonl_expected()
        out, js = self.path("r.html"), self.path("r.json")
        code, _, err = run_cli(EXAMPLE_CSV, EXAMPLE_JSONL, "-o", out, "--json", js)
        self.assertEqual(code, 0, err)
        with open(js, encoding="utf-8") as f:
            s = json.load(f)
        self.assertEqual(s["totals"]["requests"], e1["requests"] + e2["requests"])
        self.assertEqual(s["totals"]["ai_requests"], e1["ai_requests"] + e2["ai_requests"])
        self.assertEqual(s["totals"]["lines_read"], e1["lines"] + e2["lines"])
        self.assertEqual(len(s["inputs"]), 2)


class TestOutputs(TempDirTest):
    def test_json_and_csv(self):
        exp = example_csv_expected()
        out, js, cs = self.path("r.html"), self.path("r.json"), self.path("tools.csv")
        code, stdout, err = run_cli(EXAMPLE_CSV, "-o", out, "--json", js, "--csv", cs)
        self.assertEqual(code, 0, err)
        with open(js, encoding="utf-8") as f:
            s = json.load(f)
        t = s["totals"]
        self.assertEqual(t["requests"], exp["requests"])
        self.assertEqual(t["lines_read"], exp["lines"])
        self.assertEqual(t["lines_skipped"], 3)
        self.assertEqual(t["distinct_domains"], exp["distinct_domains"])
        self.assertEqual(t["ai_requests"], exp["ai_requests"])
        self.assertEqual(t["ai_tools"], exp["ai_tools"])
        self.assertEqual(t["ai_users"], exp["ai_users"])
        self.assertEqual(t["users_seen"], exp["users_seen"])
        self.assertEqual(t["ai_blocked"], exp["ai_blocked"])
        self.assertEqual(t["high_risk_or_exposure_tools"], exp["high"])
        self.assertEqual(t["flagged_tools"], exp["flagged"])
        self.assertTrue(s["period"]["start"].startswith("2026-09-21T"))
        self.assertTrue(s["period"]["end"].startswith("2026-09-27T"))
        self.assertEqual(s["period"]["days"], 7)
        self.assertEqual({r["domain"]: r["requests"] for r in s["tools"]}, exp["per_tool"])
        self.assertEqual(s["full_edition"]["url"], FULL_EDITION_URL)
        self.assertEqual(s["full_edition"]["register"], FULL_REGISTER_TEXT)
        self.assertEqual(s["tool_list"]["tools"], 300)
        self.assertIn("JSON written", stdout)
        with open(cs, encoding="utf-8", newline="") as f:
            rows = list(csv.DictReader(f))
        self.assertEqual(len(rows), exp["ai_tools"])
        self.assertEqual(list(rows[0].keys())[:4], ["tool", "domain", "category", "subcategory"])
        reqs = [int(r["requests"]) for r in rows]
        self.assertEqual(reqs, sorted(reqs, reverse=True))
        self.assertEqual({r["domain"]: int(r["requests"]) for r in rows}, exp["per_tool"])

    def test_no_user_or_time_columns(self):
        log = self.write("list.txt", "dify.ai\napi.dify.ai\nexample.org\n")
        js, out = self.path("r.json"), self.path("r.html")
        code, _, err = run_cli(log, "-o", out, "--json", js)
        self.assertEqual(code, 0, err)
        with open(js, encoding="utf-8") as f:
            s = json.load(f)
        self.assertIsNone(s["totals"]["users_seen"])
        self.assertIsNone(s["totals"]["ai_blocked"])
        self.assertEqual(s["period"]["start"], None)
        self.assertEqual(s["users"], [])
        with open(out, encoding="utf-8") as f:
            html = f.read()
        self.assertIn("no user column in the log", html)
        self.assertIn("Not stated in the log", html)
        self.assertNotIn('id="users"', html)


class TestAnonymise(TempDirTest):
    def test_pseudonyms(self):
        out, js = self.path("r.html"), self.path("r.json")
        code, _, err = run_cli(EXAMPLE_CSV, "-o", out, "--json", js, "--anonymise")
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        with open(js, encoding="utf-8") as f:
            text = f.read()
        for blob in (html, text):
            self.assertNotIn("@example.com", blob)
            self.assertNotIn("192.0.2.", blob)
        s = json.loads(text)
        names = [u["user"] for u in s["users"]]
        self.assertEqual(names, ["User %d" % (i + 1) for i in range(len(names))])
        counts = [u["ai_requests"] for u in s["users"]]
        self.assertEqual(counts, sorted(counts, reverse=True))
        self.assertIn("User 1", html)
        # stable: a second run gives the same mapping
        js2 = self.path("r2.json")
        run_cli(EXAMPLE_CSV, "-o", self.path("r2.html"), "--json", js2, "--anonymise")
        with open(js2, encoding="utf-8") as f:
            self.assertEqual(json.load(f)["users"], s["users"])

    def test_same_user_same_pseudonym(self):
        a = Analysis()
        for dom, user in (("dify.ai", "zed@x.org"), ("dust.tt", "amy@x.org"), ("api.dify.ai", "zed@x.org"),
                          ("dify.ai", "bob@x.org")):
            a.add(dom, user)
        s = build_summary(a, anonymise=True)
        self.assertEqual(s["users"][0], {"user": "User 1", "ai_requests": 2, "tools": 1,
                                         "top_tools": ["Dify"], "flagged_tools": 0})
        self.assertEqual(sorted(u["user"] for u in s["users"]), ["User 1", "User 2", "User 3"])


class TestHtmlSafety(TempDirTest):
    def test_self_contained(self):
        out = self.path("r.html")
        code, _, err = run_cli(EXAMPLE_CSV, "-o", out)
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        low = html.lower()
        for bad in ("<script", "<link", "<img", "<iframe", "@import", "url(", " src=", "@font-face"):
            self.assertNotIn(bad, low)
        links = set(re.findall(r'href="([^"]+)"', html))
        self.assertEqual({l for l in links if not l.startswith("#")},
                         {FULL_EDITION_URL, "https://creativecommons.org/licenses/by/4.0/"})
        self.assertIn(FULL_REGISTER_TEXT, html)
        self.assertIn("Free edition:</b> recognises the 300 AI tools of the free list", html)
        self.assertIn("recognises only the 300 AI tools of the free list", html)
        self.assertIn("does not mean low or zero AI use", html)
        for stale in ("200 AI tools", "200-tool", "200 tools", "public sample list", "leaves out"):
            self.assertNotIn(stale, html)
        self.assertIn("@media print", html)

    def test_escaping(self):
        log = self.write("x.csv", 'user,domain\n"<script>alert(1)</script>",dify.ai\n')
        out = self.path("r.html")
        code, _, err = run_cli(log, "-o", out, "--org", "<b>Acme & Co</b>")
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        self.assertNotIn("<script>", html)
        self.assertIn("&lt;script&gt;alert(1)&lt;/script&gt;", html)
        self.assertIn("&lt;b&gt;Acme &amp; Co&lt;/b&gt;", html)

    def test_zero_ai_notice(self):
        log = self.write("x.csv", "domain\nexample.org\nexample.com\n")
        out = self.path("r.html")
        code, stdout, err = run_cli(log, "-o", out)
        self.assertEqual(code, 0, err)
        with open(out, encoding="utf-8") as f:
            html = f.read()
        self.assertIn("No AI tools from the free list were found", html)
        self.assertIn("does not mean nobody uses AI", stdout)


class TestCli(TempDirTest):
    def test_bad_input_exit_codes(self):
        code, _, err = run_cli(self.path("missing.csv"), "-o", self.path("r.html"))
        self.assertEqual(code, 2)
        self.assertIn("file not found", err)
        code, _, err = run_cli(self.write("empty.csv", ""), "-o", self.path("r.html"))
        self.assertEqual(code, 2)
        self.assertIn("empty", err)
        code, _, err = run_cli(self.write("x.jsonl", '{"a": 1}\n'), "-o", self.path("r.html"))
        self.assertEqual(code, 2)
        self.assertIn("QueryName", err)
        code, _, err = run_cli(self.write("x.csv", "domain\ndify.ai\n"), "--format", "cloudflare-dns",
                               "-o", self.path("r.html"))
        self.assertEqual(code, 2)
        self.assertIn("no usable log lines", err)
        code, _, err = run_cli(EXAMPLE_CSV, "-o", os.path.join(self.tmp, "no", "such", "dir", "r.html"))
        self.assertEqual(code, 2)
        self.assertIn("cannot write", err)
        self.assertFalse(os.path.exists(self.path("r.html")))

    def test_usage_errors(self):
        with self.assertRaises(SystemExit) as cm:
            run_cli()
        self.assertEqual(cm.exception.code, 2)
        with self.assertRaises(SystemExit) as cm:
            run_cli(EXAMPLE_CSV, "--format", "xml")
        self.assertEqual(cm.exception.code, 2)

    def test_version(self):
        with self.assertRaises(SystemExit) as cm:
            run_cli("--version")
        self.assertEqual(cm.exception.code, 0)


class TestNoNetwork(TempDirTest):
    def test_no_network_modules_imported(self):
        pkg = os.path.dirname(shadow_ai_report.__file__)
        for name in os.listdir(pkg):
            if name.endswith(".py"):
                with open(os.path.join(pkg, name), encoding="utf-8") as f:
                    src = f.read()
                for mod in ("socket", "urllib.request", "http.client", "requests", "ssl", "ftplib", "smtplib"):
                    self.assertIsNone(re.search(r"^\s*(import|from)\s+%s\b" % re.escape(mod), src, re.M),
                                      "%s imports %s" % (name, mod))

    def test_runs_with_sockets_disabled(self):
        real = socket.socket

        def refuse(*a, **k):
            raise AssertionError("network access attempted")

        socket.socket = refuse
        try:
            code, _, err = run_cli(EXAMPLE_CSV, EXAMPLE_JSONL, "-o", self.path("r.html"),
                                   "--json", self.path("r.json"), "--csv", self.path("r.csv"))
        finally:
            socket.socket = real
        self.assertEqual(code, 0, err)


if __name__ == "__main__":
    unittest.main()