import logging
import os
from contextlib import asynccontextmanager
from functools import lru_cache
from typing import Any

import httpx
from fastembed import SparseTextEmbedding
from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from qdrant_client import AsyncQdrantClient, models

EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B"

# Ваш сервис должен считывать эти переменные из окружения (env), так как проверяющая система управляет ими
HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "8003"))

API_KEY = os.getenv("API_KEY")
EMBEDDINGS_DENSE_URL = os.getenv("EMBEDDINGS_DENSE_URL")
QDRANT_DENSE_VECTOR_NAME = os.getenv("QDRANT_DENSE_VECTOR_NAME", "dense")
QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse")
SPARSE_MODEL_NAME = "Qdrant/bm25"
RERANKER_MODEL = "nvidia/llama-nemotron-rerank-1b-v2"
RERANKER_URL = os.getenv("RERANKER_URL")
OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN")
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "evaluation")
REQUIRED_ENV_VARS = [
    "EMBEDDINGS_DENSE_URL",
    "RERANKER_URL",
    "QDRANT_URL",
]
 
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("search-service")


def validate_required_env() -> None:
    if bool(OPEN_API_LOGIN) != bool(OPEN_API_PASSWORD):
        raise RuntimeError("OPEN_API_LOGIN and OPEN_API_PASSWORD must be set together")

    if not API_KEY and not (OPEN_API_LOGIN and OPEN_API_PASSWORD):
        raise RuntimeError("Either API_KEY or OPEN_API_LOGIN and OPEN_API_PASSWORD must be set")

    missing_env_vars = [
        name for name in REQUIRED_ENV_VARS if os.getenv(name) is None or os.getenv(name) == ""
    ]
    if not missing_env_vars:
        return

    logger.error("Empty required env vars: %s", ", ".join(missing_env_vars))
    raise RuntimeError(f"Empty required env vars: {', '.join(missing_env_vars)}")


validate_required_env()


def get_upstream_request_kwargs() -> dict[str, Any]:
    headers = {"Content-Type": "application/json"}
    kwargs: dict[str, Any] = {"headers": headers}

    if OPEN_API_LOGIN and OPEN_API_PASSWORD:
        kwargs["auth"] = (OPEN_API_LOGIN, OPEN_API_PASSWORD)
        return kwargs

    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"

    return kwargs


# Модель данных, которую мы предоставляем и рассчитываем получать от вас
class DateRange(BaseModel):
    from_: str = Field(alias="from")
    to: str


class Entities(BaseModel):
    people: list[str] | None = None
    emails: list[str] | None = None
    documents: list[str] | None = None
    names: list[str] | None = None
    links: list[str] | None = None


class Question(BaseModel):
    text: str
    asker: str = ""
    asked_on: str = ""
    variants: list[str] | None = None
    hyde: list[str] | None = None
    keywords: list[str] | None = None
    entities: Entities | None = None
    date_mentions: list[str] | None = None
    date_range: DateRange | None = None
    search_text: str = ""


class SearchAPIRequest(BaseModel):
    question: Question


class SearchAPIItem(BaseModel):
    message_ids: list[str]


class SearchAPIResponse(BaseModel):
    results: list[SearchAPIItem]


class DenseEmbeddingItem(BaseModel):
    index: int
    embedding: list[float]


class DenseEmbeddingResponse(BaseModel):
    data: list[DenseEmbeddingItem]


class SparseVector(BaseModel):
    indices: list[int] = Field(default_factory=list)
    values: list[float] = Field(default_factory=list)


class SparseEmbeddingResponse(BaseModel):
    vectors: list[SparseVector]

# Метадата чанков в Qdrant'e, по которой вы можете фильтровать
class ChunkMetadata(BaseModel):
    chat_name: str
    chat_type: str # channel, group, private, thread
    chat_id: str
    chat_sn: str
    thread_sn: str | None = None
    message_ids: list[str]
    start: str
    end: str
    participants: list[str] = Field(default_factory=list)
    mentions: list[str] = Field(default_factory=list)
    contains_forward: bool = False
    contains_quote: bool = False


@lru_cache(maxsize=1)
def get_sparse_model() -> SparseTextEmbedding:
    logger.info("Loading local sparse model %s", SPARSE_MODEL_NAME)
    return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.http = httpx.AsyncClient()
    app.state.qdrant = AsyncQdrantClient(
        url=QDRANT_URL,
        api_key=API_KEY,
    )
    try:
        yield
    finally:
        await app.state.http.aclose()
        await app.state.qdrant.close()


app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)


# Внутри шаблона dense и rerank берутся из внешних HTTP endpoint'ов,
# которые предоставляет проверяющая система.
# Текущий код ниже — минимальный пример search pipeline.
DENSE_PREFETCH_K = 10
SPRASE_PREFETCH_K = 30
RETRIEVE_K = 20
RERANK_LIMIT = 10
FINAL_TOP_K = 50

