import os
import json
import logging
import re
import time
from pathlib import Path
from dotenv import load_dotenv
from flask import Flask, request, jsonify
from flask_cors import CORS
import nltk
from nltk.tokenize import sent_tokenize
from groq import Groq, RateLimitError
from functools import wraps
from threading import Lock

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

# Groq client (lazy init)
_groq_client = None
_rate_limiter = Lock()
_last_request_time = 0
_min_request_interval = 2.1  # 2.1s = ~28 req/min (safe margin)

def get_groq_client():
    """Get or create Groq client"""
    global _groq_client
    if _groq_client is None:
        if not os.getenv('GROQ_API_KEY'):
            raise ValueError("GROQ_API_KEY not configured")
        _groq_client = Groq(api_key=os.getenv('GROQ_API_KEY'))
    return _groq_client

def rate_limited_call(func):
    """Decorator to enforce rate limiting"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        global _last_request_time
        with _rate_limiter:
            current_time = time.time()
            time_since_last = current_time - _last_request_time
            if time_since_last < _min_request_interval:
                sleep_time = _min_request_interval - time_since_last
                time.sleep(sleep_time)
            _last_request_time = time.time()
        return func(*args, **kwargs)
    return wrapper

# Setup Flask
app = Flask(__name__)
CORS(app)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - [%(levelname)s] - %(name)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Download NLTK data
try:
    nltk.download('punkt', quiet=True)
    nltk.download('punkt_tab', quiet=True)
except Exception as e:
    logger.warning(f"NLTK download failed: {e}")


@app.route('/health', methods=['GET'])
def health_check():
    """Health check"""
    return jsonify({
        'status': 'healthy',
        'ai_provider': 'Groq (Llama 3.3 70B)',
        'groq_configured': bool(os.getenv('GROQ_API_KEY')),
        'endpoints': {
            'test': '/api/restructure/',
            'content': '/api/study-content/',
            'recommendations': '/api/recommendations/'
        },
        'pipeline': 'Smart Batching (5-8 calls)',
        'version': '8.0 - Production Ready'
    }), 200


@app.route('/api/restructure/', methods=['POST'])
def generate_from_studykit():
    """Test question generation - optimized batching"""
    start = time.time()
    
    try:
        if not request.is_json:
            return jsonify({'status': 'error', 'error': 'Request must be JSON'}), 400

        data = request.get_json(force=True)
        content = data.get('content', '').strip()

        if not content:
            return jsonify({'status': 'error', 'error': 'Content required'}), 400

        logger.info(f"🎯 TEST GENERATION: {len(content)} chars")

        # Tokenize
        sentences = sent_tokenize(content)
        logger.info(f"📝 Tokenized: {len(sentences)} sentences")

        # BiGNN Classification
        from question_api.classify_content import classify_sentences
        classifications = classify_sentences(sentences)
        
        # Fix tuple/dict issue
        if classifications and isinstance(classifications[0], tuple):
            classifications = [{'sentence': s, 'label': l} for s, l in classifications]
        
        logger.info(f"🔍 Classified: {len(classifications)} sentences")

        # Generate questions (optimized)
        context = "\n".join([f"- {s['sentence']}" for s in classifications[:50]])
        mcqs, theories = generate_questions_optimized(context, 40, 20)
        
        elapsed = time.time() - start
        logger.info(f"✅ Generated {len(mcqs)} MCQs + {len(theories)} theories in {elapsed:.2f}s")
        
        return jsonify({
            'status': 'success',
            'mcqs': mcqs,
            'theory': theories,
            'counts': {'mcqs': len(mcqs), 'theory': len(theories), 'total': len(mcqs) + len(theories)},
            'processing_time_seconds': round(elapsed, 2)
        }), 200

    except Exception as e:
        logger.exception("Test generation failed")
        return jsonify({'status': 'error', 'error': str(e)}), 500


@app.route('/api/study-content/', methods=['POST'])
def generate_study_content():
    """Study content generation - smart batching"""
    start = time.time()
    
    try:
        if not request.is_json:
            return jsonify({'status': 'error', 'error': 'Request must be JSON'}), 400

        data = request.get_json(force=True)
        topic = data.get('topic', '').strip()
        description = data.get('description', '').strip()

        if not topic:
            return jsonify({'status': 'error', 'error': 'Topic required'}), 400

        logger.info(f"📚 STUDY CONTENT: {topic}")

        # Generate using optimized pipeline
        result = generate_study_content_optimized(topic, description)
        
        if not result:
            return jsonify({'status': 'error', 'error': 'Generation failed'}), 500

        elapsed = time.time() - start
        logger.info(f"✅ Generated in {elapsed:.2f}s | Subtopics: {len(result['subtopics'])} | Terms: {len(result['terms'])} | Guide: {len(result['guide'])} chars")
        
        return jsonify({
            'status': 'success',
            'subtopics': result['subtopics'],
            'terms': result['terms'],
            'guide': result['guide'],
            'processing_time_seconds': round(elapsed, 2)
        }), 200

    except Exception as e:
        logger.exception("Study content failed")
        return jsonify({'status': 'error', 'error': str(e)}), 500


@app.route('/api/recommendations/', methods=['POST'])
def generate_recommendations():
    """Recommendations - 2 API calls"""
    start = time.time()
    
    try:
        if not request.is_json:
            return jsonify({'status': 'error', 'error': 'Request must be JSON'}), 400

        data = request.get_json(force=True)
        topic = data.get('topic', '').strip()

        if not topic:
            return jsonify({'status': 'error', 'error': 'Topic required'}), 400

        logger.info(f"🎥 RECOMMENDATIONS: {topic}")

        result = generate_recommendations_optimized(topic)
        
        if not result:
            return jsonify({'status': 'error', 'error': 'Generation failed'}), 500
        
        elapsed = time.time() - start
        logger.info(f"✅ Generated {len(result['videos'])} videos + {len(result['blogs'])} blogs in {elapsed:.2f}s")
        
        return jsonify({
            'status': 'success',
            'videos': result['videos'],
            'blogs': result['blogs'],
            'processing_time_seconds': round(elapsed, 2)
        }), 200
        
    except Exception as e:
        logger.exception("Recommendations failed")
        return jsonify({'status': 'error', 'error': str(e)}), 500


# ============================================================================
# OPTIMIZED GENERATION FUNCTIONS
# ============================================================================

@rate_limited_call
def call_groq(messages, temperature=0.5, max_tokens=4000):
    """Centralized Groq call with rate limiting"""
    client = get_groq_client()
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content.strip()


def generate_questions_optimized(content, target_mcqs=40, target_theories=20):
    """Generate questions in large batches (4 calls total)"""
    mcqs = []
    theories = []
    
    try:
        # 2 batches for MCQs (20 each)
        logger.info("📝 Generating MCQ batch 1/2...")
        batch1 = generate_mcq_batch(content[:1500], 20)
        mcqs.extend(batch1)
        
        logger.info("📝 Generating MCQ batch 2/2...")
        batch2 = generate_mcq_batch(content[:1500], 20)
        mcqs.extend(batch2)
        
        # 2 batches for theories (10 each)
        logger.info("📝 Generating theory batch 1/2...")
        theory1 = generate_theory_batch(content[:1500], 10)
        theories.extend(theory1)
        
        logger.info("📝 Generating theory batch 2/2...")
        theory2 = generate_theory_batch(content[:1500], 10)
        theories.extend(theory2)
        
        # Ensure counts
        mcqs = mcqs[:target_mcqs]
        theories = theories[:target_theories]
        
        # Pad if needed
        while len(mcqs) < target_mcqs:
            mcqs.append({
                "stem": f"Question {len(mcqs)+1}?",
                "key": "Answer",
                "distractors": ["A", "B", "C"]
            })
        
        while len(theories) < target_theories:
            theories.append({
                "question": f"Question {len(theories)+1}?",
                "answer": "Detailed answer required."
            })
        
        return mcqs, theories
        
    except Exception as e:
        logger.error(f"Question generation failed: {e}")
        return (
            [{"stem": f"Q{i+1}?", "key": "A", "distractors": ["B", "C", "D"]} for i in range(target_mcqs)],
            [{"question": f"Q{i+1}?", "answer": "Answer"} for i in range(target_theories)]
        )


def generate_mcq_batch(content, count):
    """Generate MCQ batch"""
    prompt = f"""Generate EXACTLY {count} multiple-choice questions.

