"""Tests for JSON import functionality."""

import json
from datetime import datetime
from pathlib import Path
from unittest.mock import MagicMock, mock_open, patch

import pytest

from quber.db.importer import (
    get_import_stats,
    import_directory,
    import_json_file,
)
from quber.db.models import Document


class TestImportJsonFile:
    """Tests for import_json_file function."""

    @pytest.fixture
    def sample_json_data(self):
        """Sample JSON data for testing."""
        return {
            "document": "test.pdf",
            "extraction_date": "2024-01-01T12:00:00",
            "total_pages": 10,
            "total_tables": 2,
            "executive_summary": "Test summary",
            "model_provider": "anthropic",
            "model": "claude-3-5-sonnet-20241022",
            "tables": [
                {
                    "table_id": 0,
                    "page": 1,
                    "procedural_title": "Table 1",
                    "llm_title": "Financial Data",
                    "llm_description": "Quarterly financials",
                    "table_markdown": "| Q1 | Q2 |\n|---|---|\n| 100 | 200 |",
                    "headers": ["Q1", "Q2"],
                    "metadata": {"key": "value"},
                },
                {
                    "table_id": 1,
                    "page": 2,
                    "procedural_title": "Table 2",
                    "llm_title": "Sales Data",
                    "llm_description": "Monthly sales",
                    "table_markdown": "| Jan | Feb |\n|---|---|\n| 50 | 75 |",
                    "headers": ["Jan", "Feb"],
                    "metadata": {},
                },
            ],
        }

    def test_import_json_file_not_found(self):
        """Test that FileNotFoundError is raised for missing file."""
        with pytest.raises(FileNotFoundError, match="JSON file not found"):
            import_json_file("/nonexistent/file.json")

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    @patch("quber.db.importer.get_session")
    def test_import_json_file_basic(self, mock_get_session, mock_exists, mock_file, sample_json_data):
        """Test basic JSON file import."""
        mock_exists.return_value = True
        mock_file.return_value.read.return_value = json.dumps(sample_json_data)

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        with patch("json.load", return_value=sample_json_data):
            doc, table_count = import_json_file("/test/file.json")

        assert doc.filename == "test.pdf"
        assert doc.total_pages == 10
        assert doc.total_tables == 2
        assert table_count == 2

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    @patch("quber.db.importer.get_session")
    def test_import_json_file_with_embeddings(
        self, mock_get_session, mock_exists, mock_file, sample_json_data
    ):
        """Test JSON import with embedding generation."""
        mock_exists.return_value = True

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        mock_embedding_service = MagicMock()
        mock_embedding_service.embed_table_metadata.return_value = MagicMock(
            tolist=MagicMock(return_value=[0.1] * 1024)
        )

        with patch("json.load", return_value=sample_json_data):
            import_json_file(
                "/test/file.json",
                session=None,
                embedding_service=mock_embedding_service,
                generate_embeddings=True,
            )

        # Should call embed_table_metadata for each table
        assert mock_embedding_service.embed_table_metadata.call_count == 2

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    @patch("quber.db.importer.get_session")
    def test_import_json_file_duplicate(self, mock_get_session, mock_exists, mock_file, sample_json_data):
        """Test importing a duplicate file (already exists)."""
        mock_exists.return_value = True

        existing_doc = Document(id=1, filename="test.pdf", extraction_date=datetime.now(), total_tables=2)
        existing_doc.tables = [MagicMock(), MagicMock()]

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = existing_doc
        mock_session.query.return_value.filter_by.return_value.count.return_value = 2
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        with patch("json.load", return_value=sample_json_data):
            doc, table_count = import_json_file("/test/file.json")

        # Should return existing document
        assert doc == existing_doc
        assert table_count == 2
        # Should not add a new document
        mock_session.add.assert_not_called()

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    @patch("quber.db.importer.get_session")
    def test_import_json_file_invalid_date(self, mock_get_session, mock_exists, mock_file, sample_json_data):
        """Test import with invalid extraction_date falls back to current time."""
        mock_exists.return_value = True
        sample_json_data["extraction_date"] = "invalid-date"

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        with patch("json.load", return_value=sample_json_data):
            doc, _ = import_json_file("/test/file.json")

        # Should still create document with current time
        assert doc.extraction_date is not None

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    def test_import_json_file_with_provided_session(self, mock_exists, mock_file, sample_json_data):
        """Test import with a provided session."""
        mock_exists.return_value = True

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        with patch("json.load", return_value=sample_json_data):
            import_json_file("/test/file.json", session=mock_session)

        # Should use provided session
        mock_session.add.assert_called_once()
        mock_session.commit.assert_called_once()

    @patch("builtins.open", new_callable=mock_open)
    @patch("pathlib.Path.exists")
    def test_import_json_file_session_error(self, mock_exists, mock_file, sample_json_data):
        """Test that session errors are handled properly."""
        mock_exists.return_value = True

        mock_session = MagicMock()
        mock_session.query.return_value.filter_by.return_value.first.return_value = None
        mock_session.commit.side_effect = Exception("Database error")

        with patch("json.load", return_value=sample_json_data):
            with pytest.raises(Exception, match="Database error"):
                import_json_file("/test/file.json", session=mock_session)

        mock_session.rollback.assert_called_once()


