"""
Knowledge Graph service for Whisper Scribe 3.0.

Builds and maintains a knowledge graph connecting videos, topics, entities,
and concepts for visualization and relationship discovery.
"""

import json
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Dict, Set, Tuple
from dataclasses import dataclass, field
from collections import defaultdict
from ..config import get_config_dir
from .project import load_projects, ProjectData


@dataclass
class GraphNode:
    """A node in the knowledge graph."""
    id: str
    type: str  # "video", "topic", "person", "organization", "technology", "concept"
    label: str
    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'type': self.type,
            'label': self.label,
            'metadata': self.metadata
        }


@dataclass
class GraphEdge:
    """An edge connecting two nodes."""
    source: str  # Node ID
    target: str  # Node ID
    type: str    # "mentions", "related_to", "part_of", "shares_topic"
    weight: float = 1.0
    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            'source': self.source,
            'target': self.target,
            'type': self.type,
            'weight': self.weight,
            'metadata': self.metadata
        }


@dataclass
class KnowledgeGraph:
    """The complete knowledge graph."""
    nodes: List[GraphNode] = field(default_factory=list)
    edges: List[GraphEdge] = field(default_factory=list)
    last_updated: str = ""

    def __post_init__(self):
        if not self.last_updated:
            self.last_updated = datetime.now().isoformat()

    def to_dict(self) -> dict:
        return {
            'nodes': [n.to_dict() for n in self.nodes],
            'edges': [e.to_dict() for e in self.edges],
            'last_updated': self.last_updated
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'KnowledgeGraph':
        return cls(
            nodes=[GraphNode(**n) for n in data.get('nodes', [])],
            edges=[GraphEdge(**e) for e in data.get('edges', [])],
            last_updated=data.get('last_updated', '')
        )

    def add_node(self, node: GraphNode) -> None:
        """Add a node if it doesn't already exist."""
        if not any(n.id == node.id for n in self.nodes):
            self.nodes.append(node)

    def add_edge(self, edge: GraphEdge) -> None:
        """Add an edge, updating weight if it already exists."""
        for e in self.edges:
            if e.source == edge.source and e.target == edge.target and e.type == edge.type:
                e.weight = max(e.weight, edge.weight)
                return
        self.edges.append(edge)

    def get_node(self, node_id: str) -> Optional[GraphNode]:
        """Get a node by ID."""
        for n in self.nodes:
            if n.id == node_id:
                return n
        return None

    def get_connected_nodes(self, node_id: str) -> List[Tuple[GraphNode, GraphEdge]]:
        """Get all nodes connected to a given node."""
        connected = []
        for edge in self.edges:
            if edge.source == node_id:
                node = self.get_node(edge.target)
                if node:
                    connected.append((node, edge))
            elif edge.target == node_id:
                node = self.get_node(edge.source)
                if node:
                    connected.append((node, edge))
        return connected

    def get_nodes_by_type(self, node_type: str) -> List[GraphNode]:
        """Get all nodes of a specific type."""
        return [n for n in self.nodes if n.type == node_type]

    def get_video_connections(self, video_id: str) -> List[dict]:
        """Get all videos connected to a given video with connection details."""
        video_node_id = f"video:{video_id}"
        connections = []
        seen_videos = set()

        for node, edge in self.get_connected_nodes(video_node_id):
            if node.type == "video" and node.id not in seen_videos:
                seen_videos.add(node.id)
                connections.append({
                    'video_id': node.id.replace('video:', ''),
                    'title': node.label,
                    'connection_type': edge.type,
                    'weight': edge.weight
                })
            elif node.type in ("topic", "person", "organization", "technology"):
                # Find other videos connected to this entity
                for sub_node, sub_edge in self.get_connected_nodes(node.id):
                    if sub_node.type == "video" and sub_node.id != video_node_id and sub_node.id not in seen_videos:
                        seen_videos.add(sub_node.id)
                        connections.append({
                            'video_id': sub_node.id.replace('video:', ''),
                            'title': sub_node.label,
                            'via': node.label,
                            'via_type': node.type,
                            'connection_type': 'shared_entity',
                            'weight': min(edge.weight, sub_edge.weight)
                        })

        # Sort by weight descending
        connections.sort(key=lambda x: x.get('weight', 0), reverse=True)
        return connections


# =============================================================================
# Storage Functions
# =============================================================================

def get_graph_file() -> Path:
    """Get the knowledge graph file path."""
    return get_config_dir() / "knowledge_graph.json"


def load_graph() -> KnowledgeGraph:
    """Load the knowledge graph from file."""
    graph_file = get_graph_file()
    if not graph_file.exists():
        return KnowledgeGraph()

    try:
        with open(graph_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
            return KnowledgeGraph.from_dict(data)
    except (json.JSONDecodeError, Exception):
        return KnowledgeGraph()


def save_graph(graph: KnowledgeGraph) -> None:
    """Save the knowledge graph to file."""
    graph.last_updated = datetime.now().isoformat()
    graph_file = get_graph_file()
    with open(graph_file, 'w', encoding='utf-8') as f:
        json.dump(graph.to_dict(), f, indent=2, ensure_ascii=False)


# =============================================================================
# Graph Building Functions
# =============================================================================

def build_graph_from_projects() -> KnowledgeGraph:
    """
    Build the complete knowledge graph from all projects.
    This rebuilds the graph from scratch.
    """
    graph = KnowledgeGraph()
    projects = load_projects()

    topic_videos: Dict[str, Set[str]] = defaultdict(set)  # topic -> set of video_ids
    entity_videos: Dict[str, Set[str]] = defaultdict(set)  # entity_key -> set of video_ids

    for p_data in projects:
        project = ProjectData.from_dict(p_data)

        # Add video node
        video_node_id = f"video:{project.video_id}"
        graph.add_node(GraphNode(
            id=video_node_id,
            type="video",
            label=project.title,
            metadata={
                'channel': project.channel,
                'duration': project.duration,
                'thumbnail': project.thumbnail,
                'url': project.url
            }
        ))

        # Process topics
        all_topics = set(project.shared_topics) | set(project.tags)

        # Also extract topics from analysis if available
        if project.analyses and project.primary_analysis_provider:
            primary_analysis = project.analyses.get(project.primary_analysis_provider, {})
            content = primary_analysis.get('content', {})
            if isinstance(content, dict):
                analysis_topics = content.get('topics', [])
                if isinstance(analysis_topics, list):
                    all_topics.update(analysis_topics)

        for topic in all_topics:
            if topic:
                topic_node_id = f"topic:{topic.lower()}"
                graph.add_node(GraphNode(
                    id=topic_node_id,
                    type="topic",
                    label=topic
                ))
                graph.add_edge(GraphEdge(
                    source=video_node_id,
                    target=topic_node_id,
                    type="has_topic",
                    weight=1.0
                ))
                topic_videos[topic.lower()].add(project.video_id)

        # Process entities
        for entity_list, entity_type in [
            (project.entities.people, "person"),
            (project.entities.organizations, "organization"),
            (project.entities.technologies, "technology"),
            (project.entities.products, "product"),
            (project.entities.places, "place"),
            (project.entities.concepts, "concept")
        ]:
            for entity in entity_list:
                entity_key = f"{entity_type}:{entity.name.lower()}"
                graph.add_node(GraphNode(
                    id=entity_key,
                    type=entity_type,
                    label=entity.name,
                    metadata={
                        'description': entity.description,
                        'wikipedia_url': entity.wikipedia_url
                    }
                ))
                graph.add_edge(GraphEdge(
                    source=video_node_id,
                    target=entity_key,
                    type="mentions",
                    weight=len(entity.mentions) if entity.mentions else 1.0
                ))
                entity_videos[entity_key].add(project.video_id)

    # Add edges between videos that share topics
    for topic, videos in topic_videos.items():
        video_list = list(videos)
        for i, v1 in enumerate(video_list):
            for v2 in video_list[i + 1:]:
                graph.add_edge(GraphEdge(
                    source=f"video:{v1}",
                    target=f"video:{v2}",
                    type="shares_topic",
                    weight=1.0,
                    metadata={'topic': topic}
                ))

    # Add edges between videos that share entities
    for entity_key, videos in entity_videos.items():
        video_list = list(videos)
        for i, v1 in enumerate(video_list):
            for v2 in video_list[i + 1:]:
                graph.add_edge(GraphEdge(
                    source=f"video:{v1}",
                    target=f"video:{v2}",
                    type="shares_entity",
                    weight=1.5,  # Higher weight for shared entities
                    metadata={'entity': entity_key}
                ))

    save_graph(graph)
    return graph


def update_graph_for_project(project: ProjectData) -> KnowledgeGraph:
    """
    Update the knowledge graph for a single project.
    More efficient than rebuilding the entire graph.
    """
    graph = load_graph()

    video_node_id = f"video:{project.video_id}"

    # Remove existing edges for this video
    graph.edges = [e for e in graph.edges if e.source != video_node_id and e.target != video_node_id]

    # Add/update video node
    existing_node = graph.get_node(video_node_id)
    if existing_node:
        existing_node.label = project.title
        existing_node.metadata = {
            'channel': project.channel,
            'duration': project.duration,
            'thumbnail': project.thumbnail,
            'url': project.url
        }
    else:
        graph.add_node(GraphNode(
            id=video_node_id,
            type="video",
            label=project.title,
            metadata={
                'channel': project.channel,
                'duration': project.duration,
                'thumbnail': project.thumbnail,
                'url': project.url
            }
        ))

    # Add topic connections
    all_topics = set(project.shared_topics) | set(project.tags)
    for topic in all_topics:
        if topic:
            topic_node_id = f"topic:{topic.lower()}"
            graph.add_node(GraphNode(
                id=topic_node_id,
                type="topic",
                label=topic
            ))
            graph.add_edge(GraphEdge(
                source=video_node_id,
                target=topic_node_id,
                type="has_topic",
                weight=1.0
            ))

    # Add entity connections
    for entity_list, entity_type in [
        (project.entities.people, "person"),
        (project.entities.organizations, "organization"),
        (project.entities.technologies, "technology"),
        (project.entities.products, "product"),
        (project.entities.places, "place"),
        (project.entities.concepts, "concept")
    ]:
        for entity in entity_list:
            entity_key = f"{entity_type}:{entity.name.lower()}"
            graph.add_node(GraphNode(
                id=entity_key,
                type=entity_type,
                label=entity.name,
                metadata={
                    'description': entity.description,
                    'wikipedia_url': entity.wikipedia_url
                }
            ))
            graph.add_edge(GraphEdge(
                source=video_node_id,
                target=entity_key,
                type="mentions",
                weight=len(entity.mentions) if entity.mentions else 1.0
            ))

    # Find connections to other videos through shared topics/entities
    for node in graph.nodes:
        if node.type == "video" and node.id != video_node_id:
            # Check for shared topics
            other_video_topics = set()
            for edge in graph.edges:
                if edge.source == node.id and edge.type == "has_topic":
                    other_video_topics.add(edge.target)

            this_video_topics = {f"topic:{t.lower()}" for t in all_topics if t}
            shared = other_video_topics & this_video_topics
            if shared:
                graph.add_edge(GraphEdge(
                    source=video_node_id,
                    target=node.id,
                    type="shares_topic",
                    weight=len(shared)
                ))

    save_graph(graph)
    return graph


def remove_video_from_graph(video_id: str) -> KnowledgeGraph:
    """Remove a video and its connections from the graph."""
    graph = load_graph()

    video_node_id = f"video:{video_id}"

    # Remove the video node
    graph.nodes = [n for n in graph.nodes if n.id != video_node_id]

    # Remove edges involving this video
    graph.edges = [e for e in graph.edges if e.source != video_node_id and e.target != video_node_id]

    # Clean up orphaned nodes (topics/entities with no video connections)
    connected_nodes = set()
    for edge in graph.edges:
        connected_nodes.add(edge.source)
        connected_nodes.add(edge.target)

    graph.nodes = [n for n in graph.nodes if n.type == "video" or n.id in connected_nodes]

    save_graph(graph)
    return graph


# =============================================================================
# Query Functions
# =============================================================================

def get_graph_stats() -> dict:
    """Get statistics about the knowledge graph."""
    graph = load_graph()

    node_counts = defaultdict(int)
    for node in graph.nodes:
        node_counts[node.type] += 1

    edge_counts = defaultdict(int)
    for edge in graph.edges:
        edge_counts[edge.type] += 1

    return {
        'total_nodes': len(graph.nodes),
        'total_edges': len(graph.edges),
        'node_counts': dict(node_counts),
        'edge_counts': dict(edge_counts),
        'last_updated': graph.last_updated
    }


def get_topic_clusters() -> List[dict]:
    """Get topics clustered with their related videos and entities."""
    graph = load_graph()

    topics = []
    for node in graph.nodes:
        if node.type == "topic":
            connected = graph.get_connected_nodes(node.id)
            videos = []
            entities = []

            for conn_node, edge in connected:
                if conn_node.type == "video":
                    videos.append({
                        'video_id': conn_node.id.replace('video:', ''),
                        'title': conn_node.label
                    })
                elif conn_node.type in ("person", "organization", "technology"):
                    entities.append({
                        'name': conn_node.label,
                        'type': conn_node.type
                    })

            topics.append({
                'topic': node.label,
                'video_count': len(videos),
                'videos': videos,
                'related_entities': entities
            })

    # Sort by video count descending
    topics.sort(key=lambda x: x['video_count'], reverse=True)
    return topics


def find_related_videos(video_id: str, limit: int = 10) -> List[dict]:
    """Find videos related to a given video through shared topics/entities."""
    graph = load_graph()
    return graph.get_video_connections(video_id)[:limit]


def get_entity_network(entity_type: str = None) -> dict:
    """Get the network of entities, optionally filtered by type."""
    graph = load_graph()

    nodes = []
    edges = []

    for node in graph.nodes:
        if entity_type is None or node.type == entity_type:
            if node.type != "video":
                nodes.append(node.to_dict())

    # Get edges between non-video nodes
    node_ids = {n['id'] for n in nodes}
    for edge in graph.edges:
        if edge.source in node_ids and edge.target in node_ids:
            edges.append(edge.to_dict())

    return {
        'nodes': nodes,
        'edges': edges
    }


def search_graph(query: str) -> List[dict]:
    """Search for nodes matching a query."""
    graph = load_graph()
    query_lower = query.lower()

    results = []
    for node in graph.nodes:
        score = 0
        if query_lower in node.label.lower():
            score = 1.0 if node.label.lower() == query_lower else 0.5
        elif query_lower in str(node.metadata).lower():
            score = 0.3

        if score > 0:
            connected = graph.get_connected_nodes(node.id)
            results.append({
                'node': node.to_dict(),
                'score': score,
                'connections': len(connected)
            })

    results.sort(key=lambda x: (x['score'], x['connections']), reverse=True)
    return results
