# routes/generation_routes.py - Combined Endpoint with Correct Task Polling

import json
import time
import base64
import requests
from flask import Blueprint, request, jsonify

from db import get_db
from config import BEAM_COMBINED_URL, BEAM_API_TOKEN, get_beam_headers_json
from utils import convert_to_mp3, strip_markdown
from auth import login_required

generation_bp = Blueprint('generation', __name__)


# ============================================
# Beam Task Polling Config
# ============================================

BEAM_TASK_API = "https://api.beam.cloud/v2/task/{task_id}/"

POLL_INTERVAL = 3
POLL_MAX_WAIT = 300


def get_beam_auth_headers():
    """Beam API headers — Bearer AND Basic উভয়ই try করবে।"""
    return {
        'Authorization': f'Bearer {BEAM_API_TOKEN}',
        'Content-Type': 'application/json',
    }


def poll_beam_task(task_id):
    """Beam task poll করে result আনে।"""
    print(f"⏳ Polling task: {task_id}")

    task_url = BEAM_TASK_API.format(task_id=task_id)
    print(f"   URL: {task_url}")

    start_time = time.time()

    # প্রথম কয়েকটা attempt এ 404 আসতে পারে — task register হতে delay
    initial_delay = True

    while True:
        elapsed = time.time() - start_time

        if elapsed > POLL_MAX_WAIT:
            print(f"❌ Polling timeout after {POLL_MAX_WAIT}s")
            return None, f'Task timed out after {int(POLL_MAX_WAIT)} seconds'

        # প্রথম ২ সেকেন্ড wait করি task register হতে
        if initial_delay and elapsed < 2:
            time.sleep(2)
            initial_delay = False
            continue

        try:
            # ★ Bearer token দিয়ে try
            resp = requests.get(task_url, headers=get_beam_auth_headers(), timeout=30)

            # Bearer fail হলে Basic try করি
            if resp.status_code in (401, 403):
                print(f"   Bearer auth failed, trying Basic...")
                basic_headers = {
                    'Authorization': f'Basic {BEAM_API_TOKEN}',
                    'Content-Type': 'application/json',
                }
                resp = requests.get(task_url, headers=basic_headers, timeout=30)

            print(f"   [{int(elapsed)}s] HTTP {resp.status_code} | Body: {len(resp.text)} chars")

            if resp.status_code == 404:
                # Task এখনও register হয়নি — wait
                if elapsed < 30:
                    print(f"   Task not found yet, waiting...")
                    time.sleep(POLL_INTERVAL)
                    continue
                else:
                    # ৩০ সেকেন্ড পরেও 404 — সমস্যা
                    print(f"❌ Task not found after {int(elapsed)}s")

                    # ★ Debug: response body দেখি
                    print(f"   404 body: {resp.text[:300]}")

                    # ★ Alternative: Beam API base URL ভিন্ন হতে পারে
                    # কিছু Beam setup এ URL format ভিন্ন
                    alt_urls = [
                        f"https://api.beam.cloud/v2/task/{task_id}/status/",
                        f"https://api.beam.cloud/v2/task/{task_id}",
                        f"https://api.beam.cloud/v1/task/{task_id}/",
                    ]
                    for alt_url in alt_urls:
                        try:
                            alt_resp = requests.get(alt_url, headers=get_beam_auth_headers(), timeout=10)
                            print(f"   Alt URL {alt_url}: HTTP {alt_resp.status_code}")
                            if alt_resp.status_code == 200:
                                print(f"   ✅ Found working URL!")
                                resp = alt_resp
                                break
                        except Exception:
                            pass

                    if resp.status_code == 404:
                        return None, f'Task {task_id} not found on Beam API after {int(elapsed)}s'

            if resp.status_code == 200 and resp.text:
                try:
                    data = resp.json()
                except Exception as e:
                    print(f"   JSON parse error: {e}")
                    print(f"   Body: {resp.text[:300]}")
                    time.sleep(POLL_INTERVAL)
                    continue

                status = data.get('status', '').upper()
                print(f"   [{int(elapsed)}s] Task status: {status}")

                if status in ('COMPLETE', 'COMPLETED', 'SUCCESS'):
                    print(f"✅ Task complete!")

                    # ★ Result বের করা — Beam বিভিন্ন জায়গায় result রাখে
                    # 1. 'output' key
                    # 2. 'result' key
                    # 3. 'outputs' list (file-based)
                    # 4. response body তেই (endpoint mode)

                    actual_result = None

                    # Check 'output' (endpoint mode — function return value)
                    if data.get('output') and isinstance(data['output'], dict):
                        actual_result = data['output']
                        print(f"   Result found in 'output' key")

                    # Check 'result'
                    elif data.get('result') and isinstance(data['result'], dict):
                        actual_result = data['result']
                        print(f"   Result found in 'result' key")

                    # Check if top-level has audio_base64 (unlikely but possible)
                    elif data.get('audio_base64'):
                        actual_result = data
                        print(f"   Result found in top-level data")

                    if actual_result and actual_result.get('audio_base64'):
                        return actual_result, None
                    elif actual_result and actual_result.get('success'):
                        return actual_result, None

                    # ★ Outputs (file-based) — need to download
                    outputs = data.get('outputs', [])
                    if outputs:
                        print(f"   Task has {len(outputs)} output files")
                        # For our use case, result should be in 'output' not files
                        # But log it for debug
                        for out in outputs:
                            print(f"   Output: {out.get('name', '?')} → {out.get('url', '?')}")

                    # No usable result found
                    print(f"   ⚠️ Task complete but no audio in response")
                    print(f"   Response keys: {list(data.keys())}")
                    print(f"   Full response (first 500): {json.dumps(data, default=str)[:500]}")

                    return None, 'Task completed but no audio data in result. Check Beam logs.'

                elif status in ('FAILED', 'ERROR'):
                    error_msg = data.get('error', 'Task failed')
                    print(f"❌ Task failed: {error_msg}")
                    return None, f'Task failed: {error_msg}'

                elif status in ('CANCELLED', 'CANCELED'):
                    return None, 'Task was cancelled'

                elif status in ('TIMEOUT', 'EXPIRED'):
                    return None, f'Task {status.lower()} on Beam. Container may not have started in time.'

                elif status in ('PENDING', 'RUNNING', 'RETRY'):
                    pass  # Keep polling

                else:
                    print(f"   Unknown status: {status}")

            elif resp.status_code != 404:
                print(f"   Unexpected HTTP {resp.status_code}: {resp.text[:200]}")

        except requests.exceptions.RequestException as e:
            print(f"   Poll error: {e}")

        time.sleep(POLL_INTERVAL)