CONTENT:
{content}

Return ONLY valid JSON (no markdown):
{{
  "mcqs": [
    {{
      "stem": "Question?",
      "key": "Correct",
      "distractors": ["Wrong1", "Wrong2", "Wrong3"]
    }}
  ]
}}

EXACTLY {count} questions. Test understanding."""

    try:
        result_text = call_groq(
            [
                {"role": "system", "content": "You generate exam questions. Return only JSON."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.6,
            max_tokens=3000
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        return result.get('mcqs', [])[:count]
        
    except Exception as e:
        logger.error(f"MCQ batch failed: {e}")
        return []


def generate_theory_batch(content, count):
    """Generate theory batch"""
    prompt = f"""Generate EXACTLY {count} theory questions.

CONTENT:
{content}

Return ONLY valid JSON (no markdown):
{{
  "theory": [
    {{
      "question": "Question?",
      "answer": "2-3 paragraph detailed answer."
    }}
  ]
}}

EXACTLY {count} questions."""

    try:
        result_text = call_groq(
            [
                {"role": "system", "content": "You generate exam questions. Return only JSON."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.6,
            max_tokens=4000
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        return result.get('theory', [])[:count]
        
    except Exception as e:
        logger.error(f"Theory batch failed: {e}")
        return []


def generate_study_content_optimized(topic, description):
    """Generate study content (5-7 calls)"""
    try:
        # CALL 1: All subtopics
        logger.info("📋 Step 1/5: Generating 20 subtopics...")
        subtopics = generate_all_subtopics(topic, description, 20)
        
        # CALL 2: All terms
        logger.info("📋 Step 2/5: Generating 30 terms...")
        terms = generate_all_terms(topic, description, 30)
        
        # CALL 3: Introduction
        logger.info("📋 Step 3/5: Generating introduction...")
        intro = generate_introduction(topic, description)
        
        # CALL 4-5: Main content in 2 sections
        logger.info("📋 Step 4/5: Generating main content part 1...")
        main1 = generate_guide_section(topic, "foundational concepts and principles", 600)
        
        logger.info("📋 Step 5/5: Generating main content part 2...")
        main2 = generate_guide_section(topic, "advanced topics and applications", 600)
        
        # Assemble guide
        guide = f"""# {topic}: Complete Study Guide

