import json
import httpx
import re
import logging
from uuid import uuid4
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.utils import timezone
from django.conf import settings

from apps.onboarding.mcp import MCPRouter
from apps.onboarding.models import AgentConfig, OnboardingFlow, OnboardingSession

logger = logging.getLogger(__name__)

class OnboardingConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.user = self.scope["user"]
                                              
        self.context_uuid = self.scope["url_route"]["kwargs"].get("session_uuid")
        
        if not self.user.is_authenticated:
            await self.close()
            return

        self.router = MCPRouter()
        await self.accept()

    async def disconnect(self, close_code):
        pass

    def _build_system_prompt(self, config):
        base_prompt = config.system_prompt or "You are a helpful onboarding assistant."
        permissions = config.tool_permissions or []
        if permissions:
            return f"{base_prompt}\n\nAllowed tools: {', '.join(str(p) for p in permissions)}"
        return base_prompt

    async def receive(self, text_data):
        try:
            data = json.loads(text_data)
            action = data.get("action")

            if action == "start_full_onboarding":
                role_uuid = data.get("role_uuid")
                if not role_uuid:
                    await self.send_log("error", "Missing role_uuid for full onboarding generation")
                    return
                await self.run_full_onboarding_generation(role_uuid)
            elif action == "progress_monitor":
                role_uuid = data.get("role_uuid") or self.context_uuid
                if not role_uuid:
                    await self.send_log("error", "Missing role_uuid for progress monitoring")
                    return
                await self.run_progress_monitor(role_uuid)
            else:
                                                                           
                user_message = data.get("query") or data.get("message")
                if not user_message:
                    await self.send_log("error", "Missing query/message payload")
                    return
                config = await self.get_config(self.context_uuid)
                ai_response = await self.orchestrate_ai(user_message, config)

                await self.send(json.dumps({
                    "type": "completed",
                    "timestamp": timezone.now().isoformat(),
                    "message": "Inference complete.",
                    "content": {
                        "response": ai_response,
                    }
                }))
        except Exception as e:
            logger.error(f"WebSocket Receive Error: {str(e)}")
            await self.send_log("error", f"Consumer encountered an error: {str(e)}")

    async def run_full_onboarding_generation(self, role_uuid):
        """
        The Master Script that builds the JSON structure sequentially.
        Pipeline: Curriculum Agent -> Knowledge Agent -> Assessment Agent
        """
        
        await self.send_log("status", "Phase 1: Generating Curriculum...", "curriculum")
        ca_config = await self.get_config_by_type(role_uuid, 'curriculum')
        if not ca_config:
            await self.send_log("error", "Missing curriculum AgentConfig for this role")
            return
        
        ca_prompt = (
            "Based on available documentation, create an onboarding curriculum for this role. "
            "Output ONLY a valid JSON array of 3-5 strings representing module titles. "
            "Example: [\"Introduction\", \"Safety\", \"Operations\"]"
        )
        ca_response = await self.orchestrate_ai(ca_prompt, ca_config)
        topics = self._extract_json_list(ca_response)
        if not topics:
            await self.send_log("error", "Curriculum generation returned no topics")
            return

        toc_lines = [f"{idx + 1}. {title}" for idx, title in enumerate(topics)]
        toc_markdown = "## Table of Contents\n\n" + "\n".join(toc_lines)

        full_structure = []

                                                
        for index, topic in enumerate(topics):
                                   
            await self.send_log("status", f"Phase 2: Researching {topic}...", "knowledge")
            ka_config = await self.get_config_by_type(role_uuid, 'knowledge')
            if not ka_config:
                await self.send_log("error", "Missing knowledge AgentConfig for this role")
                return

            knowledge_hits = await self.fetch_knowledge_context(role_uuid, topic)
            context_markdown = self.format_knowledge_context(knowledge_hits)

            page_content = await self.orchestrate_ai(
                (
                    f"Write a practical onboarding training guide for the topic '{topic}'. "
                    "Use the MCP search context provided below as the primary source. "
                    "If the context is empty, provide a concise best-practice overview and clearly say no indexed documents were found. "
                    "Use Markdown formatting and do NOT include a table of contents in this section.\n\n"
                    f"Role UUID: {role_uuid}\n"
                    f"MCP search context:\n{context_markdown}"
                ),
                ka_config
            )

            if index == 0:
                page_content = f"{toc_markdown}\n\n---\n\n{page_content}"

                                    
            await self.send_log("status", f"Phase 3: Creating quiz for {topic}...", "assessment")
            aa_config = await self.get_config_by_type(role_uuid, 'assessment')
            if not aa_config:
                await self.send_log("error", "Missing assessment AgentConfig for this role")
                return
            aa_prompt = (
                f"Based on this content: '{page_content[:1000]}', create 2 multiple choice questions. "
                "Output ONLY a JSON array of objects with keys: 'key', 'label', 'field_type' (use 'select'), "
                "'options' (array of strings), and 'required' (true)."
            )
            quiz_response = await self.orchestrate_ai(aa_prompt, aa_config)
            quiz_fields = self._extract_json_list(quiz_response)

                                     
            full_structure.append({
                "title": topic,
                "body": page_content,
                "order": index,
                "fields": quiz_fields
            })

                               
        await self.save_full_flow(role_uuid, full_structure)

                                  
        await self.send(json.dumps({
            "type": "completed",
            "timestamp": timezone.now().isoformat(),
            "message": "Onboarding pipeline complete and structure saved."
        }))

    async def run_progress_monitor(self, role_uuid):
        await self.send_log("status", "Progress Monitor is analyzing your onboarding progress...", "monitor")

        monitor_config = await self.get_config_by_type(role_uuid, 'monitor')
        if not monitor_config:
            await self.send_log("error", "Missing Progress Monitor AgentConfig for this role")
            return

        progress_context = await self.get_role_progress_context(role_uuid, self.user.id)

        monitor_prompt = (
            "You are a progress monitoring agent for onboarding. "
            "Analyze the role onboarding data below and provide concise feedback with:\n"
            "1) current status\n2) strengths\n3) gaps\n4) next actions\n"
            "Keep it short and practical.\n\n"
            f"Progress context JSON:\n{json.dumps(progress_context)}"
        )

        feedback = await self.orchestrate_ai(monitor_prompt, monitor_config)

        await self.send(json.dumps({
            "type": "completed",
            "timestamp": timezone.now().isoformat(),
            "message": "Progress analysis complete.",
            "content": {
                "role_uuid": role_uuid,
                "feedback": feedback,
                "status": progress_context.get("latest_status", "unknown"),
            }
        }))

    async def orchestrate_ai(self, user_message, config):
        """
        Handles the multi-turn ReAct loop (Reasoning + Tool Use).
        """
        messages = [
            {"role": "system", "content": self._build_system_prompt(config)},
            {"role": "user", "content": user_message}
        ]

        async with httpx.AsyncClient(timeout=60.0) as client:
            for turn in range(5):              
                await self.send_log("thought", f"Agent is thinking (Turn {turn+1})...")
                
                try:
                    response = await client.post(
                        f"{settings.INFERENCE_URL}/v1/chat/completions",
                        json={
                            "model": config.llm_config.get("model_id", "meta-llama-3.1-8b"),
                            "messages": messages,
                            "tools": self.router.get_tool_definitions(),
                            "tool_choice": "auto"
                        }
                    )
                    response.raise_for_status()
                    res_json = response.json()
                    
                    ai_message = res_json["choices"][0]["message"]
                    messages.append(ai_message)

                    if ai_message.get("tool_calls"):
                        for tool_call in ai_message["tool_calls"]:
                            fn_name = tool_call["function"]["name"]
                            fn_args = json.loads(tool_call["function"]["arguments"])
                            
                            await self.send(json.dumps({
                                "type": "tool_start",
                                "message": f"Accessing knowledge base: {fn_name}...",
                                "content": fn_args
                            }))
                            
                                          
                            result = await self.router.handle_tool_call(fn_name, fn_args)
                            
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call["id"],
                                "name": fn_name,
                                "content": json.dumps(result)
                            })
                        continue 
                    
                    else:
                        return ai_message["content"]
                        
                except Exception as e:
                    await self.send_log("error", f"Inference failed: {str(e)}")
                    return f"Error: {str(e)}"

                                

    async def fetch_knowledge_context(self, role_uuid, topic):
        query = f"onboarding training content for {topic}"
        await self.send(json.dumps({
            "type": "tool_start",
            "message": "Accessing knowledge base: search_knowledge...",
            "content": {"query": query, "role_uuid": role_uuid}
        }))

        try:
            result = await self.router.handle_tool_call(
                "search_knowledge",
                {
                    "query": query,
                    "role_uuid": role_uuid,
                },
            )

            await self.send(json.dumps({
                "type": "tool_result",
                "message": f"Retrieved {len(result) if isinstance(result, list) else 0} knowledge chunk(s)",
                "content": result,
                "timestamp": timezone.now().isoformat(),
            }))

            return result if isinstance(result, list) else []
        except Exception as exc:
            await self.send_log("error", f"Knowledge retrieval failed for topic '{topic}': {str(exc)}")
            return []

    def format_knowledge_context(self, knowledge_hits):
        if not knowledge_hits:
            return "No indexed MCP documents found for this role/topic."

        lines = []
        for idx, item in enumerate(knowledge_hits[:5]):
            source = item.get("source", "Unknown Source") if isinstance(item, dict) else "Unknown Source"
            relevance = item.get("relevance") if isinstance(item, dict) else None
            content = item.get("content", "") if isinstance(item, dict) else ""
            safe_content = str(content).strip()[:1600]
            lines.append(
                f"[{idx + 1}] Source: {source} | Relevance: {relevance}\n{safe_content}"
            )

        return "\n\n".join(lines)

    def _extract_json_list(self, text):
        """Regex helper to pull JSON out of LLM conversational filler."""
        try:
            if not text:
                return []
            match = re.search(r'\[.*\]', text, re.DOTALL)
            if match:
                return json.loads(match.group())
            return []
        except Exception:
            return []

    def _normalize_structure(self, structure):
        normalized_pages = []
        for index, page in enumerate(structure or []):
            fields = []
            for field_index, field in enumerate(page.get('fields', []) if isinstance(page, dict) else []):
                if not isinstance(field, dict):
                    continue
                key = str(field.get('key') or f'field_{field_index + 1}')
                fields.append({
                    'uuid': str(uuid4()),
                    'key': key,
                    'label': str(field.get('label') or key.replace('_', ' ').title()),
                    'field_type': str(field.get('field_type') or 'text'),
                    'required': bool(field.get('required', False)),
                    'options': field.get('options') if isinstance(field.get('options'), list) else [],
                    'default_value': field.get('default_value', ''),
                })

            page_title = page.get('title') if isinstance(page, dict) else None
            page_body = page.get('body') if isinstance(page, dict) else ''
            page_order = page.get('order') if isinstance(page, dict) else index
            normalized_pages.append({
                'uuid': str(uuid4()),
                'title': str(page_title or f'Module {index + 1}'),
                'body': str(page_body or ''),
                'order': int(page_order if isinstance(page_order, int) else index),
                'fields': fields,
            })
        return normalized_pages

    @database_sync_to_async
    def save_full_flow(self, role_uuid, structure):
        """Saves the final nested structure to the OnboardingFlow model."""
        from apps.accounts.models import Role
        role = Role.objects.get(uuid=role_uuid)
        normalized_structure = self._normalize_structure(structure)
        flow, _ = OnboardingFlow.objects.update_or_create(
            role=role,
            defaults={
                'title': f"AI Onboarding: {role.name}",
                'structure': normalized_structure,
                'is_active': True
            }
        )
        return flow

    async def send_log(self, log_type, message, content=None):
        await self.send(json.dumps({
            "type": log_type,
            "message": message,
            "content": content,
            "timestamp": timezone.now().isoformat()
        }))

    @database_sync_to_async
    def get_config(self, config_uuid):
        return AgentConfig.objects.get(uuid=config_uuid)

    @database_sync_to_async
    def get_config_by_type(self, role_uuid, agent_type):
        return AgentConfig.objects.filter(
            organization__roles__uuid=role_uuid,
            agent_type=agent_type,
        ).order_by('-updated_at').first()

    @database_sync_to_async
    def get_role_progress_context(self, role_uuid, user_id):
        from apps.accounts.models import Role

        role = Role.objects.get(uuid=role_uuid)
        sessions = OnboardingSession.objects.filter(user_id=user_id, role=role).order_by('-updated_at')
        latest_session = sessions.first()
        active_flow = OnboardingFlow.objects.filter(role=role, is_active=True).order_by('-updated_at').first()

        if not latest_session:
            return {
                "role_uuid": str(role.uuid),
                "role_name": role.name,
                "latest_status": "not_started",
                "session_count": 0,
                "flow_exists": bool(active_flow),
                "progress": 0,
                "responses_count": 0,
                "completed_modules": [],
            }

        state = latest_session.state or {}
        responses = state.get("responses", {})
        completed_modules = state.get("completed_modules", [])
        progress = state.get("progress_percentage", state.get("progress", 0))

        return {
            "role_uuid": str(role.uuid),
            "role_name": role.name,
            "latest_status": latest_session.status,
            "session_count": sessions.count(),
            "flow_exists": bool(active_flow),
            "progress": progress,
            "responses_count": len(responses) if isinstance(responses, dict) else 0,
            "completed_modules": completed_modules if isinstance(completed_modules, list) else [],
            "updated_at": latest_session.updated_at.isoformat() if latest_session.updated_at else None,
        }