import topic_taxonomy
import micro_prompts
import validators
import asyncio
import json
import cost_tracker
import re
import logging
from utils.safe_json import safe_extract_json  # V1.0: Canonical extractor (replaces local duplicate)
import prompts

logger = logging.getLogger(__name__)


class StrategyManager:
    def __init__(self, llm_model):
        self.llm = llm_model
        self.max_retries = 2
    
    async def solve_with_strategy(self, problem_text, data_anchor, grade="10", image_data=None, image_data_list=None, parent_category=None, ambiguity_warning=False, intent=None, intent_contract=None, proof_graph_steps_count=1, student_name="נסיך", student_gender="M"):
        """
        V3.1.2: Added support for adaptive failure (Soft Recovery).
        V5.8.2: Dynamic Token Budget added via proof_graph_steps_count.
        """
        self.grade = grade # V4.2.11: Persist grade for prompt building
        topic_id = topic_taxonomy.detect_topic(problem_text, grade)
        detected_category = topic_taxonomy.get_topic_category(topic_id)
        category = detected_category if detected_category != "GENERAL" else (parent_category or "GENERAL")
        
        image_pages = []
        if image_data_list:
            # Multi-image support
            for i_data in image_data_list:
                image_pages.append({"mime_type": "image/jpeg", "data": i_data})
        elif image_data:
            try:
                from ocr_strip_engine import paginate_image
                image_pages = paginate_image(image_data)
            except Exception as e: 
                logger.exception("CRITICAL FLOW ERROR")
                logger.error(f"Pagination failed: {e}")
        
        # V4.3: Single Pass (No Retries)
        try:
            strategy_config = self._get_strategy(topic_id)
            prompt = self._build_prompt(topic_id, data_anchor, strategy_config, problem_text, category, grade, student_name, student_gender)
            
            # V4.2 Intent Lockdown (Iron Law #2)
            if intent_contract and intent_contract.get("status") != "unconstrained":
                lockdown_note = f"""
📢 [V4.2] PEDAGOGICAL LOCKDOWN:
- Intent: {intent}
- Max Variables: {intent_contract.get('max_variables', 'N/A')}
- Forbidden: {", ".join(intent_contract.get('forbidden_strategies', []))}
- REQUIRED: Use ONLY basic algebraic steps. DO NOT use advanced functions or calculus.
"""
                prompt = lockdown_note + "\n" + prompt
            # V281.1: HARD OCR DISCONNECT
            # If we have image data, DO NOT inject the OCR-based function equation.
            # The LLM must rely on the Vision-First extracted Data Anchor and the image itself.
            if data_anchor.get("function_equations") and not image_data and not image_data_list:
                prompt = f"TARGET FUNCTION: {data_anchor['function_equations'][0]}\n\n" + prompt
            
            # V3.1.2: Soft Recovery Prompt Injection
            if ambiguity_warning:
                soft_recovery_note = """
System Prompt Override: הוסף לתחילת ההסבר שלך את ההערה הבאה בנימה ידידותית: 'היי! הצילום היה קצת לא ברור, אבל נראה לי שהתכוונת לביטוי הזה. אם התכוונת למשהו אחר, פשוט צלם שוב מקרוב!' - המשך להסביר את השלבים כרגיל.
"""
                prompt = soft_recovery_note + "\n" + prompt

            prompt = (prompt or "") + f"\n\n🎯 MISSION: Solve ONLY part: {problem_text}"
            llm_response = await self._call_llm(prompt, image_data, category, image_pages, proof_graph_steps_count)
            
            # V5.8.2: Guard against raw list responses
            if isinstance(llm_response, list):
                llm_response = {"steps": llm_response, "is_valid": True, "confidence_score": 1.0}
            
            return {
                "topic": topic_id,
                "category": category,
                "llm_response": llm_response,
                "validated": llm_response.get("is_valid", True),
                "confidence": llm_response.get("confidence_score", 1.0)
            }
        except asyncio.CancelledError:
            logger.warning("V4.3 Single Pass Cancelled (Client Disconnected)")
            return {
                "topic": topic_id,
                "category": category,
                "llm_response": {"solution_markdown": "בוטל על ידי המשתמש", "is_valid": False, "confidence_score": 0.0},
                "validated": False
            }
        except Exception as e:
            logger.error(f"V4.3 Single Pass Failed: {e}")
            return {
                "topic": topic_id,
                "category": category,
                "llm_response": {"solution_markdown": "המורה נתקלה בקושי טכני.", "is_valid": False, "confidence_score": 0.0},
                "validated": False
            }

    def _get_strategy(self, topic_id):
        return {"has_micro_prompt": topic_id in micro_prompts.MICRO_PROMPTS, "has_validator": topic_id in validators.VALIDATORS}

    def _build_prompt(self, topic_id, data_anchor, strategy_config, problem_text, category, grade, student_name, student_gender):
        # V7.3 Hybrid Mode: Always use Specialist Prompt for the "Solution Skin"
        specialist = prompts.get_specialist_prompt(
            category=category,
            problem_text=problem_text,
            solver_hint="Hybrid Navigation Mode",
            grade=grade,
            student_name=student_name,
            student_gender=student_gender,
            data_anchor=data_anchor
        )

        if strategy_config["has_micro_prompt"]:
            try: 
                 focus = micro_prompts.get_micro_prompt(topic_id, data_anchor, grade=self.grade)
                 return f"{focus}\n\n{specialist}"
            except: 
                 pass
        
        return specialist
        return micro_prompts.get_general_prompt(data_anchor)

    async def _call_llm(self, prompt, image_data, category, image_pages, proof_graph_steps_count=1):
        # V8.5: No More Patchwork. Use centralized V4.3.0 standard.
        v430_instruction = prompts.get_master_prompt_v430(category=category, problem_text=prompt)
        prompt = (prompt or "") + (v430_instruction or "")
        
        from google.generativeai.types import GenerationConfig
        import asyncio
        
        # V5.8.3: Pre-release Hardening (Increased Budget for Preamble)
        # V8.6: Golden Merge - Increased to 8192 to support high-density Micro-Stepping
        max_tokens = 8192
        logger.info(f"🪙 [BUDGET] Dynamic Token Budget allocated: {max_tokens} tokens for {proof_graph_steps_count} steps.")
        
        gen_config = GenerationConfig(
            temperature=0.0, 
            top_p=0.1, 
            top_k=1,
            max_output_tokens=max_tokens,
            response_mime_type="application/json"
        )
        
        try:
            if image_pages:
                payload = [f"Images are sequential. Image 1 is Header.\n{prompt}"] + image_pages
                logger.info(f"🧠 [TRACE] PROMPT SENT TO LLM: {payload[0]}")
                # V8.5: True Async Cancellation. Using current_task to handle disconnects.
                response = await asyncio.wait_for(
                    self.llm.generate_content_async(payload, generation_config=gen_config), 
                    timeout=60.0
                )
            else:
                logger.info(f"🧠 [TRACE] PROMPT SENT TO LLM: {prompt}")
                response = await asyncio.wait_for(
                    self.llm.generate_content_async(prompt, generation_config=gen_config), 
                    timeout=30.0
                )
        except asyncio.CancelledError:
            # HARD STOP: Client disconnected. Active cancellation is handled by wait_for + generate_content_async
            logger.warning("📉 [V8.5] LLM Call actively CANCELLED due to client disconnect.")
            # Repropagate to allow parent orchestrator to catch it
            raise
        except Exception as e:
            logger.error(f"Error in _call_llm generation: {e}")
            raise

        raw_text = response.text
        logger.info(f"📦 [TRACE] RAW RESPONSE RECEIVED: {raw_text}")
        
        # V8.5.1: Convert UsageMetadata (Custom Class) to dict to avoid serialization error
        usage = getattr(response, 'usage_metadata', None)
        usage_dict = {}
        if usage:
            usage_dict = {
                "prompt_token_count": getattr(usage, 'prompt_token_count', 0),
                "candidates_token_count": getattr(usage, 'candidates_token_count', 0),
                "total_token_count": getattr(usage, 'total_token_count', 0)
            }
            
        parsed = self._parse_v4_json(raw_text)
        return {**parsed, "usage_metadata": usage_dict}

    def _parse_v4_json(self, raw_text):
        # V4.7 → V1.0: Use canonical safe_extract_json (logs RAW, LaTeX shield, json_repair, fail-closed)
        if "{" not in raw_text and "[" not in raw_text:
            return {"solution_markdown": "שגיאת מבנה", "is_valid": False, "confidence_score": 0.1}
        result = safe_extract_json(raw_text, caller="STRATEGY_MANAGER")
        if isinstance(result, dict) and result.get("logic_error"):
            return {"solution_markdown": "המורה נכשלה בפירוש הפתרון", "is_valid": False, "confidence_score": 0.1}
        return result

    def _validate(self, topic_id, llm_response, data_anchor): return {"valid": True, "error": None}

    async def generate_raw(self, system_prompt: str, user_prompt: str) -> str:
        """
        V6.1: Direct raw LLM call for ProposalEngine. Unconstrained output.
        """
        from google.generativeai.types import GenerationConfig
        import asyncio
        
        full_prompt = f"{system_prompt}\n\nUser Request: {user_prompt}"
        
        gen_config = GenerationConfig(
            temperature=0.2, # Slight creativity needed to draft proofs, but still grounded
            top_p=0.8,
        )
        
        logger.info(f"🧠 [PROPOSAL-GATEWAY] Sending PROMPT to LLM.")
        try:
            response = await asyncio.wait_for(
                asyncio.to_thread(self.llm.generate_content, full_prompt, generation_config=gen_config),
                timeout=45.0
            )
            return getattr(response, 'text', '')
        except Exception as e:
            logger.error(f"❌ [PROPOSAL-GATEWAY] LLM Request failed: {e}")
            raise e