from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime


@dataclass
class Topic:
    """
    Represents a topic cluster of related sources with a lead text.

    A topic groups multiple search results around a central theme,
    with one source selected as the "lead" that best represents the topic.
    """

    id: str
    title: str
    lead_source: Dict[str, Any]  # The main source that represents this topic
    supporting_sources: List[Dict[str, Any]] = field(default_factory=list)
    rejected_sources: List[Dict[str, Any]] = field(default_factory=list)

    # Metadata
    created_at: datetime = field(default_factory=datetime.now)

    # Relationships to other topics
    related_topic_ids: List[str] = field(default_factory=list)
    parent_topic_id: Optional[str] = None
    child_topic_ids: List[str] = field(default_factory=list)

    def add_supporting_source(self, source: Dict[str, Any]) -> None:
        """Add a source that supports this topic."""
        if source not in self.supporting_sources:
            self.supporting_sources.append(source)

    def reject_source(self, source: Dict[str, Any]) -> None:
        """Move a source to the rejected list."""
        if source in self.supporting_sources:
            self.supporting_sources.remove(source)
        if source not in self.rejected_sources:
            self.rejected_sources.append(source)

    def update_lead_source(self, new_lead: Dict[str, Any]) -> None:
        """
        Change the lead source for this topic.
        The old lead becomes a supporting source.
        """
        if self.lead_source:
            self.supporting_sources.append(self.lead_source)
        self.lead_source = new_lead
        # Remove new lead from supporting if it was there
        if new_lead in self.supporting_sources:
            self.supporting_sources.remove(new_lead)

    def get_all_sources(self) -> List[Dict[str, Any]]:
        """Get all sources (lead + supporting) for this topic."""
        return [self.lead_source] + self.supporting_sources

    def to_dict(self) -> Dict[str, Any]:
        """Convert topic to dictionary for serialization."""
        return {
            "id": self.id,
            "title": self.title,
            "lead_source": self.lead_source,
            "supporting_sources": self.supporting_sources,
            "rejected_sources": self.rejected_sources,
            "created_at": self.created_at.isoformat(),
            "related_topic_ids": self.related_topic_ids,
            "parent_topic_id": self.parent_topic_id,
            "child_topic_ids": self.child_topic_ids,
        }


@dataclass
class TopicGraph:
    """
    Manages the collection of topics and their relationships.
    Provides methods for traversing and organizing topics.
    """

    topics: Dict[str, Topic] = field(default_factory=dict)

    def add_topic(self, topic: Topic) -> None:
        """Add a topic to the graph."""
        self.topics[topic.id] = topic

    def get_topic(self, topic_id: str) -> Optional[Topic]:
        """Get a topic by ID."""
        return self.topics.get(topic_id)

    def link_topics(self, topic1_id: str, topic2_id: str) -> None:
        """Create a bidirectional relationship between two topics."""
        topic1 = self.get_topic(topic1_id)
        topic2 = self.get_topic(topic2_id)

        if topic1 and topic2:
            if topic2_id not in topic1.related_topic_ids:
                topic1.related_topic_ids.append(topic2_id)
            if topic1_id not in topic2.related_topic_ids:
                topic2.related_topic_ids.append(topic1_id)

    def set_parent_child(self, parent_id: str, child_id: str) -> None:
        """Set a parent-child relationship between topics."""
        parent = self.get_topic(parent_id)
        child = self.get_topic(child_id)

        if parent and child:
            if child_id not in parent.child_topic_ids:
                parent.child_topic_ids.append(child_id)
            child.parent_topic_id = parent_id

    def get_root_topics(self) -> List[Topic]:
        """Get all topics that have no parent."""
        return [
            topic
            for topic in self.topics.values()
            if topic.parent_topic_id is None
        ]

    def get_related_topics(self, topic_id: str) -> List[Topic]:
        """Get all topics related to a given topic."""
        topic = self.get_topic(topic_id)
        if not topic:
            return []

        related = []
        for related_id in topic.related_topic_ids:
            related_topic = self.get_topic(related_id)
            if related_topic:
                related.append(related_topic)
        return related

    def merge_topics(
        self, topic1_id: str, topic2_id: str, new_title: str = None
    ) -> Optional[Topic]:
        """
        Merge two topics into one, combining their sources.
        Returns the merged topic or None if merge failed.
        """
        topic1 = self.get_topic(topic1_id)
        topic2 = self.get_topic(topic2_id)

        if not topic1 or not topic2:
            return None

        # Use topic1 as the base, update title if provided
        if new_title:
            topic1.title = new_title

        # Add topic2's sources to topic1
        for source in topic2.supporting_sources:
            topic1.add_supporting_source(source)

        for source in topic2.rejected_sources:
            if source not in topic1.rejected_sources:
                topic1.rejected_sources.append(source)

        # Update relationships
        for related_id in topic2.related_topic_ids:
            if (
                related_id != topic1_id
                and related_id not in topic1.related_topic_ids
            ):
                topic1.related_topic_ids.append(related_id)

        # Update child relationships
        for child_id in topic2.child_topic_ids:
            self.set_parent_child(topic1_id, child_id)

        # Remove topic2 from the graph
        del self.topics[topic2_id]

        # Update any references to topic2 in other topics
        for topic in self.topics.values():
            if topic2_id in topic.related_topic_ids:
                topic.related_topic_ids.remove(topic2_id)
                if topic1_id not in topic.related_topic_ids:
                    topic.related_topic_ids.append(topic1_id)

            if topic.parent_topic_id == topic2_id:
                topic.parent_topic_id = topic1_id

        return topic1

    def to_dict(self) -> Dict[str, Any]:
        """Convert the entire graph to a dictionary."""
        return {
            topic_id: topic.to_dict() for topic_id, topic in self.topics.items()
        }
