"""
Concept Graph Generator
Generates canonical concept dependency graphs for given concepts
"""

from typing import Dict, List
import os
import requests
import json

class ConceptGraphGenerator:
    def __init__(self):
        self.hf_api_key = os.getenv('HUGGINGFACE_API_KEY')
        self.llm_endpoint = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
        self._ready = True
    
    def is_ready(self) -> bool:
        return self._ready
    
    async def generate_graph(self, concept: str) -> Dict:
        """
        Generate canonical concept dependency graph
        
        Returns:
            {
                'nodes': [{'id': str, 'label': str, 'level': int}],
                'edges': [{'source': str, 'target': str, 'relationship': str}]
            }
        """
        # Use LLM to generate concept structure
        graph_structure = await self._llm_generate_structure(concept)
        
        # Validate and format
        return self._format_graph(graph_structure, concept)
    
    async def _llm_generate_structure(self, concept: str) -> Dict:
        """Use LLM to generate concept prerequisite structure"""
        prompt = f"""<s>[INST] You are a concept structure expert. For the concept "{concept}", identify the core prerequisite concepts that must be understood first, and their relationships.

Output a JSON structure with:
1. "prerequisites": list of prerequisite concepts needed to understand {concept}
2. "core_components": main parts/aspects of {concept} itself
3. "relationships": how concepts connect (prerequisite, causal, etc.)

Be precise and pedagogical. Focus on understanding order.

Output only valid JSON, no other text. [/INST]"""

        try:
            headers = {"Authorization": f"Bearer {self.hf_api_key}"}
            payload = {
                "inputs": prompt,
                "parameters": {
                    "max_new_tokens": 800,
                    "temperature": 0.4,
                    "return_full_text": False
                }
            }
            
            response = requests.post(self.llm_endpoint, headers=headers, json=payload, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                text = result[0]['generated_text'] if isinstance(result, list) else result.get('generated_text', '')
                
                # Try to parse JSON from response
                try:
                    # Extract JSON if wrapped in other text
                    start = text.find('{')
                    end = text.rfind('}') + 1
                    if start != -1 and end > start:
                        json_str = text[start:end]
                        return json.loads(json_str)
                except:
                    pass
                
                return self._fallback_graph(concept)
            else:
                return self._fallback_graph(concept)
                
        except Exception as e:
            print(f"Graph generation error: {e}")
            return self._fallback_graph(concept)
    
    def _fallback_graph(self, concept: str) -> Dict:
        """Fallback: create a basic graph structure"""
        # Predefined templates for common concepts
        templates = {
            'entropy': {
                'prerequisites': ['energy', 'system states', 'probability'],
                'core_components': ['disorder measure', 'thermodynamic entropy', 'information entropy'],
                'relationships': [
                    ('energy', 'entropy', 'prerequisite'),
                    ('system states', 'entropy', 'prerequisite'),
                    ('probability', 'entropy', 'prerequisite')
                ]
            },
            'neural networks': {
                'prerequisites': ['linear algebra', 'calculus', 'probability'],
                'core_components': ['neurons', 'layers', 'weights', 'activation functions', 'backpropagation'],
                'relationships': [
                    ('linear algebra', 'neural networks', 'prerequisite'),
                    ('neurons', 'layers', 'component'),
                    ('weights', 'neurons', 'component'),
                    ('backpropagation', 'weights', 'causal')
                ]
            },
            'photosynthesis': {
                'prerequisites': ['energy', 'chemical reactions', 'cells'],
                'core_components': ['light reactions', 'dark reactions', 'chlorophyll', 'glucose production'],
                'relationships': [
                    ('energy', 'light reactions', 'prerequisite'),
                    ('light reactions', 'dark reactions', 'causal'),
                    ('dark reactions', 'glucose production', 'causal')
                ]
            }
        }
        
        # Check if concept matches template
        concept_lower = concept.lower()
        for key, template in templates.items():
            if key in concept_lower:
                return template
        
        # Generic fallback
        return {
            'prerequisites': ['foundational knowledge'],
            'core_components': [concept, f'{concept} principles', f'{concept} applications'],
            'relationships': [
                ('foundational knowledge', concept, 'prerequisite')
            ]
        }
    
    def _format_graph(self, structure: Dict, concept: str) -> Dict:
        """Format graph structure for frontend"""
        nodes = []
        edges = []
        node_id = 0
        node_map = {}
        
        # Add prerequisite nodes
        for prereq in structure.get('prerequisites', []):
            node_map[prereq] = f'node_{node_id}'
            nodes.append({
                'id': f'node_{node_id}',
                'label': prereq,
                'level': 0,
                'type': 'prerequisite'
            })
            node_id += 1
        
        # Add main concept node
        node_map[concept] = f'node_{node_id}'
        nodes.append({
            'id': f'node_{node_id}',
            'label': concept,
            'level': 1,
            'type': 'main'
        })
        concept_node_id = f'node_{node_id}'
        node_id += 1
        
        # Add core component nodes
        for component in structure.get('core_components', []):
            node_map[component] = f'node_{node_id}'
            nodes.append({
                'id': f'node_{node_id}',
                'label': component,
                'level': 2,
                'type': 'component'
            })
            node_id += 1
        
        # Add edges from relationships
        for rel in structure.get('relationships', []):
            if len(rel) >= 3:
                source_key, target_key, rel_type = rel[0], rel[1], rel[2]
                source_id = node_map.get(source_key, node_map.get(concept))
                target_id = node_map.get(target_key, concept_node_id)
                
                edges.append({
                    'source': source_id,
                    'target': target_id,
                    'relationship': rel_type
                })
        
        # Add default prerequisite edges if none exist
        if not edges:
            for prereq in structure.get('prerequisites', []):
                edges.append({
                    'source': node_map[prereq],
                    'target': concept_node_id,
                    'relationship': 'prerequisite'
                })
        
        return {
            'nodes': nodes,
            'edges': edges,
            'concept': concept
        }