class TestImportDirectory:
    """Tests for import_directory function."""

    def test_import_directory_not_found(self):
        """Test that FileNotFoundError is raised for missing directory."""
        with pytest.raises(FileNotFoundError, match="Directory not found"):
            import_directory("/nonexistent/directory")

    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.exists")
    def test_import_directory_no_files(self, mock_exists, mock_glob):
        """Test importing directory with no matching files."""
        mock_exists.return_value = True
        mock_glob.return_value = []

        docs, total_tables = import_directory("/test/dir")

        assert docs == []
        assert total_tables == 0

    @patch("quber.db.importer.import_json_file")
    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.exists")
    def test_import_directory_multiple_files(self, mock_exists, mock_glob, mock_import_file):
        """Test importing multiple files from directory."""
        mock_exists.return_value = True
        mock_glob.return_value = [Path("/test/file1.json"), Path("/test/file2.json")]

        # Mock import_json_file to return documents
        mock_doc1 = MagicMock(spec=Document)
        mock_doc2 = MagicMock(spec=Document)
        mock_import_file.side_effect = [(mock_doc1, 2), (mock_doc2, 3)]

        docs, total_tables = import_directory("/test/dir")

        assert len(docs) == 2
        assert total_tables == 5
        assert mock_import_file.call_count == 2

    @patch("quber.db.importer.import_json_file")
    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.exists")
    def test_import_directory_with_errors(self, mock_exists, mock_glob, mock_import_file):
        """Test that directory import continues on individual file errors."""
        mock_exists.return_value = True
        mock_glob.return_value = [
            Path("/test/file1.json"),
            Path("/test/file2.json"),
            Path("/test/file3.json"),
        ]

        # First file succeeds, second fails, third succeeds
        mock_doc1 = MagicMock(spec=Document)
        mock_doc3 = MagicMock(spec=Document)
        mock_import_file.side_effect = [
            (mock_doc1, 2),
            Exception("Import error"),
            (mock_doc3, 3),
        ]

        docs, total_tables = import_directory("/test/dir")

        # Should successfully import 2 out of 3 files
        assert len(docs) == 2
        assert total_tables == 5

    @patch("quber.db.importer.get_embedding_service")
    @patch("quber.db.importer.import_json_file")
    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.exists")
    def test_import_directory_with_embeddings(
        self, mock_exists, mock_glob, mock_import_file, mock_get_embedding
    ):
        """Test directory import with embedding generation."""
        mock_exists.return_value = True
        mock_glob.return_value = [Path("/test/file1.json")]

        mock_embedding_service = MagicMock()
        mock_get_embedding.return_value = mock_embedding_service

        mock_doc = MagicMock(spec=Document)
        mock_import_file.return_value = (mock_doc, 2)

        import_directory("/test/dir", generate_embeddings=True)

        # Should initialize embedding service once
        mock_get_embedding.assert_called_once()

        # Should pass embedding service to import_json_file
        call_kwargs = mock_import_file.call_args[1]
        assert call_kwargs["embedding_service"] == mock_embedding_service
        assert call_kwargs["generate_embeddings"] is True

    @patch("quber.db.importer.import_json_file")
    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.exists")
    def test_import_directory_custom_pattern(self, mock_exists, mock_glob, mock_import_file):
        """Test directory import with custom glob pattern."""
        mock_exists.return_value = True
        mock_glob.return_value = []

        import_directory("/test/dir", pattern="*.json")

        # Should use custom pattern
        mock_glob.assert_called_once_with("*.json")