{intro}

## Part 1: Foundational Concepts

{main1}

## Part 2: Advanced Applications

{main2}

## Summary

This guide covers the essential aspects of {topic}. Focus on understanding core principles and connecting theoretical knowledge to practical applications. Regular review and active engagement lead to mastery."""

        return {'subtopics': subtopics, 'terms': terms, 'guide': guide}
        
    except Exception as e:
        logger.error(f"Study content failed: {e}")
        return None


def generate_all_subtopics(topic, description, count):
    """Generate all subtopics in one call"""
    context = f" Context: {description}" if description else ""
    
    prompt = f"""Generate EXACTLY {count} subtopics for "{topic}".{context}

Return ONLY valid JSON:
{{
  "subtopics": [
    "Title - Subtitle: 2-3 sentences (40-70 words) explaining what students learn and why it matters."
  ]
}}

EXACTLY {count} subtopics."""

    try:
        result_text = call_groq(
            [{"role": "system", "content": "You are a curriculum designer. Return only JSON."},
             {"role": "user", "content": prompt}],
            temperature=0.5,
            max_tokens=3000
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        subtopics = result.get('subtopics', [])
        
        while len(subtopics) < count:
            subtopics.append(f"Topic {len(subtopics)+1} - Study Area: Explore {topic}")
        
        return subtopics[:count]
        
    except Exception as e:
        logger.error(f"Subtopics failed: {e}")
        return [f"Subtopic {i+1}: Explore {topic}" for i in range(count)]


def generate_all_terms(topic, description, count):
    """Generate all terms in one call"""
    context = f" Context: {description}" if description else ""
    
    prompt = f"""Generate EXACTLY {count} key terms for "{topic}".{context}

