import json
import sqlite3
import unittest
from unittest.mock import patch

from core.contracts import SearchResult
from core.sources.lora_manager_archive import (
    _build_result_from_row,
    _extract_search_tokens,
    _find_hash_match_row,
    _load_full_rows_for_versions,
    _normalize_model_type,
    _query_candidate_rows,
    search_lora_manager_archive_for_file,
)


class LoraManagerArchiveTests(unittest.TestCase):

    def test_normalize_model_type(self):
        self.assertEqual(_normalize_model_type("Loras"), "lora")
        self.assertEqual(_normalize_model_type("Checkpoints"), "checkpoint")
        self.assertEqual(_normalize_model_type("Embedding"), "textualinversion")
        self.assertEqual(_normalize_model_type("Custom"), "custom")

    def test_extract_search_tokens(self):
        tokens = _extract_search_tokens("my_awesome_model_v1.safetensors")
        self.assertIn("awesome", tokens["content_terms"])
        self.assertIn("model", tokens["content_terms"])
        self.assertIn("v1", tokens["version_terms"])

        # Test that generic tokens and extensions are filtered out
        tokens_generic = _extract_search_tokens("my_awesome_model_fp8_v1.sft")
        self.assertNotIn("fp8", tokens_generic["content_terms"])
        self.assertNotIn("sft", tokens_generic["content_terms"])
        self.assertIn("awesome", tokens_generic["content_terms"])
        self.assertIn("model", tokens_generic["content_terms"])
        self.assertIn("v1", tokens_generic["version_terms"])

    def test_query_candidate_rows_preserves_all_filter_plans(self):
        class EmptyQueryResult:
            def fetchall(self):
                return []

        class RecordingConnection:
            def __init__(self):
                self.calls = []

            def execute(self, sql, params):
                self.calls.append((sql, tuple(params)))
                return EmptyQueryResult()

        connection = RecordingConnection()

        candidates = _query_candidate_rows(
            connection,
            "aurora_style_v2.safetensors",
            "lora",
            2,
        )

        self.assertEqual([], candidates)
        self.assertEqual(5, len(connection.calls))
        self.assertTrue(
            all("m.type = ?" in sql for sql, _params in connection.calls)
        )
        self.assertTrue(
            all(
                "ORDER BY m.id ASC, COALESCE(v.position, 999999) ASC LIMIT ?"
                in sql
                for sql, _params in connection.calls
            )
        )
        self.assertEqual(
            [
                ("%aurora%", "%style%", "%v2%", "LORA", 50),
                ("%aurora%", "%style%", "LORA", 50),
                ("%aurora%", "%v2%", "LORA", 50),
                ("%aurora%", "LORA", 50),
                ("%v2%", "LORA", 50),
            ],
            [params for _sql, params in connection.calls],
        )

    def test_database_operations_with_in_memory_sqlite(self):
        # Create an in-memory database
        conn = sqlite3.connect(":memory:")
        conn.row_factory = sqlite3.Row

        # Create tables
        conn.execute("""
            CREATE TABLE models (
                id INTEGER PRIMARY KEY,
                name TEXT,
                type TEXT,
                data TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE model_versions (
                id INTEGER PRIMARY KEY,
                model_id INTEGER,
                name TEXT,
                base_model TEXT,
                position INTEGER,
                data TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE model_files (
                id INTEGER PRIMARY KEY,
                version_id INTEGER,
                type TEXT,
                data TEXT
            )
        """)

        # Insert dummy data
        conn.execute("INSERT INTO models VALUES (1, 'ghost_checkpoint', 'CHECKPOINT', ?)", (json.dumps({"tags": ["anime"]}),))
        conn.execute("INSERT INTO model_versions VALUES (10, 1, 'v1.0', 'SD 1.5', 1, ?)", (json.dumps({"trainedWords": ["ghost_word"]}),))
        archive_hash = "a" * 64
        conn.execute(
            "INSERT INTO model_files VALUES (100, 10, 'Model', ?)",
            (
                json.dumps(
                    {
                        "name": "ghost_v1.safetensors",
                        "downloadUrl": "https://civitai.com/api/download/models/10",
                        "sizeKB": 3,
                        "hashes": {"SHA256": archive_hash},
                    }
                ),
            ),
        )
        conn.commit()

        # Test querying candidates
        candidates = _query_candidate_rows(conn, "ghost_v1.safetensors", "checkpoint", 10)
        self.assertEqual(len(candidates), 1)
        self.assertEqual(candidates[0]["version_id"], 10)

        # Test loading full version rows
        full_rows = _load_full_rows_for_versions(conn, [10])
        self.assertIn(10, full_rows)
        row = full_rows[10]
        self.assertEqual(row["model_name"], "ghost_checkpoint")
        self.assertEqual(row["version_name"], "v1.0")

        # Test building result from row
        result = _build_result_from_row(conn, row, "ghost_checkpoint")
        self.assertIsNotNone(result)
        self.assertEqual(result.filename, "ghost_v1.safetensors")
        self.assertEqual(result.name, "ghost_checkpoint")
        self.assertEqual(list(result.trained_words), ["ghost_word"])
        self.assertEqual(result.size, 3 * 1024)
        self.assertEqual(result.sha256, archive_hash)
        self.assertEqual(dict(result.hashes), {"SHA256": archive_hash})

        hash_match = _find_hash_match_row(conn, archive_hash)
        self.assertIsNotNone(hash_match)
        self.assertEqual(hash_match[0]["version_id"], 10)
        self.assertEqual(hash_match[1]["name"], "ghost_v1.safetensors")

        with patch(
            "core.sources.lora_manager_archive._connect_readonly",
            return_value=conn,
        ):
            hash_result = search_lora_manager_archive_for_file(
                "",
                sha256=archive_hash,
            )

        self.assertIsNotNone(hash_result)
        self.assertEqual(hash_result.match_type, "hash")
        self.assertEqual(hash_result.confidence, 100.0)
        self.assertEqual(hash_result.sha256, archive_hash)

    def test_filename_search_keeps_typed_candidates(self):
        filename = "typed-lora-archive-model.safetensors"
        candidate = SearchResult(
            source="lora_manager_archive",
            name="typed-lora-archive-model",
            version_name="v1",
            filename=filename,
            download_url="https://example.com/typed-lora-archive-model.safetensors",
        )

        with patch(
            "core.sources.lora_manager_archive.search_lora_manager_archive",
            return_value=[candidate],
        ):
            result = search_lora_manager_archive_for_file(filename)

        self.assertIsInstance(result, SearchResult)
        self.assertEqual(filename, result.filename)
        self.assertEqual(100.0, result.confidence)
        self.assertEqual("exact", result.match_type)