async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
    # Dense endpoint ожидает OpenAI-compatible body с input как списком строк.
    response = await client.post(
        EMBEDDINGS_DENSE_URL,
        **get_upstream_request_kwargs(),
        json={
            "model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
            "input": [text],
        },
    )
    response.raise_for_status()

    payload = DenseEmbeddingResponse.model_validate(response.json())
    if not payload.data:
        raise ValueError("Dense embedding response is empty")

    return payload.data[0].embedding


async def embed_sparse(text: str) -> SparseVector:
    vectors = list(get_sparse_model().embed([text]))
    if not vectors:
        raise ValueError("Sparse embedding response is empty")

    item = vectors[0]
    return SparseVector(
        indices=[int(index) for index in item.indices.tolist()],
        values=[float(value) for value in item.values.tolist()],
    )

# ПЕРЕПИСАТЬ

async def qdrant_search(
    client: AsyncQdrantClient,
    dense_vector: list[float],
    sparse_vector: SparseVector,
    question_data: Question
) -> Any | None:
    must_conditions: []

    # Фильтр по диапазону дат (поле metadata.start в Qdrant) [cite: 147, 148, 175]
    if question_data.date_range:
        must_conditions.append(
            models.FieldCondition(
                key="metadata.start",
                range=models.Range(
                    gte=question_data.date_range.from_,
                    lte=question_data.date_range.to_
                )
            )
        )

    # Фильтр по автору вопроса (поле metadata.participants) [cite: 161, 163]
    if question_data.asker:
        must_conditions.append(
            models.FieldCondition(
                key="metadata.participants",
                match=models.MatchValue(value=question_data.asker)
            )
        )

    # Создаем итоговый объект фильтра, если есть условия
    search_filter = models.Filter(must=must_conditions) if must_conditions else None

    response = await client.query_points(
        collection_name=QDRANT_COLLECTION_NAME,
        prefetch=[
            models.Prefetch(
                query=dense_vector,
                using=QDRANT_DENSE_VECTOR_NAME,
                limit=DENSE_PREFETCH_K,
                filter=search_filter,
            ),
            models.Prefetch(
                query=models.SparseVector(
                    indices=sparse_vector.indices,
                    values=sparse_vector.values,
                ),
                using=QDRANT_SPARSE_VECTOR_NAME,
                limit=SPRASE_PREFETCH_K,
                filter=search_filter,
            ),
        ],
        query=models.FusionQuery(fusion=models.Fusion.RRF),
        limit=RETRIEVE_K,
        with_payload=True,
    )

    if not response.points:
        return None

    return response.points


async def qdrant_search_dense_only(
    client: AsyncQdrantClient,
    dense_vector: list[float],
    question_data: Question,
) -> Any | None:
    must_conditions: []

    if question_data.date_range:
        must_conditions.append(
            models.FieldCondition(
                key="metadata.start",
                range=models.Range(
                    gte=question_data.date_range.from_,
                    lte=question_data.date_range.to_,
                ),
            )
        )

    if question_data.asker:
        must_conditions.append(
            models.FieldCondition(
                key="metadata.participants",
                match=models.MatchValue(value=question_data.asker),
            )
        )

    search_filter = models.Filter(must=must_conditions) if must_conditions else None

    response = await client.query_points(
        collection_name=QDRANT_COLLECTION_NAME,
        prefetch=[
            models.Prefetch(
                query=dense_vector,
                using=QDRANT_DENSE_VECTOR_NAME,
                limit=DENSE_PREFETCH_K,
                filter=search_filter,
            ),
        ],
        query=models.FusionQuery(fusion=models.Fusion.RRF),
        limit=RETRIEVE_K,
        with_payload=True,
    )

    if not response.points:
        return None

    return response.points


def collect_query_variants(question: Question) -> list[str]:
    variants: list[str] = []
    seen: set[str] = set()

    def add_query(text: str | None) -> None:
        if text is None:
            return
        normalized = text.strip()
        if not normalized:
            return
        if normalized in seen:
            return
        seen.add(normalized)
        variants.append(normalized)

    add_query(question.search_text)
    add_query(question.text)

    for variant in question.variants or []:
        add_query(variant)

    return variants


def collect_hyde_queries(question: Question, base_queries: list[str]) -> list[str]:
    hyde_queries: list[str] = []
    seen: set[str] = set(base_queries)

    for hyde_query in question.hyde or []:
        normalized = hyde_query.strip()
        if not normalized:
            continue
        if normalized in seen:
            continue
        seen.add(normalized)
        hyde_queries.append(normalized)

    return hyde_queries


def build_sparse_query_text(question: Question, fallback_query: str) -> str:
    keywords: list[str] = []
    seen: set[str] = set()

    for keyword in question.keywords or []:
        normalized = keyword.strip()
        if not normalized:
            continue
        if normalized in seen:
            continue
        seen.add(normalized)
        keywords.append(normalized)

    if keywords:
        return " ".join(keywords)

    return fallback_query


