import logging
import os
import time
from dataclasses import dataclass
from io import BytesIO

import numpy as np
import requests
from PIL import Image
from search_rdf import Data, EmbeddingIndex, FuzzyIndex, KeywordIndex
from search_rdf.model import (
    HuggingFaceImageModel,
    OpenClipModel,
    SentenceTransformerModel,
)
from universal_ml_utils.configuration import load_config
from universal_ml_utils.io import dump_json, load_json, load_text
from universal_ml_utils.logging import get_logger

from grasp.manager.cache import Cache
from grasp.manager.normalizer import Normalizer, WikidataPropertyNormalizer
from grasp.sparql.types import ObjType
from grasp.sparql.utils import find_longest_prefix, get_endpoint, load_qlever_prefixes
from grasp.utils import get_index_dir

SearchIndex = KeywordIndex | EmbeddingIndex | FuzzyIndex
EmbeddingModel = HuggingFaceImageModel | OpenClipModel | SentenceTransformerModel


@dataclass
class Index:
    description: str
    index: SearchIndex
    info_sparql: str | None = None
    cache: Cache | None = None


def load_data(index_dir: str) -> Data:
    try:
        data = Data.load(os.path.join(index_dir, "data"))
    except Exception as e:
        raise ValueError(f"Failed to load index data from {index_dir}") from e

    return data


def load_index(index_dir: str, index_type: str) -> SearchIndex | None:
    logger = get_logger("KG INDEX LOADING")
    start = time.perf_counter()

    try:
        data = load_data(index_dir)
    except Exception as e:
        logger.warning(f"Failed to load data from {index_dir}: {e}")
        return None

    index_dir = os.path.join(index_dir, index_type)
    load_kwargs = {"data": data, "index_dir": index_dir}

    if index_type == "keyword":
        index_cls = KeywordIndex
    elif index_type == "fuzzy":
        index_cls = FuzzyIndex
    elif index_type == "embedding":
        index_cls = EmbeddingIndex
        load_kwargs["embedding_path"] = os.path.join(index_dir, "embedding.safetensors")
    else:
        raise ValueError(f"Unknown index type {index_type}")

    try:
        index = index_cls.load(**load_kwargs)
    except Exception as e:
        logger.warning(f"Failed to load {index_type} index from {index_dir}: {e}")
        return None

    end = time.perf_counter()

    logger.debug(f"Loading {index_type} index from {index_dir} took {end - start:.2f}s")

    return index


def load_entity_index(kg: str, index_type: str) -> SearchIndex | None:
    index_dir = os.path.join(get_index_dir(kg), "entities")
    return load_index(index_dir, index_type)


def load_property_index(kg: str, index_type: str) -> SearchIndex | None:
    index_dir = os.path.join(get_index_dir(kg), "properties")
    return load_index(index_dir, index_type)


def load_info_sparql(
    index_dir: str,
    logger: logging.Logger | None = None,
) -> str | None:
    info_sparql_path = os.path.join(index_dir, "info.sparql")
    if os.path.exists(info_sparql_path):
        if logger is not None:
            logger.debug(f"Loaded info.sparql from {index_dir}")
        return load_text(info_sparql_path)

    if logger is not None:
        logger.debug(f"No info.sparql found at {index_dir}")
    return None


def load_info_cache(
    index_dir: str,
    logger: logging.Logger | None = None,
) -> Cache | None:
    cache_dir = os.path.join(index_dir, "info.cache", "db")
    if not os.path.exists(cache_dir):
        if logger is not None:
            logger.debug(f"No info cache found at {cache_dir}")
        return None

    try:
        start = time.perf_counter()
        cache = Cache.load(cache_dir)
        end = time.perf_counter()
        if logger is not None:
            logger.debug(f"Loaded cache from {cache_dir} in {end - start:.2f}s")
        return cache
    except Exception as e:
        if logger is not None:
            logger.warning(f"Failed to load cache from {cache_dir}: {e}")
        return None


