import json
import logging
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

# Global cache for lineage data to avoid repeated file reads
_lineage_data: Optional["LineageData"] = None


@dataclass
class Field:
    name: str
    path: str
    isLineage: bool
    relationship: Optional[Dict]


@dataclass
class Aspect:
    name: str
    fields: List[Field]


@dataclass
class Entity:
    name: str
    aspects: Dict[str, Aspect]


@dataclass
class LineageData:
    # entity name -> aspect
    entities: Dict[str, Entity]
    generated_by: str
    generated_at: str


def get_lineage_data() -> LineageData:
    """
    This is experimental internal API subject to breaking changes without prior notice.
    """
    global _lineage_data

    if _lineage_data is not None:
        return _lineage_data

    raw_data = _load_lineage_data()
    _entities = raw_data.get("entities", {})
    for entity_name, entity_data in _entities.items():
        entity = Entity(
            name=entity_name,
            aspects={},
        )
        for aspect_name, aspect_data in entity_data.items():
            entity.aspects[aspect_name] = Aspect(
                name=aspect_name,
                fields=[
                    Field(
                        name=field["name"],
                        path=field["path"],
                        isLineage=field["isLineage"],
                        relationship=field.get("relationship", None),
                    )
                    for field in aspect_data.get("fields", [])
                ],
            )
        _entities[entity_name] = entity

    _lineage_data = LineageData(
        entities=_entities,
        generated_by=raw_data.get("generated_by", ""),
        generated_at=raw_data.get("generated_at", ""),
    )
    return _lineage_data


def get_all_aspect_names() -> List[str]:
    """
    This is experimental internal API subject to breaking changes without prior notice.
    """
    entities = get_lineage_data().entities
    if not entities:
        return []
    first_entity = next(iter(entities.values()))
    return list(first_entity.aspects.keys())


def _load_lineage_data() -> Dict:
    """
    This is experimental internal API subject to breaking changes without prior notice.

    Load lineage data from the autogenerated lineage.json file.

    Returns:
        Dict containing the lineage information, or empty dict if file doesn't exist

    Raises:
        json.JSONDecodeError: If lineage.json is malformed
    """
    # Get the path to lineage.json relative to this file
    current_file = Path(__file__)
    lineage_file = current_file.parent / "lineage.json"

    if not lineage_file.exists():
        logger.warning(
            f"Lineage file not found: {lineage_file}. "
            "This may indicate a packaging issue. Lineage detection will be disabled."
        )
        return {}

    try:
        with open(lineage_file, "r") as f:
            return json.load(f)
    except json.JSONDecodeError as e:
        logger.error(
            f"Failed to parse lineage.json: {e}. Lineage detection will be disabled."
        )
        return {}


def _get_fields(entity_type: str, aspect_name: str) -> List[Dict]:
    """
    This is experimental internal API subject to breaking changes without prior notice.
    """
    lineage_data = get_lineage_data()
    entity = lineage_data.entities.get(entity_type)
    if not entity:
        return []

    aspect = entity.aspects.get(aspect_name)
    if not aspect:
        return []

    return [
        {
            "name": field.name,
            "path": field.path,
            "isLineage": field.isLineage,
            "relationship": field.relationship,
        }
        for field in aspect.fields
    ]


def _get_lineage_fields(entity_type: str, aspect_name: str) -> List[Dict]:
    """
    This is experimental internal API subject to breaking changes without prior notice.
    """
    return [
        field
        for field in _get_fields(entity_type, aspect_name)
        if field.get("isLineage", False)
    ]


@lru_cache(maxsize=128)
def is_lineage_aspect(entity_type: str, aspect_name: str) -> bool:
    """
    This is experimental internal API subject to breaking changes without prior notice.
    """
    return len(_get_lineage_fields(entity_type, aspect_name)) > 0


def clear_cache() -> None:
    """
    This is experimental internal API subject to breaking changes without prior notice.

    Clear the internal cache of lineage data.

    This is useful for testing or when the lineage.json file has been updated.
    """
    global _lineage_data
    _lineage_data = None