def deduplicate_points(points: list[Any]) -> list[Any]:
    unique_points: list[Any] = []
    seen_ids: set[str] = set()

    for point in points:
        point_id = str(getattr(point, "id", ""))
        if not point_id:
            continue
        if point_id in seen_ids:
            continue
        seen_ids.add(point_id)
        unique_points.append(point)

    return unique_points


def extract_point_score(point: Any) -> float:
    score = getattr(point, "score", 0.0)
    if score is None:
        return 0.0
    return float(score)


def extract_message_ids(point: Any) -> list[str]:
    payload = point.payload or {}
    metadata = payload.get("metadata") or {}
    message_ids = metadata.get("message_ids") or []

    return [str(message_id) for message_id in message_ids]


async def get_rerank_scores(
    client: httpx.AsyncClient,
    label: str,
    targets: list[str],
) -> list[float]:
    if not targets:
        return []

    # Rerank endpoint возвращает score для пары query -> candidate text.
    response = await client.post(
        RERANKER_URL,
        **get_upstream_request_kwargs(),
        json={
            "model": RERANKER_MODEL,
            "encoding_format": "float",
            "text_1": label,
            "text_2": targets,
        },
    )
    response.raise_for_status()

    payload = response.json()
    data = payload.get("data") or []

    return [float(sample["score"]) for sample in data]


async def rerank_points(
    client: httpx.AsyncClient,
    query: str,
    points: list[Any],
) -> list[tuple[Any, float]]:
    rerank_candidates = points[:RERANK_LIMIT]
    tail_candidates = points[RERANK_LIMIT:]
    rerank_targets = [point.payload.get("page_content") for point in rerank_candidates]
    scores = await get_rerank_scores(client, query, rerank_targets)

    reranked_candidates = [
        (point, float(score))
        for score, point in sorted(
            zip(scores, rerank_candidates, strict=True),
            key=lambda item: item[0],
            reverse=True,
        )
    ]
    tail_with_scores = [
        (point, extract_point_score(point))
        for _, point in sorted(
            [(extract_point_score(point), point) for point in tail_candidates],
            key=lambda item: item[0],
            reverse=True,
        )
    ]

    return reranked_candidates + tail_with_scores


def aggregate_message_scores(scored_points: list[tuple[Any, float]]) -> dict[str, float]:
    aggregated_scores: dict[str, float] = {}

    for point, point_score in scored_points:
        point_message_ids = set(extract_message_ids(point))
        for message_id in point_message_ids:
            aggregated_scores[message_id] = aggregated_scores.get(message_id, 0.0) + point_score

    return aggregated_scores


def select_top_message_ids(aggregated_scores: dict[str, float], limit: int) -> list[str]:
    sorted_items = sorted(
        aggregated_scores.items(),
        key=lambda item: (-item[1], item[0]),
    )
    return [message_id for message_id, _ in sorted_items[:limit]]


# Ваш сервис должен имплементировать оба этих метода
@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/search", response_model=SearchAPIResponse)
async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
    queries = collect_query_variants(payload.question)
    if not queries:
        raise HTTPException(status_code=400, detail="question.search_text or question.text is required")

    hyde_queries = collect_hyde_queries(payload.question, queries)
    query = queries[0]
    client: httpx.AsyncClient = app.state.http
    qdrant: AsyncQdrantClient = app.state.qdrant

    all_points: list[Any] = []
    for query_variant in queries:
        dense_vector = await embed_dense(client, query_variant)
        sparse_query_text = build_sparse_query_text(payload.question, query_variant)
        sparse_vector = await embed_sparse(sparse_query_text)
        points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question)
        if points:
            all_points.extend(list(points))

    for hyde_query in hyde_queries:
        hyde_dense_vector = await embed_dense(client, hyde_query)
        hyde_points = await qdrant_search_dense_only(qdrant, hyde_dense_vector, payload.question)
        if hyde_points:
            all_points.extend(list(hyde_points))

    best_points = deduplicate_points(all_points)
    if not best_points:
        return SearchAPIResponse(results=[])

    scored_points = await rerank_points(client, query, list(best_points))
    aggregated_scores = aggregate_message_scores(scored_points)
    message_ids = select_top_message_ids(aggregated_scores, FINAL_TOP_K)

    return SearchAPIResponse(
        results=[SearchAPIItem(message_ids=message_ids)]
    )


@app.exception_handler(Exception)
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
    logger.exception(exc)
    detail = str(exc) or repr(exc)

    if isinstance(exc, RequestValidationError):
        return JSONResponse(status_code=422, content={"detail": exc.errors()})

    if isinstance(exc, HTTPException):
        return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})

    return JSONResponse(status_code=500, content={"detail": detail})


def main() -> None:
    import uvicorn

    uvicorn.run(
        "main:app",
        host=HOST,
        port=PORT,
        reload=False,
    )


if __name__ == "__main__":
    main()