def load_other_indices(kg: str, indices: list[str]) -> dict[str, Index]:
    logger = get_logger("KG OTHER INDICES LOADING")
    base_index_dir = get_index_dir(kg)
    config_path = os.path.join(base_index_dir, "indices.yaml")
    if not os.path.exists(config_path):
        logger.debug(f"No indices.yaml found at {config_path}, skipping other indices")
        return {}

    config = load_config(config_path)

    others = {}
    for cfg in config["indices"]:
        name = cfg["name"]
        if name not in indices:
            logger.debug(
                f"Skipping index {name} as it's not in the specified indices list"
            )
            continue

        desc = cfg.get("description", "No description provided")

        sub_index_dir = os.path.join(base_index_dir, name)

        # normalize embedding index type, which is called embedding-with-data
        # in search-rdf, but embedding in the search-rdf python interface
        if cfg["type"].startswith("embedding"):
            cfg["type"] = "embedding"

        index = load_index(sub_index_dir, cfg["type"])
        if index is None:
            continue

        info_sparql = load_info_sparql(sub_index_dir, logger)
        info_cache = load_info_cache(sub_index_dir, logger)

        others[name] = Index(desc, index, info_sparql, info_cache)

    return others


def get_embedding_model_key(index: EmbeddingIndex) -> str:
    assert index.model is not None, "Embedding index must have model metadata"
    provider = index.provider or "sentence-transformer"
    return f"{provider}/{index.model}"


def load_embedding_model(
    index: SearchIndex,
    models: dict[str, EmbeddingModel],
) -> dict[str, EmbeddingModel]:
    if not index.index_type == "embedding":
        return models

    assert isinstance(index, EmbeddingIndex), "Expected an EmbeddingIndex"
    assert index.model is not None, "Embedding index must have model metadata"

    key = get_embedding_model_key(index)
    if key in models:
        return models

    provider = index.provider or "sentence-transformer"
    if provider == "sentence-transformer":
        model = SentenceTransformerModel(index.model)
    elif provider == "open-clip":
        model = OpenClipModel(index.model)
    elif provider == "huggingface-image":
        model = HuggingFaceImageModel(index.model)
    else:
        raise ValueError(f"Unknown embedding model provider {provider}")

    models[key] = model
    return models


def load_kg_normalizers(kg: str) -> tuple[Normalizer, Normalizer]:
    ent_normalizer = Normalizer()
    prop_normalizer = ent_normalizer
    if kg.startswith("wikidata"):
        prop_normalizer = WikidataPropertyNormalizer()
    return ent_normalizer, prop_normalizer


def load_kg_info(
    kg: str,
    endpoint: str | None = None,
) -> tuple[dict[str, str], str | None]:
    kg_index_dir = get_index_dir(kg)
    info_file = os.path.join(kg_index_dir, "info.json")
    prefix_file = os.path.join(kg_index_dir, "prefixes.json")

    # load info.json if it exists, fall back to empty
    if os.path.exists(info_file):
        info = load_json(info_file)
    else:
        info = {}

    description = info.get("description")
    prefixes = info.get("prefixes", {})

    # if no prefixes in info.json, try prefixes.json (legacy)
    if not prefixes and os.path.exists(prefix_file):
        prefixes = load_json(prefix_file)

    # if still no prefixes, try to fetch from QLever endpoint
    if not prefixes:
        try:
            prefixes = load_qlever_prefixes(endpoint or get_endpoint(kg))
        except Exception:
            prefixes = {}

    # save prefixes back to info.json for future use
    if prefixes and not info.get("prefixes"):
        info["prefixes"] = prefixes
        dump_json(info, info_file, indent=2)

    # compatibility with IRIs: strip leading < and trailing >
    prefixes = {k: v.lstrip("<").rstrip(">") for k, v in prefixes.items()}

    common_prefixes = get_common_sparql_prefixes()
    values = set(prefixes.values())

    # add common prefixes that might not be covered by the
    # specified prefixes
    for short, long in common_prefixes.items():
        if short in prefixes or long in values:
            continue

        prefixes[short] = long

    return prefixes, description


def load_kg_info_sparqls(kg: str) -> tuple[str | None, str | None]:
    logger = get_logger("KG INFO SPARQL LOADING")
    kg_index_dir = get_index_dir(kg)
    ent_info = load_info_sparql(os.path.join(kg_index_dir, "entities"), logger)
    prop_info = load_info_sparql(os.path.join(kg_index_dir, "properties"), logger)
    return ent_info, prop_info