# ============================================
# Beam Call + Smart Response Handler
# ============================================

def call_beam_and_get_result(text, voice='af_heart', speed=1.0):
    """Beam combined endpoint call + async polling।"""

    if not BEAM_COMBINED_URL:
        return None, 'BEAM_COMBINED_URL is not configured in .env'

    print(f"📞 Calling: {BEAM_COMBINED_URL}")
    print(f"   text={len(text)} chars, voice={voice}")

    response = requests.post(
        BEAM_COMBINED_URL,
        headers=get_beam_headers_json(),
        json={
            'text': text,
            'voice': voice,
            'speed': speed,
            'skip_alignment': False,
        },
        timeout=300
    )

    print(f"📡 Status: {response.status_code}")
    print(f"📡 Content-Length: {response.headers.get('Content-Length', 'N/A')}")

    task_id = response.headers.get('X-Task-Id', '')

    # ========================================
    # CASE 1: Task ID + empty/no body → Async → Poll
    # ========================================
    if task_id and (not response.text or not response.text.strip() or response.headers.get('Content-Length') == '0'):
        print(f"📋 Async mode — Task ID: {task_id}")
        return poll_beam_task(task_id)

    # ========================================
    # CASE 2: Task ID + body
    # ========================================
    if task_id and response.text and response.text.strip():
        print(f"📋 Task ID: {task_id} + body ({len(response.text)} chars)")
        try:
            result = response.json()
            if result.get('success') and result.get('audio_base64'):
                print(f"✅ Direct sync result")
                return _extract(result), None
            # Body isn't the final result — poll
            return poll_beam_task(task_id)
        except Exception:
            return poll_beam_task(task_id)

    # ========================================
    # CASE 3: No task_id + empty body → Error
    # ========================================
    if not response.text or not response.text.strip():
        return None, 'Empty response from Beam with no task ID'

    # ========================================
    # CASE 4: Synchronous response
    # ========================================
    if response.status_code != 200:
        try:
            err = response.json().get('error', response.text[:200])
        except Exception:
            err = response.text[:200]
        return None, f'Beam Error ({response.status_code}): {err}'

    try:
        result = response.json()
    except Exception as e:
        return None, f'Invalid JSON: {response.text[:100]}'

    if not result.get('success'):
        return None, result.get('error', 'Unknown error')

    return _extract(result), None


def _extract(result):
    return {
        'audio_base64': result.get('audio_base64', ''),
        'audio_format': result.get('audio_format', 'wav'),
        'sample_rate': result.get('sample_rate', 24000),
        'transcription': result.get('timestamps', []),
    }


# ============================================
# API Route: Single Block
# ============================================

