from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from django.contrib.auth.decorators import login_required
from .models import Project, MindMapNode, MindMapConnection
from ai.local_ai_api import LocalAIApi
from django.shortcuts import get_object_or_404
import json

@csrf_exempt
@login_required
def ai_chat(request, pk):
    if request.method == 'POST':
        project = get_object_or_404(Project, pk=pk, user=request.user)
        
        try:
            data = json.loads(request.body)
            user_message = data.get('message')
            
            # Context builder for the project
            nodes = list(project.nodes.values('id', 'title', 'category'))
            nodes_context = json.dumps(nodes)
            
            system_prompt = f"""You are an AI business strategist helping a user build a mind map for their project.
Project Title: {project.title}
Industry: {project.industry}
Goal: {project.goal}

Current Nodes (with IDs):
{nodes_context}

Respond ONLY with a valid JSON object in the following format:
{{
    "message": "Your text response to the user's prompt goes here.",
    "new_nodes": [
        {{"id": "temp_1", "title": "New Idea", "summary": "Short description", "category": "Strategy"}}
    ],
    "new_connections": [
        {{"source_id": 1, "target_id": "temp_1", "how": "Relates to existing node", "why": "Important reason"}}
    ]
}}
Instructions:
- If the user just asks a question, put your answer in "message" and leave "new_nodes" and "new_connections" empty.
- If the user asks to add nodes, brainstorm nodes, or expand the map, generate them and place them in "new_nodes". Use temporary string IDs for new nodes (e.g., "temp_1").
- For connections, "source_id" and "target_id" can be an existing integer ID from the Current Nodes list OR a temporary string ID from "new_nodes".
- Keep your answers concise and practical.
"""

            response = LocalAIApi.create_response({
                "input": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message}
                ],
                "text": {"format": {"type": "json_object"}}
            })

            if response.get("success"):
                ai_data = LocalAIApi.decode_json_from_response(response)
                
                if not ai_data:
                     # fallback if it's not valid JSON
                     ai_text = LocalAIApi.extract_text(response)
                     return JsonResponse({'response': ai_text})

                # Extract parts
                ai_message = ai_data.get("message", "I have updated the mind map.")
                new_nodes_data = ai_data.get("new_nodes", [])
                new_connections_data = ai_data.get("new_connections", [])
                
                added_nodes = []
                added_connections = []
                
                if new_nodes_data or new_connections_data:
                    # Save nodes
                    temp_to_real_id = {}
                    
                    for n in new_nodes_data:
                        node = MindMapNode.objects.create(
                            project=project,
                            title=n.get("title", "Untitled"),
                            summary=n.get("summary", ""),
                            category=n.get("category", "General")
                        )
                        temp_to_real_id[n.get("id")] = node
                        added_nodes.append({
                            "id": node.pk,
                            "title": node.title,
                            "summary": node.summary,
                            "category": node.category
                        })
                    
                    # Create dictionary to fetch existing nodes quickly by int ID
                    existing_nodes_map = {n.id: n for n in project.nodes.all()}
                    
                    # Save connections
                    for c in new_connections_data:
                        source_ref = c.get("source_id")
                        target_ref = c.get("target_id")
                        
                        source_node = None
                        target_node = None
                        
                        # Resolve source
                        if isinstance(source_ref, int) and source_ref in existing_nodes_map:
                            source_node = existing_nodes_map[source_ref]
                        elif isinstance(source_ref, str) and str(source_ref).isdigit() and int(source_ref) in existing_nodes_map:
                            source_node = existing_nodes_map[int(source_ref)]
                        elif source_ref in temp_to_real_id:
                            source_node = temp_to_real_id[source_ref]
                            
                        # Resolve target
                        if isinstance(target_ref, int) and target_ref in existing_nodes_map:
                            target_node = existing_nodes_map[target_ref]
                        elif isinstance(target_ref, str) and str(target_ref).isdigit() and int(target_ref) in existing_nodes_map:
                            target_node = existing_nodes_map[int(target_ref)]
                        elif target_ref in temp_to_real_id:
                            target_node = temp_to_real_id[target_ref]
                            
                        if source_node and target_node:
                            conn = MindMapConnection.objects.create(
                                project=project,
                                source=source_node,
                                target=target_node,
                                how=c.get("how", ""),
                                why=c.get("why", "")
                            )
                            added_connections.append({
                                "id": conn.pk,
                                "source_id": source_node.pk,
                                "target_id": target_node.pk,
                                "how": conn.how,
                                "why": conn.why
                            })

                return JsonResponse({
                    'response': ai_message,
                    'added_nodes': added_nodes,
                    'added_connections': added_connections
                })
            else:
                return JsonResponse({'response': "Sorry, I had trouble processing that request. Please try again."})
                
        except json.JSONDecodeError:
            return JsonResponse({'error': 'Invalid request body'}, status=400)
            
    return JsonResponse({'error': 'Invalid request method'}, status=405)