import os
import logging
from pathlib import Path
from dotenv import load_dotenv
import google.generativeai as genai
from concurrent.futures import ThreadPoolExecutor, as_completed
import sys

# Load .env file
env_path = Path(__file__).parent.parent.parent / '.env'
if env_path.exists():
    load_dotenv(env_path)
else:
    load_dotenv()

sys.path.append(os.path.dirname(__file__))
from utils import parse_restructured_text

logger = logging.getLogger(__name__)

# Get Gemini API key
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise EnvironmentError("GEMINI_API_KEY not set")

# Configure Gemini
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-1.5-flash-latest")


def build_batch_prompt(mcqs, theories, batch_num):
    """Build OPTIMIZED prompt for one batch."""
    prompt = f"""Reformat into {len(mcqs)} MCQs and {len(theories)} theory questions. Be CONCISE.

MCQs ({len(mcqs)} total):
{chr(10).join(f'{i+1}. {q}' for i, q in enumerate(mcqs))}

Theory ({len(theories)} total):
{chr(10).join(f'{i+1}. {q}' for i, q in enumerate(theories))}

FORMAT (NO extra text):

MCQ 1
Stem: [question]
Key: [correct answer]
Distractors:
- [wrong 1]
- [wrong 2]
- [wrong 3]

Theory 1
Question: [question]
Answer: [2-3 sentences]

Continue for all questions. NO commentary.
"""
    return prompt


def process_batch(batch_data):
    """Process one batch through Gemini with SPEED optimization."""
    batch_num, mcqs, theories = batch_data
    
    try:
        logger.info(f"Processing batch {batch_num}: {len(mcqs)} MCQs, {len(theories)} theories")
        
        prompt = build_batch_prompt(mcqs, theories, batch_num)
        
        # SPEED OPTIMIZED: Lower temperature, higher top_p, reduced tokens
        response = model.generate_content(
            prompt,
            generation_config={
                'temperature': 0.5,      # REDUCED from 0.7 for faster, more deterministic output
                'top_p': 0.9,            # REDUCED from 0.95
                'max_output_tokens': 3072,  # REDUCED from 4096
            },
            request_options={'timeout': 30}  # 30 second timeout per batch
        )
        
        text = getattr(response, "text", "").strip()
        
        if not text:
            raise ValueError(f"Empty response from Gemini for batch {batch_num}")
        
        # Parse the structured output
        mcqs_parsed, theories_parsed = parse_restructured_text(text)
        
        logger.info(f"Batch {batch_num} complete: {len(mcqs_parsed)} MCQs, {len(theories_parsed)} theories")
        
        return mcqs_parsed, theories_parsed
        
    except Exception as e:
        logger.error(f"Batch {batch_num} failed: {str(e)}")
        raise


def restructure_questions_parallel(generated_mcqs, generated_theory):
    """
    OPTIMIZED: Restructure 40 MCQs + 20 Theory questions faster.
    
    Uses 2 parallel batches with optimized Gemini settings.
    Target time: 20-30 seconds (down from 60+ seconds)
    """
    if len(generated_mcqs) < 40:
        logger.warning(f"Expected 40 MCQs, got {len(generated_mcqs)}")
    if len(generated_theory) < 20:
        logger.warning(f"Expected 20 theories, got {len(generated_theory)}")
    
    # Ensure we have exactly 40 MCQs and 20 theories
    generated_mcqs = generated_mcqs[:40]
    generated_theory = generated_theory[:20]
    
    # Split into 2 batches
    batches = [
        (1, generated_mcqs[:20], generated_theory[:10]),
        (2, generated_mcqs[20:40], generated_theory[10:20]),
    ]
    
    final_mcqs = []
    final_theory = []
    
    logger.info("Starting parallel Gemini processing (2 batches)")
    
    try:
        # OPTIMIZED: Process both batches in parallel with 60s total timeout
        with ThreadPoolExecutor(max_workers=2) as executor:
            future_to_batch = {
                executor.submit(process_batch, batch): batch[0] 
                for batch in batches
            }
            
            for future in as_completed(future_to_batch, timeout=60):  # REDUCED from 120s
                batch_num = future_to_batch[future]
                try:
                    mcqs, theories = future.result()
                    final_mcqs.extend(mcqs)
                    final_theory.extend(theories)
                    logger.info(f"Collected batch {batch_num} results")
                except Exception as e:
                    logger.error(f"Failed to collect batch {batch_num}: {str(e)}")
                    raise
    
    except TimeoutError:
        raise RuntimeError("Gemini processing timeout after 60 seconds")
    
    logger.info(f"Complete: {len(final_mcqs)} MCQs, {len(final_theory)} theories")
    
    return {
        "mcqs": final_mcqs,
        "theory": final_theory
    }