@generation_bp.route('/api/generate', methods=['POST'])
@login_required
def generate_audio():
    data = request.json
    text = data.get('text', '')
    voice = data.get('voice', 'af_heart')
    block_id = data.get('block_id')

    if not text:
        return jsonify({'error': 'No text provided'}), 400

    stripped = text.strip()
    if stripped.startswith('![') and '](' in stripped and stripped.endswith(')'):
        return jsonify({'error': 'Cannot generate audio for image content'}), 400

    clean_text = strip_markdown(text)
    if not clean_text.strip():
        return jsonify({'error': 'No speakable text content'}), 400

    try:
        print(f"")
        print(f"{'='*60}")
        print(f"🔊 GENERATE REQUEST")
        print(f"   Voice: {voice}, Text: {len(clean_text)} chars")
        print(f"   Preview: {clean_text[:100]}...")
        print(f"{'='*60}")

        result, error = call_beam_and_get_result(clean_text, voice)

        if error:
            print(f"❌ Failed: {error}")
            return jsonify({'error': error}), 500

        audio_base64 = result.get('audio_base64', '')
        source_format = result.get('audio_format', 'wav')
        transcription = result.get('transcription', [])

        if not audio_base64:
            return jsonify({'error': 'No audio data received'}), 500

        if source_format != 'mp3':
            audio_base64 = convert_to_mp3(audio_base64, source_format)

        if block_id:
            db = get_db()
            cursor = db.cursor()
            cursor.execute('''
                UPDATE markdown_blocks 
                SET audio_data = ?, audio_format = 'mp3', transcription = ?
                WHERE id = ?
            ''', (audio_base64, json.dumps(transcription), block_id))
            db.commit()

        print(f"✅ DONE: audio={len(audio_base64)} bytes, words={len(transcription)}")
        print(f"{'='*60}")

        return jsonify({
            'success': True,
            'audio_data': audio_base64,
            'audio_format': 'mp3',
            'transcription': transcription
        })

    except requests.exceptions.ConnectionError as e:
        print(f"❌ CONNECTION: {e}")
        return jsonify({'error': 'Cannot connect to Beam Cloud.'}), 500
    except requests.exceptions.Timeout:
        return jsonify({'error': 'Request timed out. Try again in 1-2 minutes.'}), 500
    except requests.exceptions.RequestException as e:
        return jsonify({'error': f'API error: {str(e)}'}), 500
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


# ============================================
# API Route: Chapter
# ============================================

@generation_bp.route('/api/generate-chapter', methods=['POST'])
@login_required
def generate_chapter_audio():
    data = request.json
    chapter_id = data.get('chapter_id')
    voice = data.get('voice', 'af_heart')

    if not chapter_id:
        return jsonify({'error': 'Chapter ID required'}), 400

    db = get_db()
    cursor = db.cursor()

    cursor.execute('''
        SELECT id, content, tts_text, block_type FROM markdown_blocks 
        WHERE chapter_id = ? ORDER BY block_order
    ''', (chapter_id,))
    blocks = cursor.fetchall()

    if not blocks:
        return jsonify({'error': 'No blocks found'}), 404

    results = []
    success_count = 0
    error_count = 0
    total = len(blocks)

    print(f"\n{'='*60}")
    print(f"📖 CHAPTER: {total} blocks, voice={voice}")
    print(f"{'='*60}")

    for idx, block in enumerate(blocks):
        block_id = block['id']
        block_type = block['block_type'] if 'block_type' in block.keys() else 'paragraph'
        content = block['content'] or ''
        text = block['tts_text'] if block['tts_text'] else content

        if block_type == 'image':
            results.append({'block_id': block_id, 'success': True, 'skipped': True})
            continue

        stripped = text.strip()
        if stripped.startswith('![') and '](' in stripped and stripped.endswith(')'):
            results.append({'block_id': block_id, 'success': True, 'skipped': True})
            continue

        clean_text = strip_markdown(text)
        if not clean_text.strip():
            results.append({'block_id': block_id, 'success': True, 'skipped': True})
            continue

        print(f"\n📖 Block {idx+1}/{total}: {len(clean_text)} chars")

        try:
            result, error = call_beam_and_get_result(clean_text, voice)

            if error:
                print(f"❌ Block {block_id}: {error}")
                results.append({'block_id': block_id, 'success': False, 'error': error})
                error_count += 1
                continue

            audio_base64 = result.get('audio_base64', '')
            source_format = result.get('audio_format', 'wav')
            transcription = result.get('transcription', [])

            if not audio_base64:
                results.append({'block_id': block_id, 'success': False, 'error': 'No audio'})
                error_count += 1
                continue

            if source_format != 'mp3':
                audio_base64 = convert_to_mp3(audio_base64, source_format)

            cursor.execute('''
                UPDATE markdown_blocks 
                SET audio_data = ?, audio_format = 'mp3', transcription = ?
                WHERE id = ?
            ''', (audio_base64, json.dumps(transcription), block_id))

            results.append({
                'block_id': block_id,
                'success': True,
                'audio_data': audio_base64,
                'transcription': transcription
            })
            success_count += 1
            print(f"✅ Block {idx+1} done")

        except Exception as e:
            print(f"❌ Block {block_id}: {e}")
            results.append({'block_id': block_id, 'success': False, 'error': str(e)})
            error_count += 1

    db.commit()

    skipped = total - success_count - error_count
    print(f"\n📖 COMPLETE: {success_count} ok, {error_count} fail, {skipped} skip")
    print(f"{'='*60}\n")

    return jsonify({
        'success': True,
        'results': results,
        'summary': {
            'total': total,
            'generated': success_count,
            'failed': error_count,
            'skipped': skipped
        }
    })