def load_kg_info_caches(kg: str) -> tuple[Cache | None, Cache | None]:
    logger = get_logger("KG INFO CACHE LOADING")
    kg_index_dir = get_index_dir(kg)
    ent_cache = load_info_cache(os.path.join(kg_index_dir, "entities"), logger)
    prop_cache = load_info_cache(os.path.join(kg_index_dir, "properties"), logger)
    return ent_cache, prop_cache


def load_kg_indices(
    kg: str,
    entities_type: str,
    properties_type: str,
) -> tuple[SearchIndex | None, SearchIndex | None]:
    ent_index = load_entity_index(kg, entities_type)
    prop_index = load_property_index(kg, properties_type)
    return ent_index, prop_index


def get_common_sparql_prefixes() -> dict[str, str]:
    return {
        "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
        "rdfs": "http://www.w3.org/2000/01/rdf-schema#",
        "owl": "http://www.w3.org/2002/07/owl#",
        "xsd": "http://www.w3.org/2001/XMLSchema#",
        "foaf": "http://xmlns.com/foaf/0.1/",
        "skos": "http://www.w3.org/2004/02/skos/core#",
        "dct": "http://purl.org/dc/terms/",
        "dc": "http://purl.org/dc/elements/1.1/",
        "prov": "http://www.w3.org/ns/prov#",
        "schema": "http://schema.org/",
        "geo": "http://www.opengis.net/ont/geosparql#",
        "geosparql": "http://www.opengis.net/ont/geosparql#",
        "geof": "http://www.opengis.net/def/function/geosparql/",
        "gn": "http://www.geonames.org/ontology#",
        "bd": "http://www.bigdata.com/rdf#",
        "hint": "http://www.bigdata.com/queryHints#",
        "wikibase": "http://wikiba.se/ontology#",
        "qb": "http://purl.org/linked-data/cube#",
        "void": "http://rdfs.org/ns/void#",
    }


def find_obj_type_from_prefixes(
    iri: str,
    prefixes: dict[str, str],
    common_prefixes: dict[str, str],
) -> ObjType:
    # we have three cases:
    # 1. the IRI matches a common prefix but not a known prefix -> COMMON (e.g. rdf:type)
    # 2. the IRI matches a known prefix -> UNINDEXED (e.g. wd:Q42)
    # 3. the IRI matches neither -> UNKNOWN (e.g. fb:en.barack_obama in Wikidata)
    is_common = find_longest_prefix(iri, common_prefixes) is not None
    is_known = find_longest_prefix(iri, prefixes) is not None

    if is_common:
        return ObjType.COMMON
    elif is_known:
        return ObjType.UNINDEXED
    else:
        return ObjType.UNKNOWN


def load_image_from_url(url: str) -> np.ndarray:
    try:
        if url.startswith("file://"):
            path = url[len("file://") :]
            image = Image.open(path).convert("RGB")
        else:
            response = requests.get(url, headers={"User-Agent": "grasp-rdf"})
            response.raise_for_status()
            image = Image.open(BytesIO(response.content)).convert("RGB")
        return np.array(image)
    except Exception as e:
        raise IOError(f"Failed to load image from {url}: {e}") from e


def format_index_meta(index: SearchIndex) -> str:
    parts = [f'type="{index.index_type}"']
    if isinstance(index, EmbeddingIndex):
        modalities = index.modality or ["text"]
        parts.append(f'modalities="{"+".join(modalities)}"')
    return ", ".join(parts)


def describe_index_type(index_type: str) -> str:
    if index_type == "keyword":
        return "Retrieves items by overlap between their label words and \
the query keywords. The query keywords can match label words exactly or \
as prefixes. No special query operators like AND/OR are supported."

    elif index_type == "fuzzy":
        return "Retrieves items by overlap between their label words and \
the query keywords. The query keywords must not match label words exactly, but \
some fuzziness is allowed. The longer a query keyword is, the more it can deviate \
from a label word and still be considered a match, though it will also contribute \
less to the overall score. No special query operators like AND/OR are supported."

    elif index_type == "embedding":
        return "Retrieves items by cosine similarity between their \
embeddings and the query embedding. The embedding model used depends on the \
index and may support text, images, or both."

    else:
        raise ValueError(f"Unknown index type {index_type}")
