from __future__ import annotations

import json
import math
import re
import urllib.request
import struct
from collections import Counter
from pathlib import Path
from typing import Any

import yaml


ROOT = Path(__file__).resolve().parents[1]
MODEL_CARDS_DIR = ROOT / "registry" / "model-cards"
TOKEN_RE = re.compile(r"[A-Za-zА-Яа-яЁё0-9_]+", re.UNICODE)
RUSSIAN_ENDINGS = (
    "иями",
    "ями",
    "ами",
    "ого",
    "ему",
    "ыми",
    "ими",
    "ой",
    "ей",
    "ых",
    "их",
    "ую",
    "юю",
    "ая",
    "яя",
    "ое",
    "ее",
    "ом",
    "ем",
    "ам",
    "ям",
    "ах",
    "ях",
    "ы",
    "и",
    "а",
    "я",
    "е",
    "у",
    "ю",
)
TOKEN_ALIASES = {
    "1с": ["1c", "bsl", "конфигурация"],
    "1c": ["1с", "bsl", "configuration"],
    "бсл": ["bsl", "1с"],
    "bsl": ["бсл", "1с"],
    "справочник": ["catalog", "справочники"],
    "справочники": ["справочник", "catalog"],
    "документ": ["documents", "документы"],
    "документы": ["документ", "documents"],
    "регистр": ["register", "регистры"],
    "регистры": ["регистр", "register"],
    "реквизит": ["attribute", "реквизиты"],
    "реквизиты": ["реквизит", "attribute"],
    "табличная": ["табличные", "часть"],
    "табличные": ["табличная", "часть"],
    "запрос": ["query", "read", "select", "выбрать"],
    "выбрать": ["запрос", "query", "select"],
    "форма": ["forms", "управляемая"],
    "модуль": ["module", "bsl"],
    "метаданные": ["metadata", "схема", "snapshot"],
    "схема": ["metadata", "метаданные"],
    "номенклатура": ["справочник", "catalog"],
}


def read_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8-sig") as handle:
        data = json.load(handle)
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a JSON object")
    return data


def read_jsonl(path: Path) -> list[dict]:
    if not path.exists():
        return []

    records: list[dict] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: invalid JSONL: {exc}") from exc
            if not isinstance(record, dict):
                raise ValueError(f"{path}:{line_number}: record must be an object")
            records.append(record)
    return records


