"""Tests for OpenMemo User Memory Layer — Store / Extractor / Recall / Ranker / API"""

import time
import pytest


# ── Fixtures ──────────────────────────────────────────────────────────────────

@pytest.fixture
def db_path(tmp_path):
    return str(tmp_path / "test_um.db")


@pytest.fixture
def store(db_path):
    from openmemo.user_memory.user_memory_store import UserMemoryStore
    s = UserMemoryStore(db_path=db_path)
    yield s
    s.close()


@pytest.fixture
def um(db_path):
    from openmemo.user_memory import UserMemory
    m = UserMemory(db_path=db_path)
    yield m
    m.close()


@pytest.fixture
def client(db_path):
    from flask import Flask
    from openmemo.user_memory.api import create_user_memory_blueprint
    app = Flask(__name__)
    app.register_blueprint(create_user_memory_blueprint(db_path=db_path))
    return app.test_client()


# ── UserMemoryStore ───────────────────────────────────────────────────────────

class TestStore:
    def test_put_and_get(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        node = UserMemoryNode(user_id="alice", type="preference",
                              key="backend", value="FastAPI")
        mid = store.put(node)
        fetched = store.get(mid)
        assert fetched is not None
        assert fetched.value == "FastAPI"

    def test_upsert_same_key(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        store.put(UserMemoryNode(user_id="alice", type="preference",
                                  key="backend", value="Flask"))
        store.put(UserMemoryNode(user_id="alice", type="preference",
                                  key="backend", value="FastAPI"))
        assert store.count("alice") == 1
        assert store.get_by_user("alice")[0].value == "FastAPI"

    def test_count(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        for i in range(3):
            store.put(UserMemoryNode(user_id="bob", type="preference",
                                     key=f"tool_{i}", value=f"Tool{i}"))
        assert store.count("bob") == 3
        assert store.count("alice") == 0

    def test_delete(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        node = UserMemoryNode(user_id="alice", type="constraint",
                              key="avoid_k8s", value="Kubernetes")
        store.put(node)
        assert store.delete(node.memory_id)
        assert store.get(node.memory_id) is None

    def test_delete_by_user(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        for i in range(4):
            store.put(UserMemoryNode(user_id="carol", type="preference",
                                     key=f"k{i}", value=f"v{i}"))
        assert store.delete_by_user("carol") == 4
        assert store.count("carol") == 0

    def test_filter_by_type(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        store.put(UserMemoryNode(user_id="dave", type="preference",
                                  key="backend", value="FastAPI"))
        store.put(UserMemoryNode(user_id="dave", type="constraint",
                                  key="avoid_k8s", value="Kubernetes"))
        prefs = store.get_by_user("dave", mem_type="preference")
        assert len(prefs) == 1
        assert prefs[0].type == "preference"

    def test_increment_usage(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        node = UserMemoryNode(user_id="alice", type="preference",
                              key="lang", value="Python")
        store.put(node)
        store.increment_usage(node.memory_id)
        store.increment_usage(node.memory_id)
        assert store.get(node.memory_id).usage_count == 2


# ── UserMemoryExtractor ───────────────────────────────────────────────────────

class TestExtractor:
    @pytest.fixture
    def ext(self):
        from openmemo.user_memory.user_memory_extractor import UserMemoryExtractor
        return UserMemoryExtractor()

    def test_preference_prefer(self, ext):
        nodes = ext.extract("u1", "I prefer FastAPI for backend")
        assert len(nodes) >= 1
        assert any(n.type == "preference" for n in nodes)

    def test_constraint_avoid(self, ext):
        nodes = ext.extract("u1", "avoid Kubernetes in all deployments")
        assert len(nodes) >= 1
        assert any(n.type == "constraint" for n in nodes)

    def test_constraint_never_use(self, ext):
        nodes = ext.extract("u1", "never use AWS lambda")
        assert any(n.type == "constraint" for n in nodes)

    def test_environment_kv(self, ext):
        nodes = ext.extract("u1", "database: PostgreSQL")
        env_nodes = [n for n in nodes if n.type == "environment"]
        assert len(env_nodes) >= 1
        assert env_nodes[0].value == "PostgreSQL"

    def test_multi_sentence(self, ext):
        text = "I prefer FastAPI. avoid Kubernetes. database: PostgreSQL"
        nodes = ext.extract("u1", text)
        assert len(nodes) >= 2

    def test_no_match(self, ext):
        nodes = ext.extract("u1", "hello world how are you")
        assert isinstance(nodes, list)

    def test_deduplication(self, ext):
        nodes = ext.extract("u1", "prefer FastAPI. I prefer FastAPI")
        keys = [n.key for n in nodes]
        assert len(keys) == len(set(keys))

    def test_explicit_store(self, ext):
        nodes = ext.extract("u1", "Use docker-compose instead of Kubernetes")
        assert isinstance(nodes, list)


# ── UserMemoryRecall ──────────────────────────────────────────────────────────

class TestRecall:
    def test_recall_empty(self, store):
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        recall = UserMemoryRecall(store)
        assert recall.recall("nobody") == []

    def test_recall_returns_top_n(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        for i in range(10):
            store.put(UserMemoryNode(user_id="alice", type="preference",
                                     key=f"tool_{i}", value=f"Tool{i}",
                                     confidence=0.9 - i * 0.05))
        recall = UserMemoryRecall(store)
        top = recall.recall("alice", top_n=3)
        assert len(top) == 3

    def test_ranking_by_score(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        store.put(UserMemoryNode(user_id="u", type="preference",
                                  key="low_conf", value="LowConf", confidence=0.3))
        store.put(UserMemoryNode(user_id="u", type="preference",
                                  key="high_conf", value="HighConf", confidence=0.95))
        recall = UserMemoryRecall(store)
        top = recall.recall("u", top_n=2)
        assert top[0].confidence >= top[1].confidence

    def test_format_for_prompt(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        store.put(UserMemoryNode(user_id="alice", type="preference",
                                  key="backend", value="FastAPI", confidence=0.9))
        store.put(UserMemoryNode(user_id="alice", type="constraint",
                                  key="avoid_k8s", value="Kubernetes", confidence=0.95))
        recall = UserMemoryRecall(store)
        block = recall.format_for_prompt("alice")
        assert "USER PROFILE" in block
        assert "FastAPI" in block
        assert "Kubernetes" in block

    def test_format_empty_user(self, store):
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        recall = UserMemoryRecall(store)
        block = recall.format_for_prompt("nobody")
        assert block == ""

    def test_to_dict_structure(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_recall import UserMemoryRecall
        store.put(UserMemoryNode(user_id="bob", type="preference",
                                  key="lang", value="Python"))
        recall = UserMemoryRecall(store)
        d = recall.to_dict("bob")
        assert "preferences" in d
        assert "constraints" in d
        assert "environment" in d
        assert d["user_id"] == "bob"


# ── UserMemoryRanker ──────────────────────────────────────────────────────────

class TestRanker:
    def test_rank_empty_workflows(self, store):
        from openmemo.user_memory.user_memory_ranker import UserMemoryRanker
        ranker = UserMemoryRanker(store)
        assert ranker.rank_workflows("u", []) == []

    def test_preference_boosts_matching_workflow(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_ranker import UserMemoryRanker
        store.put(UserMemoryNode(user_id="u", type="preference",
                                  key="deployment", value="docker-compose",
                                  confidence=0.95))
        ranker = UserMemoryRanker(store)
        wfs = [
            {"id": "wf_docker", "steps": ["docker-compose up"], "score": 0.7},
            {"id": "wf_k8s",    "steps": ["kubectl apply"],     "score": 0.7},
        ]
        ranked = ranker.rank_workflows("u", wfs)
        assert ranked[0]["id"] == "wf_docker"

    def test_constraint_penalizes_workflow(self, store):
        from openmemo.user_memory.user_memory_store import UserMemoryNode
        from openmemo.user_memory.user_memory_ranker import UserMemoryRanker
        store.put(UserMemoryNode(user_id="u", type="constraint",
                                  key="avoid_k8s", value="Kubernetes",
                                  confidence=0.95))
        ranker = UserMemoryRanker(store)
        # Both workflows have equal base score so penalty decides the winner
        wfs = [
            {"id": "wf_docker", "steps": ["docker-compose up"], "score": 0.8},
            {"id": "wf_k8s",    "steps": ["kubectl apply -f kubernetes/"], "score": 0.8},
        ]
        ranked = ranker.rank_workflows("u", wfs)
        assert ranked[0]["id"] == "wf_docker"

    def test_final_score_formula(self, store):
        from openmemo.user_memory.user_memory_ranker import UserMemoryRanker
        ranker = UserMemoryRanker(store)
        wfs = [{"id": "x", "steps": ["deploy"], "score": 0.8}]
        ranked = ranker.rank_workflows("nobody", wfs)
        assert 0.0 <= ranked[0]["final_score"] <= 1.0
        assert "user_alignment" in ranked[0]

    def test_alignment_score_range(self, store):
        from openmemo.user_memory.user_memory_ranker import UserMemoryRanker
        ranker = UserMemoryRanker(store)
        score = ranker.compute_alignment("nobody", "deploy with docker-compose")
        assert 0.0 <= score <= 1.0


# ── UserMemory Facade ─────────────────────────────────────────────────────────

class TestUserMemoryFacade:
    def test_store_text(self, um):
        ids = um.store("alice", "I prefer FastAPI. avoid Kubernetes")
        assert isinstance(ids, list)

    def test_put_explicit(self, um):
        mid = um.put("alice", "preference", "backend", "FastAPI")
        assert mid
        assert um.count("alice") == 1

    def test_recall(self, um):
        um.put("alice", "preference", "backend", "FastAPI")
        nodes = um.recall("alice")
        assert len(nodes) >= 1

    def test_format_for_prompt(self, um):
        um.put("alice", "preference", "backend", "FastAPI", confidence=0.9)
        um.put("alice", "constraint", "avoid_k8s", "Kubernetes", confidence=0.9)
        block = um.format_for_prompt("alice")
        assert "USER PROFILE" in block
        assert "FastAPI" in block

    def test_rank_workflows(self, um):
        um.put("alice", "preference", "deploy_tool", "docker-compose")
        wfs = [
            {"id": "wf1", "steps": ["docker-compose up"], "score": 0.7},
            {"id": "wf2", "steps": ["kubectl apply"],     "score": 0.7},
        ]
        ranked = um.rank_workflows("alice", wfs)
        assert ranked[0]["id"] == "wf1"

    def test_delete(self, um):
        mid = um.put("alice", "preference", "backend", "FastAPI")
        assert um.delete(mid)
        assert um.count("alice") == 0

    def test_delete_user(self, um):
        um.put("alice", "preference", "k1", "v1")
        um.put("alice", "constraint", "k2", "v2")
        n = um.delete_user("alice")
        assert n == 2
        assert um.count("alice") == 0

    def test_alignment_score(self, um):
        um.put("alice", "preference", "deploy", "docker-compose")
        score = um.alignment_score("alice", "run docker-compose up in production")
        assert 0.0 <= score <= 1.0

    def test_full_flow(self, um):
        um.store("bob", "I prefer FastAPI. avoid Kubernetes. database: PostgreSQL")
        block = um.format_for_prompt("bob")
        assert isinstance(block, str)
        d = um.to_dict("bob")
        assert "preferences" in d


# ── REST API ──────────────────────────────────────────────────────────────────

class TestRestAPI:
    def test_store_text(self, client):
        resp = client.post("/user-memory/store", json={
            "user_id": "alice",
            "text": "I prefer FastAPI and avoid Kubernetes",
        })
        assert resp.status_code == 201
        data = resp.get_json()
        assert data["status"] == "extracted"
        assert "memory_ids" in data
        assert "latency_ms" in data

    def test_store_explicit(self, client):
        resp = client.post("/user-memory/store", json={
            "user_id": "alice",
            "type": "preference",
            "key": "backend",
            "value": "FastAPI",
            "confidence": 0.95,
        })
        assert resp.status_code == 201
        assert resp.get_json()["status"] == "stored"

    def test_store_invalid_type(self, client):
        resp = client.post("/user-memory/store", json={
            "user_id": "alice",
            "type": "invalid_type",
            "key": "backend",
            "value": "FastAPI",
        })
        assert resp.status_code == 400

    def test_store_missing_body(self, client):
        resp = client.post("/user-memory/store", json={})
        assert resp.status_code == 400

    def test_recall(self, client):
        client.post("/user-memory/store", json={
            "user_id": "bob",
            "key": "backend", "value": "FastAPI", "type": "preference",
        })
        resp = client.get("/user-memory/recall?user_id=bob")
        assert resp.status_code == 200
        data = resp.get_json()
        assert "preferences" in data
        assert "latency_ms" in data

    def test_recall_by_type(self, client):
        client.post("/user-memory/store", json={
            "user_id": "carol",
            "key": "deploy", "value": "docker-compose", "type": "preference",
        })
        resp = client.get("/user-memory/recall?user_id=carol&type=preference")
        assert resp.status_code == 200

    def test_prompt(self, client):
        client.post("/user-memory/store", json={
            "user_id": "dave",
            "key": "backend", "value": "FastAPI", "type": "preference",
        })
        resp = client.get("/user-memory/prompt?user_id=dave")
        assert resp.status_code == 200
        data = resp.get_json()
        assert "prompt_block" in data
        assert "FastAPI" in data["prompt_block"]

    def test_count(self, client):
        client.post("/user-memory/store", json={
            "user_id": "eve", "key": "k", "value": "v", "type": "preference",
        })
        resp = client.get("/user-memory/count?user_id=eve")
        assert resp.status_code == 200
        assert resp.get_json()["count"] >= 1

    def test_delete_memory(self, client):
        r = client.post("/user-memory/store", json={
            "user_id": "frank", "key": "k", "value": "v", "type": "preference",
        })
        mid = r.get_json()["memory_id"]
        resp = client.delete(f"/user-memory/{mid}")
        assert resp.status_code == 200
        assert resp.get_json()["status"] == "deleted"

    def test_delete_nonexistent(self, client):
        resp = client.delete("/user-memory/nonexistent-id-999")
        assert resp.status_code == 404

    def test_delete_user(self, client):
        for i in range(3):
            client.post("/user-memory/store", json={
                "user_id": "grace", "key": f"k{i}",
                "value": f"v{i}", "type": "preference",
            })
        resp = client.delete("/user-memory/user/grace")
        assert resp.status_code == 200
        assert resp.get_json()["deleted"] == 3