class TestGetImportStats:
    """Tests for get_import_stats function."""

    @patch("quber.db.importer.get_session")
    def test_get_import_stats_empty_db(self, mock_get_session):
        """Test getting stats from empty database."""
        mock_session = MagicMock()
        mock_session.query.return_value.count.return_value = 0
        mock_session.query.return_value.all.return_value = []
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        stats = get_import_stats()

        assert stats["total_documents"] == 0
        assert stats["total_tables"] == 0
        assert stats["avg_tables_per_doc"] == 0
        assert stats["documents"] == []

    @patch("quber.db.importer.get_session")
    def test_get_import_stats_with_data(self, mock_get_session):
        """Test getting stats with data in database."""
        mock_doc1 = MagicMock(spec=Document)
        mock_doc1.filename = "doc1.pdf"
        mock_doc1.total_pages = 10
        mock_doc1.tables = [MagicMock(), MagicMock()]

        mock_doc2 = MagicMock(spec=Document)
        mock_doc2.filename = "doc2.pdf"
        mock_doc2.total_pages = 20
        mock_doc2.tables = [MagicMock(), MagicMock(), MagicMock()]

        mock_session = MagicMock()

        # Set up query mocks for different queries
        def query_side_effect(model):
            if model == Document:
                mock_query = MagicMock()
                mock_query.count.return_value = 2
                mock_query.all.return_value = [mock_doc1, mock_doc2]
                return mock_query
            else:  # ExtractedTable
                mock_query = MagicMock()
                mock_query.count.return_value = 5
                return mock_query

        mock_session.query.side_effect = query_side_effect
        mock_session.__enter__ = MagicMock(return_value=mock_session)
        mock_session.__exit__ = MagicMock(return_value=None)
        mock_get_session.return_value = mock_session

        stats = get_import_stats()

        assert stats["total_documents"] == 2
        assert stats["total_tables"] == 5
        assert stats["avg_tables_per_doc"] == 2.5
        assert len(stats["documents"]) == 2
        assert stats["documents"][0]["filename"] == "doc1.pdf"
        assert stats["documents"][0]["tables"] == 2
        assert stats["documents"][1]["tables"] == 3

    def test_get_import_stats_with_provided_session(self):
        """Test getting stats with a provided session."""
        mock_doc = MagicMock(spec=Document)
        mock_doc.filename = "test.pdf"
        mock_doc.total_pages = 5
        mock_doc.tables = [MagicMock()]

        mock_session = MagicMock()

        def query_side_effect(model):
            if model == Document:
                mock_query = MagicMock()
                mock_query.count.return_value = 1
                mock_query.all.return_value = [mock_doc]
                return mock_query
            else:
                mock_query = MagicMock()
                mock_query.count.return_value = 1
                return mock_query

        mock_session.query.side_effect = query_side_effect

        stats = get_import_stats(session=mock_session)

        assert stats["total_documents"] == 1
        assert stats["total_tables"] == 1