def write_json(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def read_yaml_mapping(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a YAML mapping")
    return data


def iter_model_card_paths(*, include_examples: bool = False) -> list[Path]:
    paths = sorted(MODEL_CARDS_DIR.rglob("*.yaml")) + sorted(MODEL_CARDS_DIR.rglob("*.yml"))
    if include_examples:
        return paths
    return [path for path in paths if "examples" not in path.parts]


def load_model_card(card_id: str) -> dict[str, Any]:
    for suffix in (".yaml", ".yml"):
        path = MODEL_CARDS_DIR / f"{card_id}{suffix}"
        if path.exists():
            return read_yaml_mapping(path)
    raise FileNotFoundError(f"Model card not found for id `{card_id}` in {MODEL_CARDS_DIR}")


def localize_workspace_path(path: str) -> Path:
    if path.startswith("/workspace/"):
        if Path("/workspace").exists():
            return Path(path)
        return ROOT / path.removeprefix("/workspace/")
    if path.startswith("/models/"):
        if Path("/models").exists():
            return Path(path)
        return ROOT / "models" / path.removeprefix("/models/")
    return Path(path)


def stem_russian_token(token: str) -> str:
    if not re.search(r"[а-яё]", token, flags=re.IGNORECASE) or len(token) < 6:
        return token
    for ending in RUSSIAN_ENDINGS:
        if token.endswith(ending) and len(token) - len(ending) >= 4:
            return token[: -len(ending)]
    return token


def normalize_token(token: str) -> str:
    token = token.lower().replace("ё", "е")
    return stem_russian_token(token)


def tokenize(text: str, *, expand_aliases: bool = False) -> list[str]:
    tokens: list[str] = []
    for raw_token in TOKEN_RE.findall(text):
        token = normalize_token(raw_token)
        tokens.append(token)
        if expand_aliases:
            tokens.extend(normalize_token(alias) for alias in TOKEN_ALIASES.get(token, []))
    return tokens


def corpus_content_hash(records: list[dict]) -> str:
    import hashlib

    hasher = hashlib.sha256()
    for record in records:
        stable = {
            "id": record.get("id"),
            "document_id": record.get("document_id"),
            "source_path": record.get("source_path"),
            "chunk_index": record.get("chunk_index"),
            "title": record.get("title"),
            "content": record.get("content"),
            "metadata": record.get("metadata") or {},
        }
        hasher.update(json.dumps(stable, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8"))
        hasher.update(b"\n")
    return hasher.hexdigest()


def hashing_embedding(text: str, *, dimensions: int = 384) -> list[float]:
    import hashlib

    if dimensions < 8:
        raise ValueError("dimensions must be >= 8")
    vector = [0.0] * dimensions
    tokens = tokenize(text, expand_aliases=True)
    if not tokens:
        return vector
    counts = Counter(tokens)
    for token, count in counts.items():
        digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
        bucket = int.from_bytes(digest[:4], "little") % dimensions
        sign = 1.0 if digest[4] & 1 else -1.0
        vector[bucket] += sign * (1.0 + math.log(float(count)))
    norm = math.sqrt(sum(value * value for value in vector))
    if norm <= 0:
        return vector
    return [value / norm for value in vector]


def pack_float_vector(vector: list[float]) -> bytes:
    return struct.pack(f"<{len(vector)}f", *vector)


def unpack_float_vector(data: bytes, dimensions: int) -> list[float]:
    expected_size = dimensions * 4
    if len(data) != expected_size:
        raise ValueError(f"Vector blob has {len(data)} bytes, expected {expected_size}")
    return list(struct.unpack(f"<{dimensions}f", data))


def cosine_similarity(left: list[float], right: list[float]) -> float:
    if not left or not right or len(left) != len(right):
        return 0.0
    return float(sum(a * b for a, b in zip(left, right)))


def build_lexical_index(records: list[dict]) -> dict:
    documents = []
    document_frequency: Counter[str] = Counter()
    total_length = 0

    for record in records:
        content = record.get("content") or ""
        tokens = tokenize(content)
        title_tokens = tokenize(str(record.get("title") or ""))
        term_frequency = Counter(tokens)
        document_frequency.update(term_frequency.keys())
        total_length += len(tokens)
        documents.append(
            {
                "id": record.get("id"),
                "document_id": record.get("document_id"),
                "source_path": record.get("source_path"),
                "title": record.get("title"),
                "chunk_index": record.get("chunk_index"),
                "content": content,
                "metadata": record.get("metadata") or {},
                "length": len(tokens),
                "title_tokens": title_tokens,
                "term_frequency": dict(term_frequency),
            }
        )

    doc_count = len(documents)
    idf = {
        token: math.log((1 + doc_count) / (1 + frequency)) + 1
        for token, frequency in document_frequency.items()
    }

    return {
        "schema_version": 2,
        "type": "lexical-bm25",
        "doc_count": doc_count,
        "avg_doc_length": round(total_length / doc_count, 4) if doc_count else 0,
        "idf": idf,
        "documents": documents,
    }


def score_lexical_document(query_tf: Counter[str], document: dict, idf: dict, avg_doc_length: float) -> float:
    doc_tf = document.get("term_frequency") or {}
    doc_length = max(float(document.get("length") or 0), 1.0)
    avg_doc_length = max(float(avg_doc_length or doc_length), 1.0)
    title_tokens = set(document.get("title_tokens") or [])
    score = 0.0
    k1 = 1.4
    b = 0.72
    for token, query_count in query_tf.items():
        doc_count = doc_tf.get(token, 0)
        if doc_count:
            numerator = doc_count * (k1 + 1)
            denominator = doc_count + k1 * (1 - b + b * (doc_length / avg_doc_length))
            score += query_count * float(idf.get(token, 1.0)) * (numerator / denominator)
            if token in title_tokens:
                score += 0.35 * query_count
    return score


def search_lexical_index(
    index: dict,
    query: str,
    limit: int,
    *,
    candidate_limit: int | None = None,
    dedupe_by_document: bool = False,
    min_score: float = 0.0,
    source_types: list[str] | None = None,
    metadata_filters: dict[str, str] | None = None,
) -> list[dict]:
    query_tf = Counter(tokenize(query, expand_aliases=True))
    if not query_tf:
        return []

    results = []
    idf = index.get("idf") or {}
    avg_doc_length = float(index.get("avg_doc_length") or 0)
    allowed_source_types = {source_type.lower() for source_type in source_types or []}
    exact_filters = {key: str(value).lower() for key, value in (metadata_filters or {}).items() if str(value).strip()}
    for document in index.get("documents") or []:
        metadata = document.get("metadata") or {}
        source_type = str(metadata.get("source_type") or "").lower()
        if allowed_source_types and source_type not in allowed_source_types:
            continue
        if exact_filters and any(str(metadata.get(key) or "").lower() != value for key, value in exact_filters.items()):
            continue
        score = score_lexical_document(query_tf, document, idf, avg_doc_length)
        if score > min_score:
            results.append({"score": score, "document": document})

    results.sort(key=lambda item: item["score"], reverse=True)
    if candidate_limit:
        results = results[:candidate_limit]
    if dedupe_by_document:
        deduped = []
        seen = set()
        for result in results:
            document_id = result["document"].get("document_id") or result["document"].get("source_path")
            if document_id in seen:
                continue
            seen.add(document_id)
            deduped.append(result)
        results = deduped
    return results[:limit]


def call_chat_completion(
    base_url: str,
    model: str,
    messages: list[dict],
    *,
    temperature: float = 0.2,
    max_tokens: int = 1000,
    timeout: int = 180,
) -> str:
    url = f"{base_url.rstrip('/')}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        data = json.loads(response.read().decode("utf-8"))

    choices = data.get("choices") or []
    if not choices:
        raise ValueError("chat response has no choices")

    message = choices[0].get("message") or {}
    content = (message.get("content") or "").strip()
    if not content:
        raise ValueError("chat response content is empty")
    return content