"""
Relationship Extraction Mixin

Provides relationship extraction from dataclass type hints
for the ModuleWrapper system.
"""

import dataclasses
import inspect
import logging
from collections import defaultdict
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Set,
    Tuple,
    Union,
    get_args,
    get_origin,
    get_type_hints,
)

from .core import BUILTIN_PREFIXES, PRIMITIVE_TYPES

logger = logging.getLogger(__name__)


def _get_nl_relationship_patterns() -> Dict[tuple, str]:
    """Get NL relationship patterns (lazy load to avoid circular imports)."""
    try:
        from gchat.smart_card_builder import SmartCardBuilder

        return SmartCardBuilder.NL_RELATIONSHIP_PATTERNS
    except ImportError:
        return {
            ("Image", "OnClick"): "clickable image, image with click action",
            ("Button", "OnClick"): "button click action, button that opens link",
        }


# =============================================================================
# RELATIONSHIPS MIXIN
# =============================================================================


class RelationshipsMixin:
    """
    Mixin providing relationship extraction functionality.

    Expects the following attributes on self:
    - components: Dict[str, ModuleComponent]
    - client: Qdrant client
    - collection_name: str
    - embedder, colbert_embedder: Embedding models
    """

    # Cached properties
    _cached_relationships: Optional[Dict[str, List[str]]] = None

    def _is_dataclass_type(self, cls: type) -> bool:
        """Check if a class is a dataclass."""
        return dataclasses.is_dataclass(cls) and isinstance(cls, type)

    def _unwrap_optional(self, field_type: type) -> Tuple[type, bool]:
        """
        Unwrap Optional[X] to get X and whether it was optional.

        Returns:
            Tuple of (unwrapped_type, is_optional)
        """
        origin = get_origin(field_type)

        # Handle Optional[X] which is Union[X, None]
        if origin is Union:
            args = get_args(field_type)
            non_none_args = [t for t in args if t is not type(None)]
            if len(non_none_args) == 1:
                return non_none_args[0], True
            return field_type, False

        return field_type, False

    def _unwrap_container_type(self, field_type: type) -> List[type]:
        """
        Extract inner types from container types like List[T], Tuple[T, ...], Set[T].

        Returns:
            List of inner types (or the original type if not a container)
        """
        origin = get_origin(field_type)

        # Handle List[T], Tuple[T, ...], Set[T], FrozenSet[T]
        if origin in (list, tuple, set, frozenset):
            args = get_args(field_type)
            if args:
                return [a for a in args if a is not Ellipsis and a is not type(None)]
            return []

        # Handle Optional[List[T]]
        if origin is Union:
            args = get_args(field_type)
            non_none_args = [t for t in args if t is not type(None)]
            if len(non_none_args) == 1:
                return self._unwrap_container_type(non_none_args[0])

        return [field_type]

    def _is_component_type(self, field_type: type) -> bool:
        """Check if a type represents a component (not primitive)."""
        if field_type in PRIMITIVE_TYPES:
            return False

        if not inspect.isclass(field_type):
            return False

        module = getattr(field_type, "__module__", "")
        if any(module.startswith(prefix) for prefix in BUILTIN_PREFIXES):
            return False

        return True

    def _to_camel_case(self, snake_str: str) -> str:
        """Convert snake_case to camelCase."""
        components = snake_str.split("_")
        return components[0] + "".join(x.title() for x in components[1:])

    def _generate_nl_description(self, parent: str, child: str) -> str:
        """Generate natural language description for a relationship."""
        patterns = _get_nl_relationship_patterns()
        key = (parent, child)
        if key in patterns:
            return patterns[key]

        return f"{parent.lower()} with {child.lower()}, {parent.lower()} containing {child.lower()}"

    def _extract_relationships_from_class(
        self,
        cls: type,
        visited: Optional[Set[str]] = None,
        depth: int = 0,
        max_depth: int = 5,
        path_prefix: str = "",
    ) -> List[Dict[str, Any]]:
        """
        Extract field relationships from a class using type hints.

        Args:
            cls: The class to analyze
            visited: Set of visited class names to prevent cycles
            depth: Current recursion depth
            max_depth: Maximum recursion depth
            path_prefix: Prefix for nested paths

        Returns:
            List of relationship dictionaries
        """
        if visited is None:
            visited = set()

        class_name = cls.__name__
        if class_name in visited or depth > max_depth:
            return []

        visited.add(class_name)
        relationships = []

        try:
            hints = get_type_hints(cls)
        except Exception as e:
            logger.debug(f"Could not get type hints for {class_name}: {e}")
            return []

        for field_name, field_type in hints.items():
            if field_name.startswith("_"):
                continue

            unwrapped_type, is_optional = self._unwrap_optional(field_type)
            inner_types = self._unwrap_container_type(unwrapped_type)

            for inner_type in inner_types:
                if not self._is_component_type(inner_type):
                    continue

                child_name = inner_type.__name__
                child_module = getattr(inner_type, "__module__", "")

                if path_prefix:
                    rel_path = f"{path_prefix}.{field_name}"
                else:
                    rel_path = f"{class_name.lower()}.{field_name}"

                json_field = self._to_camel_case(field_name)
                if path_prefix:
                    json_path = f"{path_prefix}.{json_field}"
                else:
                    json_path = f"{class_name[0].lower()}{class_name[1:]}.{json_field}"

                relationship = {
                    "parent_class": class_name,
                    "parent_module": cls.__module__,
                    "parent_path": f"{cls.__module__}.{class_name}",
                    "child_class": child_name,
                    "child_module": child_module,
                    "child_path": f"{child_module}.{child_name}",
                    "field_name": field_name,
                    "is_optional": is_optional,
                    "depth": depth + 1,
                    "relationship_path": rel_path,
                    "json_path": json_path,
                    "nl_description": self._generate_nl_description(
                        class_name, child_name
                    ),
                }
                relationships.append(relationship)

                # Recursively extract from child class
                if self._is_dataclass_type(inner_type):
                    child_relationships = self._extract_relationships_from_class(
                        inner_type,
                        visited=visited.copy(),
                        depth=depth + 1,
                        max_depth=max_depth,
                        path_prefix=json_path,
                    )
                    for child_rel in child_relationships:
                        child_rel["root_parent"] = class_name
                        child_rel["intermediate_path"] = rel_path
                    relationships.extend(child_relationships)

        return relationships

    def extract_relationships(self, max_depth: int = 5) -> List[Dict[str, Any]]:
        """
        Extract all component relationships from indexed components.

        Args:
            max_depth: Maximum nesting depth to traverse

        Returns:
            List of relationship dictionaries
        """
        all_relationships = []
        processed_classes = set()

        logger.info(
            f"Extracting relationships from {len(self.components)} components..."
        )

        for full_path, component in self.components.items():
            if component.component_type != "class":
                continue

            if component.name in processed_classes:
                continue

            cls = component.obj
            if cls is None or not inspect.isclass(cls):
                continue

            if not self._is_dataclass_type(cls):
                continue

            processed_classes.add(component.name)
            relationships = self._extract_relationships_from_class(
                cls, max_depth=max_depth
            )
            all_relationships.extend(relationships)

        logger.info(f"Extracted {len(all_relationships)} relationships")
        return all_relationships

    def extract_relationships_by_child(
        self, child_class: str, max_depth: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Extract relationships filtered by child class name.

        Args:
            child_class: Name of the child class to filter by
            max_depth: Maximum nesting depth to traverse

        Returns:
            List of relationships where the child matches
        """
        all_rels = self.extract_relationships(max_depth=max_depth)
        filtered = [r for r in all_rels if r["child_class"] == child_class]
        logger.info(f"Found {len(filtered)} {child_class} relationships")
        return filtered

    def extract_relationships_by_parent(
        self, max_depth: int = 5
    ) -> Dict[str, List[Dict[str, Any]]]:
        """
        Extract relationships grouped by parent component name.

        Args:
            max_depth: Maximum nesting depth to traverse

        Returns:
            Dict mapping parent_class name to list of its child relationships
        """
        all_rels = self.extract_relationships(max_depth=max_depth)

        by_parent: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
        for rel in all_rels:
            parent = rel["parent_class"]
            by_parent[parent].append(
                {
                    "child_class": rel["child_class"],
                    "field_name": rel["field_name"],
                    "json_path": rel["json_path"],
                    "depth": rel["depth"],
                    "is_optional": rel["is_optional"],
                    "nl_description": rel["nl_description"],
                    "root_parent": rel.get("root_parent"),
                }
            )

        logger.info(f"Grouped relationships for {len(by_parent)} parent components")
        return dict(by_parent)

    @property
    def relationships(self) -> Dict[str, List[str]]:
        """Get cached relationships grouped by parent component."""
        if self._cached_relationships is None:
            raw_rels = self.extract_relationships_by_parent(max_depth=5)
            self._cached_relationships = {
                parent: [r["child_class"] for r in children]
                for parent, children in raw_rels.items()
            }
        return self._cached_relationships

    @property
    def relationship_metadata(self) -> Dict[str, Any]:
        """Get metadata about relationships."""
        rels = self.relationships
        return {
            "total_parents": len(rels),
            "total_relationships": sum(len(children) for children in rels.values()),
            "parents_with_children": list(rels.keys())[:10],
        }

    def enrich_components_with_relationships(
        self,
        max_depth: int = 5,
        collection_name: Optional[str] = None,
    ) -> int:
        """
        Enrich existing component points with relationship metadata.

        Args:
            max_depth: Maximum nesting depth to traverse
            collection_name: Collection to update

        Returns:
            Number of components enriched
        """
        if collection_name is None:
            collection_name = self.collection_name

        if not self.client:
            logger.warning("No Qdrant client - cannot enrich components")
            return 0

        by_parent = self.extract_relationships_by_parent(max_depth=max_depth)

        if not by_parent:
            logger.info("No relationships to add")
            return 0

        try:
            # Scan collection for class components
            logger.info(
                f"Scanning collection {collection_name} for class components..."
            )
            class_name_to_ids: Dict[str, List[str]] = defaultdict(list)

            offset = None
            total_scanned = 0
            while True:
                scroll_result = self.client.scroll(
                    collection_name=collection_name,
                    limit=100,
                    offset=offset,
                    with_payload=True,
                )
                points, next_offset = scroll_result

                if not points:
                    break

                for point in points:
                    payload = point.payload or {}
                    if payload.get("type") == "class":
                        name = payload.get("name")
                        if name:
                            class_name_to_ids[name].append(str(point.id))
                    total_scanned += 1

                if next_offset is None:
                    break
                offset = next_offset

            logger.info(
                f"Scanned {total_scanned} points, found {len(class_name_to_ids)} unique class names"
            )

            enriched_count = 0

            for parent_class, relationships in by_parent.items():
                point_ids = class_name_to_ids.get(parent_class, [])

                if not point_ids:
                    continue

                nl_descriptions = ", ".join(r["nl_description"] for r in relationships)

                relationships_payload = {
                    "children": relationships,
                    "child_classes": list(set(r["child_class"] for r in relationships)),
                    "max_depth": max(r["depth"] for r in relationships)
                    if relationships
                    else 0,
                    "nl_descriptions": nl_descriptions,
                }

                for point_id in point_ids:
                    try:
                        self.client.set_payload(
                            collection_name=collection_name,
                            payload={"relationships": relationships_payload},
                            points=[point_id],
                        )
                        enriched_count += 1
                    except Exception as e:
                        logger.warning(
                            f"Failed to enrich {parent_class} ({point_id}): {e}"
                        )

            logger.info(
                f"Enriched {enriched_count} components with relationship metadata"
            )
            return enriched_count

        except Exception as e:
            logger.error(f"Failed to enrich components: {e}", exc_info=True)
            return 0

    def get_relationship_tree(self, child_class: str) -> Dict[str, Any]:
        """
        Get a tree of all parents that can contain a child class.

        Args:
            child_class: Name of the child class

        Returns:
            Tree structure showing containment hierarchy
        """
        all_rels = self.extract_relationships()

        # Find all direct parents of this child
        direct_parents = [
            r["parent_class"] for r in all_rels if r["child_class"] == child_class
        ]

        # Build tree recursively
        def build_parent_tree(class_name: str, visited: Set[str]) -> Dict[str, Any]:
            if class_name in visited:
                return {"name": class_name, "cyclic": True}

            visited.add(class_name)

            parents = [
                r["parent_class"] for r in all_rels if r["child_class"] == class_name
            ]

            return {
                "name": class_name,
                "parents": [build_parent_tree(p, visited.copy()) for p in parents],
            }

        return {
            "child": child_class,
            "direct_parents": direct_parents,
            "tree": build_parent_tree(child_class, set()),
        }

    def print_relationship_tree(self, child_class: str) -> str:
        """
        Get a printable tree of parents for a child class.

        Args:
            child_class: Name of the child class

        Returns:
            Formatted tree string
        """
        tree = self.get_relationship_tree(child_class)

        lines = [f"Containment tree for: {child_class}"]
        lines.append(f"Direct parents: {', '.join(tree['direct_parents'])}")
        lines.append("")

        def format_tree(node: Dict, indent: int = 0) -> List[str]:
            result = []
            prefix = "  " * indent
            if node.get("cyclic"):
                result.append(f"{prefix}↺ {node['name']} (cyclic)")
            else:
                result.append(f"{prefix}→ {node['name']}")
                for parent in node.get("parents", []):
                    result.extend(format_tree(parent, indent + 1))
            return result

        lines.extend(format_tree(tree["tree"]))
        return "\n".join(lines)

    def invalidate_caches(self) -> None:
        """Invalidate all cached data including relationships."""
        self._cached_relationships = None
        # Also invalidate symbol caches if they exist
        if hasattr(self, "_symbol_mapping"):
            self._symbol_mapping = None
        if hasattr(self, "_reverse_symbol_mapping"):
            self._reverse_symbol_mapping = None
        if hasattr(self, "_dsl_metadata_cache"):
            self._dsl_metadata_cache = None
        logger.info("Caches invalidated")


# Export for convenience
__all__ = [
    "RelationshipsMixin",
]
