"""Diagram Manager for generating and managing Mermaid diagrams."""
import json
from datetime import datetime
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass, asdict
import logging

logger = logging.getLogger(__name__)

@dataclass
class Diagram:
    """Represents a Mermaid diagram with metadata."""
    id: str
    title: str
    description: str
    diagram_type: str  # 'flowchart', 'sequence', 'class', etc.
    content: str
    created_at: str
    updated_at: str
    tags: List[str]
    category: str

class DiagramManager:
    """Manages generation and storage of Mermaid diagrams."""
    
    def __init__(self, context_manager):
        """Initialize the Diagram Manager."""
        self.context_manager = context_manager
        self.diagrams: Dict[str, Diagram] = {}
        self.load_diagram_data()

    def load_diagram_data(self):
        """Load diagram data from context database."""
        context = self.context_manager.get_project_context()
        diagram_data = context.get('diagram_data', {})
        
        for d_dict in diagram_data.get('diagrams', []):
            diagram = Diagram(**d_dict)
            self.diagrams[diagram.id] = diagram

    def save_diagram_data(self):
        """Save diagram data to context database."""
        diagram_data = {
            'diagrams': [asdict(d) for d in self.diagrams.values()]
        }
        
        self.context_manager.update_project_context({
            'diagram_data': diagram_data
        })

    def create_flowchart(self, title: str, nodes: List[Dict[str, Any]],
                        edges: List[Dict[str, Any]], description: str = "",
                        tags: List[str] = None) -> Diagram:
        """Create a flowchart diagram."""
        diagram_id = f"diagram_{len(self.diagrams) + 1}"
        now = datetime.now().isoformat()
        
        # Generate Mermaid flowchart content
        content = ["graph TB"]
        
        # Add nodes
        for node in nodes:
            node_def = f"    {node['id']}[{node['label']}]"
            if node.get('style'):
                node_def += f":::${node['style']}"
            content.append(node_def)
        
        # Add edges
        for edge in edges:
            edge_def = f"    {edge['from']} --> {edge['to']}"
            if edge.get('label'):
                edge_def = f"    {edge['from']} -- {edge['label']} --> {edge['to']}"
            content.append(edge_def)
        
        diagram = Diagram(
            id=diagram_id,
            title=title,
            description=description,
            diagram_type='flowchart',
            content="\n".join(content),
            created_at=now,
            updated_at=now,
            tags=tags or [],
            category='system'
        )
        
        self.diagrams[diagram_id] = diagram
        self.save_diagram_data()
        
        return diagram

    def create_sequence_diagram(self, title: str, participants: List[str],
                              sequences: List[Dict[str, Any]], description: str = "",
                              tags: List[str] = None) -> Diagram:
        """Create a sequence diagram."""
        diagram_id = f"diagram_{len(self.diagrams) + 1}"
        now = datetime.now().isoformat()
        
        # Generate Mermaid sequence diagram content
        content = ["sequenceDiagram"]
        
        # Add participants
        for participant in participants:
            content.append(f"    participant {participant}")
        
        # Add sequences
        for seq in sequences:
            seq_def = f"    {seq['from']} ->> {seq['to']}: {seq['message']}"
            if seq.get('response'):
                seq_def += f"\n    {seq['to']} -->> {seq['from']}: {seq['response']}"
            content.append(seq_def)
        
        diagram = Diagram(
            id=diagram_id,
            title=title,
            description=description,
            diagram_type='sequence',
            content="\n".join(content),
            created_at=now,
            updated_at=now,
            tags=tags or [],
            category='interaction'
        )
        
        self.diagrams[diagram_id] = diagram
        self.save_diagram_data()
        
        return diagram

    def create_class_diagram(self, title: str, classes: List[Dict[str, Any]],
                           relationships: List[Dict[str, Any]], description: str = "",
                           tags: List[str] = None) -> Diagram:
        """Create a class diagram."""
        diagram_id = f"diagram_{len(self.diagrams) + 1}"
        now = datetime.now().isoformat()
        
        # Generate Mermaid class diagram content
        content = ["classDiagram"]
        
        # Add classes
        for cls in classes:
            class_def = [f"    class {cls['name']} {{"]
            
            # Add attributes
            for attr in cls.get('attributes', []):
                class_def.append(f"        {attr['visibility']}{attr['name']}: {attr['type']}")
            
            # Add methods
            for method in cls.get('methods', []):
                method_def = f"        {method['visibility']}{method['name']}("
                if method.get('params'):
                    method_def += ", ".join(f"{p['name']}: {p['type']}" for p in method['params'])
                method_def += f") {method.get('return_type', 'void')}"
                class_def.append(method_def)
            
            class_def.append("    }")
            content.extend(class_def)
        
        # Add relationships
        for rel in relationships:
            rel_def = f"    {rel['from']} {rel['type']} {rel['to']}"
            if rel.get('label'):
                rel_def += f" : {rel['label']}"
            content.append(rel_def)
        
        diagram = Diagram(
            id=diagram_id,
            title=title,
            description=description,
            diagram_type='class',
            content="\n".join(content),
            created_at=now,
            updated_at=now,
            tags=tags or [],
            category='structure'
        )
        
        self.diagrams[diagram_id] = diagram
        self.save_diagram_data()
        
        return diagram

    def get_diagram(self, diagram_id: str) -> Optional[Diagram]:
        """Get a specific diagram by ID."""
        return self.diagrams.get(diagram_id)

    def search_diagrams(self, query: str) -> List[Diagram]:
        """Search diagrams by title, description, or tags."""
        query = query.lower()
        results = []
        
        for diagram in self.diagrams.values():
            if (query in diagram.title.lower() or
                query in diagram.description.lower() or
                any(query in tag.lower() for tag in diagram.tags)):
                results.append(diagram)
        
        return results

    def update_diagram(self, diagram_id: str, updates: Dict[str, Any]) -> Diagram:
        """Update an existing diagram."""
        if diagram_id not in self.diagrams:
            raise ValueError(f"Diagram {diagram_id} not found")
            
        diagram = self.diagrams[diagram_id]
        diagram_dict = asdict(diagram)
        
        for key, value in updates.items():
            if key in diagram_dict:
                diagram_dict[key] = value
        
        diagram_dict['updated_at'] = datetime.now().isoformat()
        self.diagrams[diagram_id] = Diagram(**diagram_dict)
        self.save_diagram_data()
        
        return self.diagrams[diagram_id]

    def delete_diagram(self, diagram_id: str) -> bool:
        """Delete a diagram."""
        if diagram_id in self.diagrams:
            del self.diagrams[diagram_id]
            self.save_diagram_data()
            return True
        return False

    def export_diagram(self, diagram_id: str, format: str = 'mermaid') -> str:
        """Export a diagram in the specified format."""
        if diagram_id not in self.diagrams:
            raise ValueError(f"Diagram {diagram_id} not found")
            
        diagram = self.diagrams[diagram_id]
        
        if format == 'mermaid':
            return diagram.content
        elif format == 'json':
            return json.dumps(asdict(diagram), indent=2)
        else:
            raise ValueError(f"Unsupported format: {format}") 