Return ONLY valid JSON:
{{
  "terms": [
    "Term - Category: Definition + Context + Example + Significance (50-100 words)"
  ]
}}

EXACTLY {count} terms."""

    try:
        result_text = call_groq(
            [{"role": "system", "content": "You are a terminology expert. Return only JSON."},
             {"role": "user", "content": prompt}],
            temperature=0.5,
            max_tokens=4000
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        terms = result.get('terms', [])
        
        while len(terms) < count:
            terms.append(f"Term {len(terms)+1} - Concept: Important concept in {topic}")
        
        return terms[:count]
        
    except Exception as e:
        logger.error(f"Terms failed: {e}")
        return [f"Term {i+1}: Key concept in {topic}" for i in range(count)]


def generate_introduction(topic, description):
    """Generate introduction"""
    context = f"\n\nContext: {description}" if description else ""
    
    prompt = f"""Write a comprehensive introduction for "{topic}".{context}

FORMAT REQUIREMENTS:
- Markdown only
- Start with ## Introduction
- Use many short paragraphs
- Each paragraph 3-5 sentences
- Blank line between paragraphs
- 300-500 words
"""

    try:
        return call_groq(
            [{"role": "system", "content": "You are an academic writer."},
             {"role": "user", "content": prompt}],
            temperature=0.6,
            max_tokens=800
        )
    except Exception as e:
        logger.error(f"Introduction failed: {e}")
        return f"## Introduction\n\nThis guide covers {topic}."


def generate_guide_section(topic, focus, target_words):
    """Generate guide section"""
    prompt = f"""Write educational content about {topic}, focusing on {focus}.

FORMAT REQUIREMENTS:
- Markdown only
- Start with ## {focus}
- Exactly 3 subsections
- Each subsection starts with ###
- Each subsection has 3–4 paragraphs
- Blank line between every paragraph
- No bullets
"""

    try:
        return call_groq(
            [{"role": "system", "content": "You are an expert educator."},
             {"role": "user", "content": prompt}],
            temperature=0.6,
            max_tokens=2000
        )
    except Exception as e:
        logger.error(f"Section failed: {e}")
        return f"### Content\n\nDetailed exploration of {focus} in {topic}."


def generate_recommendations_optimized(topic):
    """Generate recommendations (2 calls)"""
    try:
        logger.info("🎥 Generating 6 videos...")
        videos = generate_all_videos(topic, 6)
        
        logger.info("📰 Generating 4 blogs...")
        blogs = generate_all_blogs(topic, 4)
        
        return {'videos': videos, 'blogs': blogs}
    except Exception as e:
        logger.error(f"Recommendations failed: {e}")
        return None


def generate_all_videos(topic, count):
    """Generate all videos in one call"""
    prompt = f"""Generate EXACTLY {count} YouTube video recommendations for "{topic}".

Return ONLY valid JSON:
{{
  "videos": [
    {{
      "title": "Video Title",
      "url": "https://youtube.com/watch?v=ID",
      "channel": "Channel Name",
      "thumbnail": "https://img.youtube.com/vi/ID/maxresdefault.jpg",
      "description": "40-60 word description"
    }}
  ]
}}

