from __future__ import annotations
import hashlib, json, time, math
from sqlalchemy import select, delete as sa_delete
from app.models.exact_cache import ExactCache
from app.models.semantic_cache import SemanticCache
from app.services.storage import async_session_maker


def compute_embedding(text: str) -> list[float]:
    """256-dim n-gram embedding using char-level n-grams (1-4)."""
    dim = 256
    vec = [0.0] * dim
    text_lower = text.lower()
    for n in range(1, 5):
        for i in range(len(text_lower) - n + 1):
            gram = text_lower[i:i + n]
            idx = int(hashlib.md5(gram.encode()).hexdigest()[:8], 16) % dim
            vec[idx] += 1.0
    norm = math.sqrt(sum(v * v for v in vec))
    if norm > 0:
        vec = [v / norm for v in vec]
    return vec


def cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(ai * bi for ai, bi in zip(a, b))
    na = math.sqrt(sum(ai * ai for ai in a))
    nb = math.sqrt(sum(bi * bi for bi in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


def make_cache_key(model: str, messages: list, **kwargs) -> str:
    raw = json.dumps({"model": model, "messages": messages, **kwargs}, sort_keys=True)
    return hashlib.sha256(raw.encode()).hexdigest()


async def get_exact(session, key: str) -> dict | None:
    result = await session.execute(select(ExactCache).where(ExactCache.cache_key == key))
    row = result.scalar_one_or_none()
    if not row:
        return None
    if time.time() >= row.expires_at:
        await session.delete(row)
        await session.commit()
        return None
    row.hit_count += 1
    row.last_hit = time.time()
    await session.commit()
    return {"response": row.response, "hit_count": row.hit_count, "created_at": row.created_at}


async def store_exact(session, key: str, response: str, ttl: int = 3600) -> None:
    now = time.time()
    existing = await session.execute(select(ExactCache).where(ExactCache.cache_key == key))
    row = existing.scalar_one_or_none()
    if row:
        row.response = response
        row.expires_at = now + ttl
        row.created_at = now
    else:
        session.add(ExactCache(
            cache_key=key, response=response,
            created_at=now, expires_at=now + ttl,
        ))
    await session.commit()


async def get_semantic(session, model: str, query: str, threshold: float = 0.95) -> dict | None:
    query_emb = compute_embedding(query)
    now = time.time()
    result = await session.execute(
        select(SemanticCache).where(
            SemanticCache.model == model,
            SemanticCache.expires_at > now,
        )
    )
    for row in result.scalars().all():
        cached_emb = json.loads(row.embedding)
        sim = cosine_similarity(query_emb, cached_emb)
        if sim >= threshold:
            row.hit_count += 1
            await session.commit()
            return {"response": row.response, "similarity": sim, "hit_count": row.hit_count}
    return None


async def store_semantic(session, model: str, query: str, response: str, ttl: int = 3600) -> None:
    now = time.time()
    emb = compute_embedding(query)
    key = hashlib.sha256(json.dumps({"model": model, "query": query}, sort_keys=True).encode()).hexdigest()
    session.add(SemanticCache(
        cache_key=key, model=model, query_text=query,
        response=response, embedding=json.dumps(emb),
        created_at=now, expires_at=now + ttl,
    ))
    await session.commit()


async def get_stats(session) -> dict:
    exact_count = await session.execute(select(ExactCache.cache_key))
    exact_rows = exact_count.scalars().all()
    exact_active = 0
    now = time.time()
    for row in exact_rows:
        r = await session.execute(select(ExactCache).where(ExactCache.cache_key == row))
        obj = r.scalar_one_or_none()
        if obj and obj.expires_at > now:
            exact_active += 1

    sem_count = await session.execute(select(SemanticCache.cache_key))
    sem_rows = sem_count.scalars().all()
    sem_active = 0
    for row in sem_rows:
        r = await session.execute(select(SemanticCache).where(SemanticCache.cache_key == row))
        obj = r.scalar_one_or_none()
        if obj and obj.expires_at > now:
            sem_active += 1

    return {
        "exact_total": len(exact_rows),
        "exact_active": exact_active,
        "semantic_total": len(sem_rows),
        "semantic_active": sem_active,
    }


async def clear_exact(session) -> None:
    await session.execute(sa_delete(ExactCache))
    await session.commit()


async def clear_semantic(session) -> None:
    await session.execute(sa_delete(SemanticCache))
    await session.commit()


async def clear_all(session) -> None:
    await clear_exact(session)
    await clear_semantic(session)


async def evict_expired(session) -> None:
    now = time.time()
    await session.execute(sa_delete(ExactCache).where(ExactCache.expires_at <= now))
    await session.execute(sa_delete(SemanticCache).where(SemanticCache.expires_at <= now))
    await session.commit()