"""Tests for Finding Annotator."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pytest

from mcpaudit.annotator import (
    Annotation,
    AnnotationStore,
    DEFAULT_ANNOTATIONS_FILE,
    EXCLUDE_STATUSES,
    VALID_STATUSES,
    annotate_finding,
    bulk_annotate,
    filter_annotated,
    filter_by_status,
    format_annotations_json,
    format_annotations_text,
    load_and_annotate,
    load_annotations,
    make_finding_id,
    make_finding_id_from_dict,
    remove_annotation,
    save_annotations,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _sample_findings() -> list[dict[str, Any]]:
    return [
        {"rule_id": "CMD-001", "file": "server.py", "line": 10,
         "severity": "HIGH", "title": "Shell injection"},
        {"rule_id": "AUTH-001", "file": "server.py", "line": 20,
         "severity": "MEDIUM", "title": "Missing auth"},
        {"rule_id": "RATE-001", "file": "api.py", "line": 5,
         "severity": "LOW", "title": "No rate limit"},
    ]


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

class TestConstants:
    def test_valid_statuses(self):
        expected = {
            "false_positive", "accepted_risk", "todo",
            "in_progress", "wont_fix",
        }
        assert VALID_STATUSES == expected

    def test_exclude_statuses(self):
        assert "false_positive" in EXCLUDE_STATUSES
        assert "accepted_risk" in EXCLUDE_STATUSES
        assert "wont_fix" in EXCLUDE_STATUSES
        assert "todo" not in EXCLUDE_STATUSES
        assert "in_progress" not in EXCLUDE_STATUSES


# ---------------------------------------------------------------------------
# make_finding_id
# ---------------------------------------------------------------------------

class TestMakeFindingId:
    def test_basic(self):
        fid = make_finding_id("CMD-001", "server.py", 42)
        assert fid == "CMD-001:server.py:42"

    def test_with_path(self):
        fid = make_finding_id("CMD-001", "/path/to/server.py", 10)
        assert fid == "CMD-001:/path/to/server.py:10"

    def test_from_dict(self):
        finding = {
            "rule_id": "CMD-001",
            "file": "server.py",
            "line": 42,
        }
        fid = make_finding_id_from_dict(finding)
        assert fid == "CMD-001:server.py:42"

    def test_from_dict_missing_keys(self):
        fid = make_finding_id_from_dict({})
        assert fid == "::0"

    def test_from_dict_partial(self):
        fid = make_finding_id_from_dict({"rule_id": "CMD-001"})
        assert fid == "CMD-001::0"


# ---------------------------------------------------------------------------
# Annotation
# ---------------------------------------------------------------------------

class TestAnnotation:
    def test_create(self):
        ann = Annotation(
            finding_id="CMD-001:server.py:42",
            status="false_positive",
            note="Checked manually",
            author="carlos",
        )
        assert ann.finding_id == "CMD-001:server.py:42"
        assert ann.status == "false_positive"
        assert ann.note == "Checked manually"
        assert ann.created_at != ""

    def test_validate_valid(self):
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="false_positive",
        )
        assert ann.validate() == []

    def test_validate_empty_finding_id(self):
        ann = Annotation(finding_id="", status="false_positive")
        errors = ann.validate()
        assert any("finding_id" in e for e in errors)

    def test_validate_empty_status(self):
        ann = Annotation(finding_id="CMD-001:a.py:1", status="")
        errors = ann.validate()
        assert any("status" in e for e in errors)

    def test_validate_invalid_status(self):
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="invalid_status",
        )
        errors = ann.validate()
        assert any("Invalid status" in e for e in errors)

    def test_all_valid_statuses(self):
        for status in VALID_STATUSES:
            ann = Annotation(
                finding_id="CMD-001:a.py:1",
                status=status,
            )
            assert ann.validate() == []

    def test_to_dict(self):
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
            note="Fix later",
            author="test",
        )
        d = ann.to_dict()
        assert d["finding_id"] == "CMD-001:a.py:1"
        assert d["status"] == "todo"
        assert d["note"] == "Fix later"
        assert d["author"] == "test"
        assert "created_at" in d
        assert "updated_at" in d

    def test_from_dict(self):
        data = {
            "finding_id": "CMD-001:a.py:1",
            "status": "false_positive",
            "note": "OK",
            "author": "test",
            "created_at": "2024-01-01T00:00:00Z",
            "updated_at": "2024-01-01T00:00:00Z",
        }
        ann = Annotation.from_dict(data)
        assert ann.finding_id == "CMD-001:a.py:1"
        assert ann.status == "false_positive"
        assert ann.created_at == "2024-01-01T00:00:00Z"

    def test_from_dict_missing_keys(self):
        ann = Annotation.from_dict({})
        assert ann.finding_id == ""
        assert ann.status == ""

    def test_auto_timestamps(self):
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
        )
        assert ann.created_at != ""
        assert ann.updated_at != ""
        assert ann.created_at == ann.updated_at

    def test_preserve_timestamps(self):
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
            created_at="2024-01-01T00:00:00Z",
            updated_at="2024-01-02T00:00:00Z",
        )
        assert ann.created_at == "2024-01-01T00:00:00Z"
        assert ann.updated_at == "2024-01-02T00:00:00Z"


# ---------------------------------------------------------------------------
# AnnotationStore
# ---------------------------------------------------------------------------

class TestAnnotationStore:
    def test_empty_store(self):
        store = AnnotationStore()
        assert store.count == 0
        assert store.list_all() == []

    def test_add_annotation(self):
        store = AnnotationStore()
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
        )
        store.annotations[ann.finding_id] = ann
        assert store.count == 1

    def test_get(self):
        store = AnnotationStore()
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
        )
        store.annotations[ann.finding_id] = ann
        got = store.get("CMD-001:a.py:1")
        assert got is not None
        assert got.status == "todo"

    def test_get_missing(self):
        store = AnnotationStore()
        assert store.get("missing") is None

    def test_has(self):
        store = AnnotationStore()
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
        )
        store.annotations[ann.finding_id] = ann
        assert store.has("CMD-001:a.py:1") is True
        assert store.has("missing") is False

    def test_is_excluded_false_positive(self):
        store = AnnotationStore()
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="false_positive",
        )
        store.annotations[ann.finding_id] = ann
        assert store.is_excluded("CMD-001:a.py:1") is True

    def test_is_excluded_todo(self):
        store = AnnotationStore()
        ann = Annotation(
            finding_id="CMD-001:a.py:1",
            status="todo",
        )
        store.annotations[ann.finding_id] = ann
        assert store.is_excluded("CMD-001:a.py:1") is False

    def test_is_excluded_missing(self):
        store = AnnotationStore()
        assert store.is_excluded("missing") is False

    def test_filter_by_status(self):
        store = AnnotationStore()
        store.annotations["a"] = Annotation(
            finding_id="a", status="todo",
        )
        store.annotations["b"] = Annotation(
            finding_id="b", status="false_positive",
        )
        store.annotations["c"] = Annotation(
            finding_id="c", status="todo",
        )
        todos = store.filter_by_status("todo")
        assert len(todos) == 2

    def test_list_all_sorted(self):
        store = AnnotationStore()
        store.annotations["z"] = Annotation(
            finding_id="z", status="todo",
        )
        store.annotations["a"] = Annotation(
            finding_id="a", status="todo",
        )
        result = store.list_all()
        assert result[0].finding_id == "a"
        assert result[1].finding_id == "z"

    def test_to_dict(self):
        store = AnnotationStore()
        store.annotations["a"] = Annotation(
            finding_id="a", status="todo",
        )
        d = store.to_dict()
        assert "a" in d
        assert d["a"]["status"] == "todo"

    def test_from_dict(self):
        data = {
            "CMD-001:a.py:1": {
                "finding_id": "CMD-001:a.py:1",
                "status": "false_positive",
                "note": "OK",
            },
        }
        store = AnnotationStore.from_dict(data)
        assert store.count == 1
        assert store.get("CMD-001:a.py:1") is not None

    def test_from_dict_empty(self):
        store = AnnotationStore.from_dict({})
        assert store.count == 0

    def test_from_dict_invalid_entries(self):
        data = {
            "a": "not a dict",
            "b": {"status": "todo"},
        }
        store = AnnotationStore.from_dict(data)
        # "a" is invalid, "b" has empty finding_id set from key
        assert store.count >= 1

    def test_stats(self):
        store = AnnotationStore()
        store.annotations["a"] = Annotation(
            finding_id="a", status="todo",
        )
        store.annotations["b"] = Annotation(
            finding_id="b", status="todo",
        )
        store.annotations["c"] = Annotation(
            finding_id="c", status="false_positive",
        )
        stats = store.stats()
        assert stats["todo"] == 2
        assert stats["false_positive"] == 1


# ---------------------------------------------------------------------------
# annotate_finding
# ---------------------------------------------------------------------------

class TestAnnotateFinding:
    def test_basic(self):
        ann, errors = annotate_finding(
            "CMD-001:a.py:1", "false_positive",
        )
        assert errors == []
        assert ann.status == "false_positive"

    def test_with_note_and_author(self):
        ann, errors = annotate_finding(
            "CMD-001:a.py:1", "accepted_risk",
            note="Checked", author="carlos",
        )
        assert errors == []
        assert ann.note == "Checked"
        assert ann.author == "carlos"

    def test_invalid_status(self):
        ann, errors = annotate_finding(
            "CMD-001:a.py:1", "invalid",
        )
        assert len(errors) > 0

    def test_empty_finding_id(self):
        ann, errors = annotate_finding("", "todo")
        assert len(errors) > 0

    def test_updates_store(self):
        store = AnnotationStore()
        ann, errors = annotate_finding(
            "CMD-001:a.py:1", "todo",
            store=store,
        )
        assert errors == []
        assert store.count == 1
        assert store.get("CMD-001:a.py:1") is not None

    def test_updates_existing(self):
        store = AnnotationStore()
        annotate_finding(
            "CMD-001:a.py:1", "todo",
            note="First", store=store,
        )
        annotate_finding(
            "CMD-001:a.py:1", "false_positive",
            note="Updated", store=store,
        )
        assert store.count == 1
        ann = store.get("CMD-001:a.py:1")
        assert ann.status == "false_positive"
        assert ann.note == "Updated"

    def test_preserves_created_at_on_update(self):
        store = AnnotationStore()
        ann1, _ = annotate_finding(
            "CMD-001:a.py:1", "todo",
            store=store,
        )
        original_created = ann1.created_at
        ann2, _ = annotate_finding(
            "CMD-001:a.py:1", "false_positive",
            store=store,
        )
        assert ann2.created_at == original_created


# ---------------------------------------------------------------------------
# bulk_annotate
# ---------------------------------------------------------------------------

class TestBulkAnnotate:
    def test_basic(self):
        store = AnnotationStore()
        results, errors = bulk_annotate(
            ["a", "b", "c"], "todo",
            store=store,
        )
        assert errors == []
        assert len(results) == 3
        assert store.count == 3

    def test_invalid_status(self):
        results, errors = bulk_annotate(
            ["a", "b"], "invalid",
        )
        assert len(errors) > 0
        assert len(results) == 0

    def test_with_note(self):
        store = AnnotationStore()
        results, errors = bulk_annotate(
            ["a", "b"], "accepted_risk",
            note="Bulk accepted", author="bot",
            store=store,
        )
        assert errors == []
        for ann in store.list_all():
            assert ann.note == "Bulk accepted"
            assert ann.author == "bot"

    def test_empty_list(self):
        store = AnnotationStore()
        results, errors = bulk_annotate([], "todo", store=store)
        assert results == []
        assert errors == []


# ---------------------------------------------------------------------------
# remove_annotation
# ---------------------------------------------------------------------------

class TestRemoveAnnotation:
    def test_remove_existing(self):
        store = AnnotationStore()
        annotate_finding("a", "todo", store=store)
        assert remove_annotation("a", store) is True
        assert store.count == 0

    def test_remove_missing(self):
        store = AnnotationStore()
        assert remove_annotation("missing", store) is False


# ---------------------------------------------------------------------------
# filter_annotated
# ---------------------------------------------------------------------------

class TestFilterAnnotated:
    def test_no_annotations(self):
        findings = _sample_findings()
        result = filter_annotated(findings, {})
        assert len(result) == 3

    def test_filter_false_positive(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[0])
        annotate_finding(fid, "false_positive", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 2

    def test_filter_accepted_risk(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[1])
        annotate_finding(fid, "accepted_risk", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 2

    def test_filter_wont_fix(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[0])
        annotate_finding(fid, "wont_fix", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 2

    def test_todo_not_filtered(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[0])
        annotate_finding(fid, "todo", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 3

    def test_in_progress_not_filtered(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[0])
        annotate_finding(fid, "in_progress", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 3

    def test_custom_exclude_statuses(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid = make_finding_id_from_dict(findings[0])
        annotate_finding(fid, "todo", store=store)
        result = filter_annotated(
            findings, store,
            exclude_statuses=frozenset({"todo"}),
        )
        assert len(result) == 2

    def test_filter_with_dict(self):
        findings = _sample_findings()
        fid = make_finding_id_from_dict(findings[0])
        annotations = {
            fid: {
                "finding_id": fid,
                "status": "false_positive",
            },
        }
        result = filter_annotated(findings, annotations)
        assert len(result) == 2

    def test_filter_all(self):
        findings = _sample_findings()
        store = AnnotationStore()
        for f in findings:
            fid = make_finding_id_from_dict(f)
            annotate_finding(fid, "false_positive", store=store)
        result = filter_annotated(findings, store)
        assert len(result) == 0


# ---------------------------------------------------------------------------
# filter_by_status
# ---------------------------------------------------------------------------

class TestFilterByStatus:
    def test_basic(self):
        findings = _sample_findings()
        store = AnnotationStore()
        fid0 = make_finding_id_from_dict(findings[0])
        fid1 = make_finding_id_from_dict(findings[1])
        annotate_finding(fid0, "todo", store=store)
        annotate_finding(fid1, "false_positive", store=store)

        todos = filter_by_status(findings, store, "todo")
        assert len(todos) == 1

        fps = filter_by_status(findings, store, "false_positive")
        assert len(fps) == 1

    def test_no_matches(self):
        findings = _sample_findings()
        store = AnnotationStore()
        result = filter_by_status(findings, store, "todo")
        assert len(result) == 0

    def test_with_dict(self):
        findings = _sample_findings()
        fid = make_finding_id_from_dict(findings[0])
        annotations = {
            fid: {"finding_id": fid, "status": "todo"},
        }
        result = filter_by_status(findings, annotations, "todo")
        assert len(result) == 1


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

class TestPersistence:
    def test_save_and_load(self, tmp_path):
        path = tmp_path / "annotations.json"
        store = AnnotationStore()
        annotate_finding("a", "todo", note="Test", store=store)

        save_annotations(store, path)
        assert path.exists()

        loaded = load_annotations(path)
        assert loaded.count == 1
        ann = loaded.get("a")
        assert ann is not None
        assert ann.status == "todo"
        assert ann.note == "Test"

    def test_load_nonexistent(self, tmp_path):
        path = tmp_path / "missing.json"
        store = load_annotations(path)
        assert store.count == 0

    def test_load_invalid_json(self, tmp_path):
        path = tmp_path / "bad.json"
        path.write_text("not valid json", encoding="utf-8")
        store = load_annotations(path)
        assert store.count == 0

    def test_load_non_dict(self, tmp_path):
        path = tmp_path / "array.json"
        path.write_text("[1, 2, 3]", encoding="utf-8")
        store = load_annotations(path)
        assert store.count == 0

    def test_save_creates_file(self, tmp_path):
        path = tmp_path / "new.json"
        store = AnnotationStore()
        save_annotations(store, path)
        assert path.exists()

    def test_roundtrip_preserves_data(self, tmp_path):
        path = tmp_path / "rt.json"
        store = AnnotationStore()
        annotate_finding(
            "CMD-001:a.py:1", "false_positive",
            note="Checked", author="carlos",
            store=store,
        )
        save_annotations(store, path)
        loaded = load_annotations(path)
        ann = loaded.get("CMD-001:a.py:1")
        assert ann is not None
        assert ann.status == "false_positive"
        assert ann.note == "Checked"
        assert ann.author == "carlos"


# ---------------------------------------------------------------------------
# load_and_annotate
# ---------------------------------------------------------------------------

class TestLoadAndAnnotate:
    def test_basic(self, tmp_path):
        path = tmp_path / "ann.json"
        ann, errors = load_and_annotate(
            "CMD-001:a.py:1", "todo",
            note="Fix this", path=path,
        )
        assert errors == []
        assert ann.status == "todo"
        assert path.exists()

        # Verify it was saved
        loaded = load_annotations(path)
        assert loaded.count == 1

    def test_update_existing(self, tmp_path):
        path = tmp_path / "ann.json"
        load_and_annotate(
            "CMD-001:a.py:1", "todo", path=path,
        )
        load_and_annotate(
            "CMD-001:a.py:1", "false_positive",
            note="Verified", path=path,
        )
        loaded = load_annotations(path)
        assert loaded.count == 1
        ann = loaded.get("CMD-001:a.py:1")
        assert ann.status == "false_positive"

    def test_invalid_does_not_save(self, tmp_path):
        path = tmp_path / "ann.json"
        ann, errors = load_and_annotate(
            "CMD-001:a.py:1", "invalid",
            path=path,
        )
        assert len(errors) > 0
        assert not path.exists()


# ---------------------------------------------------------------------------
# Formatting
# ---------------------------------------------------------------------------

class TestFormatAnnotationsText:
    def test_empty(self):
        store = AnnotationStore()
        text = format_annotations_text(store)
        assert "No annotations" in text

    def test_with_annotations(self):
        store = AnnotationStore()
        annotate_finding(
            "CMD-001:a.py:1", "todo",
            note="Fix this", author="carlos",
            store=store,
        )
        text = format_annotations_text(store)
        assert "CMD-001:a.py:1" in text
        assert "todo" in text
        assert "Fix this" in text
        assert "carlos" in text

    def test_stats_shown(self):
        store = AnnotationStore()
        annotate_finding("a", "todo", store=store)
        annotate_finding("b", "false_positive", store=store)
        text = format_annotations_text(store)
        assert "todo: 1" in text
        assert "false_positive: 1" in text


class TestFormatAnnotationsJson:
    def test_empty(self):
        store = AnnotationStore()
        j = format_annotations_json(store)
        data = json.loads(j)
        assert data["count"] == 0
        assert data["annotations"] == {}

    def test_with_annotations(self):
        store = AnnotationStore()
        annotate_finding("a", "todo", store=store)
        j = format_annotations_json(store)
        data = json.loads(j)
        assert data["count"] == 1
        assert "a" in data["annotations"]

    def test_stats_included(self):
        store = AnnotationStore()
        annotate_finding("a", "todo", store=store)
        annotate_finding("b", "false_positive", store=store)
        j = format_annotations_json(store)
        data = json.loads(j)
        assert data["stats"]["todo"] == 1
        assert data["stats"]["false_positive"] == 1