Use real channels: 3Blue1Brown, Khan Academy, Crash Course, etc.
EXACTLY {count} videos."""

    try:
        result_text = call_groq(
            [{"role": "system", "content": "You are a resource curator. Return only JSON."},
             {"role": "user", "content": prompt}],
            temperature=0.4,
            max_tokens=2000
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        videos = result.get('videos', [])
        
        for video in videos:
            if 'url' in video and 'youtube.com' in video['url']:
                video_id = extract_youtube_id(video['url'])
                if video_id:
                    video['thumbnail'] = f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
        
        return videos[:count]
        
    except Exception as e:
        logger.error(f"Videos failed: {e}")
        return [
            {
                "title": f"Learn {topic} - Video {i+1}",
                "url": f"https://youtube.com/results?search_query={topic.replace(' ', '+')}",
                "channel": "Educational Channel",
                "thumbnail": "https://img.youtube.com/vi/PLACEHOLDER/maxresdefault.jpg",
                "description": f"Educational video about {topic}."
            }
            for i in range(count)
        ]


def generate_all_blogs(topic, count):
    """Generate all blogs in one call"""
    prompt = f"""Generate EXACTLY {count} blog recommendations for "{topic}".

Return ONLY valid JSON:
{{
  "blogs": [
    {{
      "title": "Article Title",
      "url": "https://site.com/article",
      "site": "Site Name",
      "description": "40-60 word description"
    }}
  ]
}}

Use real sites: Medium, Towards Data Science, etc.
EXACTLY {count} articles."""

    try:
        result_text = call_groq(
            [{"role": "system", "content": "You are a resource curator. Return only JSON."},
             {"role": "user", "content": prompt}],
            temperature=0.4,
            max_tokens=1500
        )
        
        result_text = clean_json(result_text)
        result = json.loads(result_text)
        blogs = result.get('blogs', [])
        
        for blog in blogs:
            if 'site' not in blog and 'url' in blog:
                blog['site'] = extract_domain(blog['url'])
        
        return blogs[:count]
        
    except Exception as e:
        logger.error(f"Blogs failed: {e}")
        return [
            {
                "title": f"Article {i+1}: {topic}",
                "url": f"https://www.google.com/search?q={topic.replace(' ', '+')}",
                "site": "Educational Website",
                "description": f"Article about {topic}."
            }
            for i in range(count)
        ]


# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

def clean_json(text):
    """Remove markdown from JSON"""
    text = re.sub(r'^```json\s*\n?', '', text, flags=re.MULTILINE)
    text = re.sub(r'^```\s*\n?', '', text, flags=re.MULTILINE)
    text = re.sub(r'\n?```\s*$', '', text, flags=re.MULTILINE)
    return text.strip()


def extract_youtube_id(url):
    """Extract YouTube ID"""
    match = re.search(r'(?:v=|/)([a-zA-Z0-9_-]{11})', url)
    return match.group(1) if match else None


def extract_domain(url):
    """Extract domain"""
    match = re.search(r'https?://(?:www\.)?([^/]+)', url)
    return match.group(1) if match else "Website"


@app.errorhandler(404)
def not_found(e):
    return jsonify({'status': 'error', 'error': 'Not found'}), 404


@app.errorhandler(500)
def internal_error(e):
    logger.error(f"500 error: {e}")
    return jsonify({'status': 'error', 'error': 'Internal error'}), 500


if __name__ == '__main__':
    if not os.getenv('GROQ_API_KEY'):
        print("\n⚠️  WARNING: GROQ_API_KEY not found!\n")
    
    print("\n" + "="*60)
    print("🚀 GROQ OPTIMIZED API v8.0 - PRODUCTION READY")
    print("="*60)
    print("Architecture: Smart Batching (5-8 calls)")
    print("Rate Limit: 2.1s spacing (~28 req/min)")
    print("Expected Time: 15-25s per request")
    print("\nEndpoints:")
    print("  POST /api/restructure/ - Test generation")
    print("  POST /api/study-content/ - Study content")
    print("  POST /api/recommendations/ - Recommendations")
    print("  GET /health - Health check")
    print("="*60 + "\n")
    
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)