# ============================================
# AutoArchitect — Main Flask Server
# ============================================

import sys
import io

# Force UTF-8 on Windows cp1252 terminals so emoji in print() don't crash.
# line_buffering=True ensures every \n flushes immediately — without it
# the TextIOWrapper is block-buffered and nothing appears in real time.
if hasattr(sys.stdout, 'buffer'):
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8',
                                  errors='replace', line_buffering=True)
if hasattr(sys.stderr, 'buffer'):
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8',
                                  errors='replace', line_buffering=True)

from flask import Flask, render_template, request, jsonify, send_file
from flask_cors import CORS
import json
import os
import time
import uuid
import threading
import zipfile
from dotenv import load_dotenv
from api.self_trainer import self_train

load_dotenv()
os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN', '')
os.environ['HUGGING_FACE_HUB_TOKEN'] = os.getenv('HF_TOKEN', '')

# === Download models from HF Model repo on first boot ===
from huggingface_hub import snapshot_download

MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
os.makedirs(MODELS_DIR, exist_ok=True)

# Only download if models aren't already there
if not os.path.exists(os.path.join(MODELS_DIR, "bert", "model.safetensors")):
    print("[STARTUP] Downloading models from HF Model repo...")
    try:
        snapshot_download(
            repo_id="dhanrajpandya/dhanrajpandyaautoarchitectmodels",
            repo_type="model",
            local_dir=MODELS_DIR,
            token=os.environ.get("HF_TOKEN")
        )
        print("[STARTUP] Models downloaded successfully")
    except Exception as e:
        print(f"[STARTUP] Model download failed: {e}")
        print("[STARTUP] Continuing without preloaded models — will fall back at runtime")
else:
    print("[STARTUP] Models already cached, skipping download")
# === End model download block ===

sys.path.append(os.path.dirname(__file__))
from api.analyzer      import ProblemAnalyzer
from api.nas_engine    import run_quick_nas
from api.cache_manager import (
    check_cache, save_to_cache,
    increment_use_count, get_cache_stats,
    find_similar_cached
)
from api.auto_trainer  import (
    select_base_model, train_new_model,
    run_yolo_detection
)
from api.orchestrator  import AutoArchitectOrchestrator
from api.data_uploader import (
    process_user_data, train_on_user_data,
    predict_with_user_model
)

import logging

# Ensure werkzeug HTTP request logs (the "127.0.0.1 POST /api/..." lines)
# are always visible.  Flask debug mode sets this up, but our custom
# sys.stderr wrapper can confuse the handler — re-attach it explicitly.
_wz_logger = logging.getLogger('werkzeug')
if not _wz_logger.handlers:
    _wz_handler = logging.StreamHandler(sys.stderr)
    _wz_handler.setFormatter(
        logging.Formatter('%(message)s'))
    _wz_logger.addHandler(_wz_handler)
_wz_logger.setLevel(logging.INFO)

app = Flask(__name__)
CORS(app)

# In-memory job store for non-blocking orchestration
_jobs = {}
_JOBS_TTL = 600  # seconds — prune finished jobs after 10 min

def _cleanup_jobs():
    """Background thread: remove jobs older than _JOBS_TTL to prevent memory growth."""
    while True:
        time.sleep(120)
        cutoff = time.time() - _JOBS_TTL
        stale  = [jid for jid, j in list(_jobs.items())
                  if j['started_at'] < cutoff
                  and j['status'] in ('complete', 'error')]
        for jid in stale:
            _jobs.pop(jid, None)
        if stale:
            print(f"[JOBS] Pruned {len(stale)} stale job(s)", flush=True)

threading.Thread(target=_cleanup_jobs, daemon=True).start()

def _watchdog_jobs():
    """Force-error jobs beyond mode-specific time caps (quick=6 min, others=12 min)."""
    while True:
        time.sleep(60)
        now = time.time()
        for jid, j in list(_jobs.items()):
            if j.get('status') != 'running':
                continue
            elapsed = now - j.get('started_at', now)
            mode    = j.get('training_mode', 'balanced')
            if mode == 'quick' and elapsed > 360:
                j['status'] = 'error'
                j['error']  = 'Quick mode exceeded 6-minute limit. Try again or use Balanced mode.'
                print(f"[WATCHDOG] Force-erroring stuck QUICK job {jid} after {elapsed:.0f}s", flush=True)
            elif mode != 'quick' and elapsed > 720:
                j['status'] = 'error'
                j['error']  = 'Training exceeded 12-minute limit. Try Quick Demo mode.'
                print(f"[WATCHDOG] Force-erroring stuck job {jid} after {elapsed:.0f}s", flush=True)

threading.Thread(target=_watchdog_jobs, daemon=True).start()

# ── ANAS toggle ────────────────────────────────────────────────────────────
# Set to False to revert topology/workflow selection to template matching
# and strategy library (useful for debugging or A/B comparison).
USE_ANAS = True

analyzer     = ProblemAnalyzer()
orchestrator = AutoArchitectOrchestrator(
    groq_api_key=os.getenv("GROQ_API_KEY", "")
)

# Wire USE_ANAS into the topology designer and workflow generator
if orchestrator.topology_enabled and orchestrator.topology_designer:
    orchestrator.topology_designer.use_anas = USE_ANAS
if orchestrator.brain_enabled and orchestrator.brain:
    orchestrator.brain.use_anas = USE_ANAS


# ============================================
# ROUTES
# ============================================

@app.route('/')
def index():
    return render_template('index.html')


@app.route('/api/analyze', methods=['POST'])
def analyze():
    data    = request.json
    problem = data.get('problem', '')
    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    result            = analyzer.analyze(problem)
    result['problem'] = problem
    return jsonify(result)


@app.route('/api/search', methods=['POST'])
def search():
    data        = request.json
    num_classes = data.get('num_classes', 10)
    problem     = data.get('problem', '')
    category    = data.get('category', 'image')
    confidence  = data.get('confidence', 0)

    cached = check_cache(problem)
    if cached['found']:
        print(f"⚡ Cache hit! Loading: {problem[:40]}")
        increment_use_count(cached)
        meta         = cached['metadata']
        architecture = meta.get('architecture', [])
        if not architecture:
            architecture = [
                {'cell': 1, 'operations': [
                    {'operation': 'conv5x5', 'confidence': 95.0,
                     'weights': {'skip':0.01,'conv3x3':0.15,
                                 'conv5x5':0.75,'maxpool':0.05,
                                 'avgpool':0.04}}]},
                {'cell': 2, 'operations': [
                    {'operation': 'conv5x5', 'confidence': 99.0,
                     'weights': {'skip':0.0,'conv3x3':0.01,
                                 'conv5x5':0.99,'maxpool':0.0,
                                 'avgpool':0.0}}]},
                {'cell': 3, 'operations': [
                    {'operation': 'maxpool', 'confidence': 100.0,
                     'weights': {'skip':0.0,'conv3x3':0.0,
                                 'conv5x5':0.0,'maxpool':1.0,
                                 'avgpool':0.0}}]}
            ]
        return jsonify({
            'architecture': architecture,
            'parameters':   meta.get('parameters', 105910),
            'search_time':  meta.get('search_time', meta.get('time', 0)),
            'status':       'success',
            'from_cache':   True,
            'use_count':    meta.get('use_count', 1),
            'trained_at':   meta.get('trained_at', ''),
            'message':      'Loaded from knowledge base!'
        })

    print(f"🔍 New problem! Training: {problem[:40]}")
    results = run_quick_nas(num_classes=num_classes)
    save_to_cache(
        problem      = problem,
        category     = category,
        confidence   = confidence,
        architecture = results['architecture'],
        parameters   = results['parameters'],
        search_time  = results['search_time']
    )
    results['from_cache'] = False
    results['message']    = 'New model trained and saved!'
    return jsonify(results)


@app.route('/api/cache/stats', methods=['GET'])
def cache_stats():
    stats = get_cache_stats()
    try:
        from api.cloud_cache import get_cloud_cache
        cloud = get_cloud_cache()
        trained_dir = os.path.join(os.path.dirname(__file__), 'models', 'trained')
        local_count = len(os.listdir(trained_dir)) if os.path.exists(trained_dir) else 0
        stats['cloud_enabled'] = cloud.enabled
        stats['local_agents']  = local_count
    except Exception:
        pass
    return jsonify(stats)


@app.route('/api/self-train', methods=['POST'])
def self_train_route():
    data     = request.json
    problem  = data.get('problem', '')
    category = data.get('category', 'image')
    if not problem:
        return jsonify({'error': 'No problem'}), 400
    print(f"\n🤖 Self-training agent activated!")
    try:
        results = self_train(problem=problem, category=category, epochs=3)
        return jsonify({
            'status':         'success',
            'problem':        problem,
            'dataset':        results['dataset'],
            'train_accuracy': results['train_accuracy'],
            'test_accuracy':  results['test_accuracy'],
            'parameters':     results['parameters'],
            'train_size':     results['train_size'],
            'time':           results['time'],
            'classes':        results['classes'],
            'self_trained':   True,
            'message':        f"Trained on {results['train_size']} samples!"
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/train', methods=['POST'])
def auto_train():
    data     = request.json
    problem  = data.get('problem', '')
    category = data.get('category', 'image')
    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    selected = select_base_model(problem, category)
    results  = train_new_model(problem, category)
    nas      = run_quick_nas(num_classes=10)
    save_to_cache(
        problem      = problem,
        category     = category,
        confidence   = data.get('confidence', 0),
        architecture = nas['architecture'],
        parameters   = nas['parameters'],
        search_time  = nas['search_time']
    )
    return jsonify({
        'status':      'success',
        'base_model':  results['base_model'],
        'description': results['description'],
        'accuracy':    results['accuracy'],
        'train_time':  results['train_time'],
        'from_cache':  False,
        'message':     'New model trained and cached!'
    })


@app.route('/api/detect', methods=['POST'])
def detect():
    import base64
    from PIL import Image
    data     = request.json
    img_data = data.get('image', '')
    problem  = data.get('problem', '').lower()
    try:
        img_bytes = base64.b64decode(img_data.split(',')[1])
        img       = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        tmp_path  = os.path.join(os.path.dirname(__file__), 'tmp_detect.jpg')
        img.save(tmp_path)
        results = run_yolo_detection(tmp_path)
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        if results['boxes']:
            label = 'pothole'  if 'pothole' in problem else \
                    'disease'  if 'disease' in problem or 'crop' in problem else \
                    'person'   if 'person'  in problem or 'face' in problem else \
                    'fire'     if 'fire'    in problem else \
                    'defect'   if 'defect'  in problem else \
                    results['boxes'][0]['label']
            for box in results['boxes']:
                box['label'] = label
        return jsonify(results)
    except Exception as e:
        return jsonify({'error': str(e), 'boxes': []}), 500


@app.route('/api/explain', methods=['POST'])
def explain():
    data    = request.json
    problem = data.get('problem', '')
    arch    = data.get('architecture', [])
    params  = data.get('parameters', 0)
    return jsonify({'explanation': generate_fallback_explanation(problem, arch, params)})


@app.route('/api/download', methods=['GET'])
def download_model():
    model_path = os.path.join(os.path.dirname(__file__), 'models', 'nas_model.pth')
    if os.path.exists(model_path):
        return send_file(model_path, as_attachment=True,
                         download_name='autoarchitect_model.pth',
                         mimetype='application/octet-stream')
    return jsonify({'error': 'Model not found'}), 404


@app.route('/api/predict', methods=['POST'])
def predict():
    import torchvision.transforms as transforms
    from PIL import Image
    import base64
    import torch
    from api.nas_engine import DARTSNet

    data     = request.json
    img_data = data.get('image', '')
    category = data.get('category', 'image')
    classes  = {
        'image':    ['No issue','Minor damage','Moderate damage','Severe damage',
                     'Object detected','Pattern detected','Anomaly','Clear path',
                     'Obstruction','Requires attention'],
        'medical':  ['Normal','Mild concern','Moderate concern','Severe concern',
                     'Critical','Infection','Inflammation','Healthy tissue',
                     'Requires review','Urgent'],
        'text':     ['Positive','Negative','Neutral','Spam','Urgent','Normal',
                     'Important','Low priority','High priority','Flagged'],
        'security': ['Safe','Suspicious','Attack detected','Fraud detected',
                     'Malware','Intrusion','Phishing','DDoS','Normal traffic','Danger']
    }
    try:
        img_bytes  = base64.b64decode(img_data.split(',')[1])
        img        = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        transform  = transforms.Compose([
            transforms.Resize((32, 32)), transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
        ])
        tensor     = transform(img).unsqueeze(0)
        model_path = os.path.join(os.path.dirname(__file__), 'models', 'nas_model.pth')
        model      = DARTSNet(C=16, num_cells=3, num_classes=10)
        model.load_state_dict(torch.load(model_path,
            map_location=torch.device('cpu'), weights_only=True))
        model.eval()
        with torch.no_grad():
            output     = model(tensor)
            probs      = torch.softmax(output, dim=1)
            pred_idx   = probs.argmax().item()
            confidence = round(probs.max().item() * 100, 1)
        return jsonify({
            'label':      classes.get(category, classes['image'])[pred_idx],
            'confidence': confidence,
            'status':     'success'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# ============================================
# UNIVERSAL ORCHESTRATOR
# ============================================
@app.route('/api/orchestrate', methods=['POST'])
def orchestrate():
    data       = request.json
    problem    = data.get('problem', '').strip()
    image_data = data.get('image', '')
    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    print(f"\n🚀 Orchestrating: {problem[:50]}")
    try:
        result = orchestrator.solve(problem=problem, image_data=image_data)
        return jsonify(result)
    except Exception as e:
        print(f"❌ Orchestrator error: {e}")
        return jsonify({'error': str(e)}), 500


# ============================================
# USER DATA UPLOAD
# ============================================
@app.route('/api/upload-data', methods=['POST'])
def upload_data():
    data     = request.json
    problem  = data.get('problem', '').strip()
    category = data.get('category', 'image')
    files    = data.get('files', [])
    labels   = data.get('labels', [])

    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    if not files:
        return jsonify({'error': 'No files uploaded'}), 400
    if len(files) != len(labels):
        return jsonify({'error': 'Files and labels must match'}), 400
    if len(set(labels)) < 2:
        return jsonify({'error': 'Need at least 2 different classes'}), 400

    print(f"\n📦 User data upload: {problem[:40]}")
    print(f"   Files: {len(files)}, Classes: {list(set(labels))}")

    try:
        data_info = process_user_data(
            files=files, labels=labels,
            problem=problem, category=category
        )
        epochs = min(15, max(5, len(files) // 3))
        result = train_on_user_data(
            data_info=data_info,
            problem=problem,
            category=category,
            epochs=epochs
        )
        nas = run_quick_nas(num_classes=data_info['n_classes'])
        save_to_cache(
            problem         = problem,
            category        = category,
            confidence      = 95.0,
            architecture    = nas['architecture'],
            parameters      = result.get('parameters', 0),
            search_time     = result.get('time', 0),
            result_type     = 'user_trained',
            agents_used     = [category],
            self_trained    = True,
            avg_accuracy    = result.get('test_accuracy', 0),
            all_accuracies  = {category: result.get('test_accuracy', 0)},
            user_model_path = result.get('model_path', ''),
            classes         = data_info['classes'],
        )
        return jsonify({
            'status':         'success',
            'problem':        problem,
            'category':       category,
            'classes':        data_info['classes'],
            'total_files':    data_info['total_files'],
            'train_accuracy': result.get('train_accuracy', 0),
            'test_accuracy':  result.get('test_accuracy', 0),
            'train_size':     result.get('train_size', 0),
            'architecture':   result.get('architecture', 'ResNet18'),
            'parameters':     result.get('parameters', 0),
            'time':           result.get('time', 0),
            'model_path':     result.get('model_path', ''),
            'user_trained':   True,
            'message': (
                f"Trained on YOUR {len(files)} examples! "
                f"Accuracy: {result.get('test_accuracy', 0)}%"
            )
        })
    except Exception as e:
        print(f"❌ Upload error: {e}")
        return jsonify({'error': str(e)}), 500


# ============================================
# REAL INFERENCE
# ============================================
@app.route('/api/predict-user', methods=['POST'])
def predict_user():
    data       = request.json
    problem    = data.get('problem', '')
    image_data = data.get('image', '')
    text_data  = data.get('text', '')
    category   = data.get('category', 'image')

    cached = check_cache(problem)
    if not cached['found']:
        return jsonify({'error': 'No trained model found. Upload your data first!'}), 404

    meta       = cached['metadata']
    model_path = meta.get('user_model_path', '')
    classes    = meta.get('classes', [])

    if not model_path or not os.path.exists(model_path):
        return jsonify({'error': 'Model file not found. Please retrain.'}), 404

    try:
        input_data = image_data or text_data
        result     = predict_with_user_model(
            model_path = model_path,
            input_data = input_data,
            category   = category,
            classes    = classes
        )
        return jsonify({
            'status':     'success',
            'label':      result['label'],
            'confidence': result['confidence'],
            'all_scores': result.get('all_scores', {}),
            'classes':    classes,
            'problem':    problem,
            'message':    f"Prediction: {result['label']} ({result['confidence']}%)"
        })
    except Exception as e:
        print(f"❌ Prediction error: {e}")
        return jsonify({'error': str(e)}), 500


# ============================================
# RETRAIN ON NEW DATA
# ============================================
@app.route('/api/retrain', methods=['POST'])
def retrain_model():
    data          = request.json or {}
    problem       = data.get('problem', '').strip()
    domain        = data.get('domain', 'image')
    new_data_path = data.get('new_data_path', '').strip()

    if not problem:
        return jsonify({'error': 'problem is required'}), 400
    if not new_data_path:
        return jsonify({'error': 'new_data_path is required'}), 400
    if not os.path.exists(new_data_path):
        return jsonify({'error': f'new_data_path not found: {new_data_path}'}), 400

    print(f"\n[Retrain] Endpoint: {problem[:50]}")
    try:
        from api.brain.network_zip_generator import NetworkZipGenerator
        result = NetworkZipGenerator().retrain(
            problem       = problem,
            domain        = domain,
            new_data_path = new_data_path,
        )
        return jsonify(result)
    except FileNotFoundError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        print(f"[Retrain] Error: {e}")
        return jsonify({'error': str(e)}), 500


# ============================================
# DOWNLOAD MULTI-AGENT NAS PACKAGE
# ============================================
@app.route('/api/download/multi-nas', methods=['POST'])
def download_multi_nas():
    from api.brain.agent_generator import (
        generate_agent_code,
        generate_api_server,
        generate_predict_cli,
        generate_readme,
        generate_requirements,
    )

    data    = request.json
    problem = data.get('problem', 'custom_problem')

    # ── Get metadata from cache if available ──
    from api.cache_manager import check_cache
    cached   = check_cache(problem)
    meta     = cached.get('metadata', {}) if cached.get('found') else {}

    category    = meta.get('category',    'image')
    classes     = meta.get('classes',     ['class_0', 'class_1'])
    accuracy    = meta.get('test_accuracy',
                  meta.get('avg_accuracy', 0))
    method      = meta.get('method',      'transfer_learning_resnet18')
    agents_used = meta.get('agents_used', [category])

    if not classes:
        classes = ['class_0', 'class_1']

    zip_buffer = io.BytesIO()
    base       = os.path.dirname(__file__)

    with zipfile.ZipFile(zip_buffer, 'w',
                         zipfile.ZIP_DEFLATED) as zf:

        zf.writestr('agent.py',
            generate_agent_code(
                problem     = problem,
                category    = category,
                classes     = classes,
                accuracy    = accuracy,
                method      = method,
                agents_used = agents_used,
            ))
        zf.writestr('api_server.py', generate_api_server())
        zf.writestr('predict.py',
            generate_predict_cli(
                problem  = problem,
                category = category,
                classes  = classes,
            ))
        zf.writestr('README.md',
            generate_readme(
                problem     = problem,
                category    = category,
                classes     = classes,
                accuracy    = accuracy,
                method      = method,
                agents_used = agents_used,
            ))
        zf.writestr('requirements.txt', generate_requirements())

        model_saved = False
        if cached.get('found'):
            cache_model = os.path.join(cached.get('path', ''), 'model.pth')
            if os.path.exists(cache_model):
                zf.write(cache_model, 'models/agent_model.pth')
                model_saved = True
        if not model_saved:
            mp = os.path.join(base, 'models', 'nas_model.pth')
            if os.path.exists(mp):
                zf.write(mp, 'models/agent_model.pth')

        core_files = [
            'api/nas_engine.py', 'api/dataset_fetcher.py',
            'api/transfer_trainer.py', 'api/self_trainer.py',
            'api/dataset_manager.py', 'api/analyzer.py',
            'api/cache_manager.py', 'api/workflow_engine.py',
        ]
        for f in core_files:
            path = os.path.join(base, f)
            if os.path.exists(path):
                zf.write(path, f)

        for agent in ['image_agent', 'text_agent', 'medical_agent',
                      'security_agent', 'fusion_agent', 'evaluator_agent']:
            path = os.path.join(base, f'api/agents/{agent}.py')
            if os.path.exists(path):
                zf.write(path, f'api/agents/{agent}.py')

        for bf in ['strategy_library.py', 'performance_tracker.py',
                   'workflow_generator.py', 'meta_learner.py',
                   'output_generator.py', '__init__.py']:
            path = os.path.join(base, f'api/brain/{bf}')
            if os.path.exists(path):
                zf.write(path, f'api/brain/{bf}')

        sp = os.path.join(base, 'brain_data', 'strategies.json')
        if os.path.exists(sp):
            zf.write(sp, 'brain_data/strategies.json')

        zf.writestr('api/__init__.py',        '')
        zf.writestr('api/agents/__init__.py', '')
        zf.writestr('api/brain/__init__.py',  '')
        zf.writestr('brain_data/.gitkeep',    '')
        zf.writestr('cache/.gitkeep',         '')
        zf.writestr('datasets/.gitkeep',      '')
        zf.writestr('models/.gitkeep',        '')
        zf.writestr('.env.template',
            'GROQ_API_KEY=your_groq_key_here\n'
            '# Get free key at: console.groq.com\n')
        zf.writestr('run_nas.py', _run_script(problem))

    zip_buffer.seek(0)
    safe = problem[:30].replace(' ', '_').lower()
    return send_file(
        zip_buffer,
        as_attachment = True,
        download_name = f'autoarchitect_{safe}.zip',
        mimetype      = 'application/zip'
    )


# ============================================
# DOWNLOAD FULL AGENT NETWORK (NEW)
# ============================================
@app.route('/api/download/network', methods=['POST'])
def download_network():
    data    = request.json or {}
    problem = data.get('problem', '') or orchestrator._last_problem

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

    print(f"\n📦 Network download requested: {problem[:50]}")

    try:
        # generate_network_zip now returns 3 values (includes eval_result)
        zip_bytes, topology, eval_result = orchestrator.generate_network_zip(
            problem         = problem,
            workflow_result = orchestrator._last_workflow_result,
        )

        if not zip_bytes:
            return jsonify({'error': 'Network generation failed'}), 500

        agents    = topology.get('agents',   ['agent'])
        topo_type = topology.get('topology', 'network')
        safe_name = problem[:25].replace(' ', '_').lower()
        filename  = f"{safe_name}_{topo_type}_network.zip"

        print(f"   ✅ Generated: {filename}")
        print(f"   Agents:  {' → '.join(agents)}")
        print(f"   Score:   {eval_result.get('score', 'N/A')}/100")
        print(f"   Size:    {len(zip_bytes):,} bytes")

        return send_file(
            io.BytesIO(zip_bytes),
            mimetype      = 'application/zip',
            as_attachment = True,
            download_name = filename,
        )

    except Exception as e:
        print(f"❌ Network download error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


# ============================================
# TOPOLOGY PREVIEW
# ============================================
@app.route('/api/topology/preview', methods=['POST'])
def topology_preview():
    data    = request.json or {}
    problem = data.get('problem', '') or orchestrator._last_problem

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

    try:
        domain   = orchestrator._last_workflow_result.get('domain', 'text')
        topology = orchestrator.topology_designer.design(
            problem = problem,
            domain  = domain,
        ) if orchestrator.topology_enabled else {}

        return jsonify({
            'status':      'success',
            'problem':     problem,
            'agents':      topology.get('agents', []),
            'topology':    topology.get('topology', ''),
            'connections': topology.get('connections', []),
            'agent_roles': topology.get('agent_roles', {}),
            'confidence':  topology.get('confidence', 0),
            'source':      topology.get('source', ''),
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# ============================================
# BRAIN STATUS
# ============================================
@app.route('/api/brain/status', methods=['GET'])
def brain_status():
    return jsonify(orchestrator.brain.get_brain_status())


# ============================================
# SELF-EVAL STATUS (NEW)
# ============================================
@app.route('/api/brain/eval-stats', methods=['GET'])
def eval_stats():
    if orchestrator.self_evaluator_enabled:
        return jsonify(orchestrator.self_evaluator.stats())
    return jsonify({'error': 'Self evaluator not enabled'}), 404


# ============================================
# HELPER FUNCTIONS
# ============================================

def generate_fallback_explanation(problem, arch, params):
    ops    = [op['operation'] for cell in arch
              for op in cell.get('operations', [])]
    top_op = max(set(ops), key=ops.count) if ops else 'conv5x5'
    return (
        f'AutoArchitect designed a custom model with {params:,} '
        f'parameters for: "{problem}". '
        f'Discovered {top_op} operations work best. '
        f'Expected accuracy: 70-85%.'
    )


def _run_script(problem: str) -> str:
    return '''#!/usr/bin/env python3
# ============================================
# AutoArchitect — Full Multi-Agent NAS Pipeline
# Problem: "''' + problem + '''"
#
# Usage:
#   python run_nas.py
#   python run_nas.py --problem "detect fraud in transactions"
#   python run_nas.py --problem "classify spam messages" --epochs 5
# ============================================

import sys, os, argparse, time
sys.path.append(os.path.dirname(__file__))
from dotenv import load_dotenv
load_dotenv()


def main():
    parser = argparse.ArgumentParser(description="AutoArchitect Multi-Agent NAS")
    parser.add_argument("--problem", type=str,
                        default="''' + problem + '''",
                        help="Describe your problem in plain English")
    parser.add_argument("--epochs", type=int, default=3,
                        help="Training epochs (default: 3)")
    args = parser.parse_args()

    problem = args.problem
    epochs  = args.epochs

    print("=" * 60)
    print("  AutoArchitect - Self-Learning Multi-Agent NAS")
    print("  Oakland Research Showcase 2026")
    print("=" * 60)
    print(f"  Problem : {problem}")
    print(f"  Epochs  : {epochs}")
    print("=" * 60)

    start = time.time()

    # Step 1: Classify problem
    print("\\n[1/5] Classifying problem with BERT...")
    try:
        from api.analyzer import ProblemAnalyzer
        analysis = ProblemAnalyzer().analyze(problem)
        category = analysis.get("category", "image")
        print(f"      Category: {category} ({analysis.get('confidence',0)}%)")
    except Exception as e:
        print(f"      Defaulting to image ({e})")
        category = "image"

    # Step 2: Fetch smart dataset
    print("\\n[2/5] Fetching best dataset from HuggingFace...")
    from api.dataset_fetcher import fetch_dataset
    data = fetch_dataset(problem, category, subset_size=2000)
    print(f"      Dataset : {data['name']} ({'REAL' if data.get('real_dataset') else 'generic'})")
    print(f"      Samples : {data['train_size']} train, {data['test_size']} test")
    print(f"      Classes : {data['classes']}")

    # Step 3: Train
    print("\\n[3/5] Training model...")
    import torch
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if category in ("image", "medical"):
        print("      Method: ResNet18 Transfer Learning")
        from api.transfer_trainer import train_transfer
        tr        = train_transfer(problem, data, epochs=epochs, device=device)
        model     = tr["model"]
        train_acc = tr["train_accuracy"]
        test_acc  = tr["test_accuracy"]
        params    = tr["parameters"]
        method    = "ResNet18 Transfer Learning"
    else:
        print("      Method: DARTS NAS")
        import torch.nn as nn
        import torch.optim as optim
        from api.nas_engine import DARTSNet
        model     = DARTSNet(C=16, num_cells=3,
                             num_classes=data["num_classes"]).to(device)
        params    = sum(p.numel() for p in model.parameters())
        opt       = optim.Adam(
            [p for n,p in model.named_parameters() if "arch_weights" not in n],
            lr=0.001)
        criterion = nn.CrossEntropyLoss()
        acc = 0
        for epoch in range(epochs):
            model.train()
            correct = total = 0
            for imgs, labels in data["train_loader"]:
                if imgs.shape[1] == 1: imgs = imgs.repeat(1,3,1,1)
                imgs, labels = imgs.to(device), labels.to(device)
                opt.zero_grad()
                out  = model(imgs)
                loss = criterion(out, labels)
                loss.backward()
                opt.step()
                correct += (out.argmax(1) == labels).sum().item()
                total   += labels.size(0)
            acc = round(100 * correct / total, 2)
            print(f"      Epoch {epoch+1}/{epochs} -> {acc}%")
        model.eval()
        correct = total = 0
        with torch.no_grad():
            for imgs, labels in data["test_loader"]:
                if imgs.shape[1] == 1: imgs = imgs.repeat(1,3,1,1)
                imgs, labels = imgs.to(device), labels.to(device)
                correct += (model(imgs).argmax(1) == labels).sum().item()
                total   += labels.size(0)
        train_acc = acc
        test_acc  = round(100 * correct / total, 2)
        method    = "DARTS NAS"

    # Step 4: Evaluate architecture
    print("\\n[4/5] Evaluating architecture quality...")
    from api.nas_engine             import run_quick_nas
    from api.agents.fusion_agent    import FusionAgent
    from api.agents.evaluator_agent import EvaluatorAgent
    nas        = run_quick_nas(num_classes=data["num_classes"])
    fused      = FusionAgent().fuse(
        [{"domain": category, "architecture": nas["architecture"],
          "parameters": nas["parameters"], "search_time": nas["search_time"]}],
        problem)
    evaluation = EvaluatorAgent().evaluate(fused, problem)
    score      = evaluation["avg_score"]
    verdict    = evaluation["verdict"]
    print(f"      Architecture quality: {score}% ({verdict})")

    # Step 5: Save model
    print("\\n[5/5] Saving trained model...")
    os.makedirs("output_models", exist_ok=True)
    safe      = problem[:30].replace(" ", "_").lower()
    save_path = f"output_models/{safe}_model.pth"
    torch.save(model.state_dict(), save_path)
    print(f"      Saved: {save_path}")

    # Summary
    elapsed = round(time.time() - start, 1)
    print()
    print("=" * 60)
    print("  RESULTS")
    print("=" * 60)
    print(f"  Problem    : {problem}")
    print(f"  Category   : {category}")
    print(f"  Dataset    : {data['name']} ({'REAL' if data.get('real_dataset') else 'generic'})")
    print(f"  Method     : {method}")
    print(f"  Parameters : {params:,}")
    print(f"  Train acc  : {train_acc}%")
    print(f"  Test acc   : {test_acc}%")
    print(f"  Arch score : {score}% ({verdict})")
    print(f"  Time       : {elapsed}s")
    print(f"  Model      : {save_path}")
    print("=" * 60)
    print("  AutoArchitect AI - Oakland Research Showcase 2026")
    print("  AI beats human baseline by +22.33% at $0 cost")
    print("=" * 60)


if __name__ == "__main__":
    main()
'''


def _requirements() -> str:
    return '''torch
torchvision
numpy
scikit-learn
transformers
datasets
huggingface-hub
requests
python-dotenv
flask
flask-cors
Pillow
'''


def _readme(problem: str) -> str:
    return '''# AutoArchitect - Multi-Agent NAS Pipeline
## Problem: "''' + problem + '''"

## Quick Start

### 1. Install
pip install -r requirements.txt

### 2. Add Groq key (free at console.groq.com)
cp .env.template .env
# Edit .env → add your GROQ_API_KEY

### 3. Run
python run_nas.py

### 4. Custom problem
python run_nas.py --problem "detect defects in factory images"
python run_nas.py --problem "classify spam emails" --epochs 5

## What happens
1. BERT classifies your problem automatically
2. Fetches the RIGHT dataset from HuggingFace (free)
3. Runs DARTS Neural Architecture Search
4. Trains with ResNet18 transfer learning (images)
   or NAS (text/security)
5. Saves trained model to output_models/

## Research Results
| Metric        | Human  | AutoArchitect |
|---------------|--------|---------------|
| Accuracy      | 52.56% | 74.89%        |
| Parameters    | 1.6M   | 105K          |
| Cache speedup | 1x     | 2066x         |
| Cost          | $$$    | $0            |

## AutoArchitect AI - Oakland Research Showcase 2026
'''


# ============================================
# NON-BLOCKING JOB SYSTEM (Day 23)
# ============================================

@app.route('/api/orchestrate/start', methods=['POST'])
def orchestrate_start():
    """Start orchestration in background; return job_id immediately."""
    data          = request.json or {}
    problem       = data.get('problem', '').strip()
    training_mode = data.get('training_mode', 'balanced')
    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    if len(problem) < 10:
        return jsonify({'error': 'Please describe your classification task more specifically.'}), 400
    if len(problem) > 500:
        return jsonify({'error': 'Description too long. Keep under 500 characters.'}), 400
    _clf_kws = ['classify', 'detect', 'predict', 'identify', 'recognize',
                'categorize', 'sort', 'filter', 'flag', 'grade', 'score',
                'diagnose', 'analyze', 'monitor', 'track']
    if not any(kw in problem.lower() for kw in _clf_kws):
        print(f"[WARN] Vague input — may produce poor results: '{problem[:60]}'", flush=True)

    job_id = str(uuid.uuid4())[:8]
    _jobs[job_id] = {
        'status':           'running',
        'result':           None,
        'error':            None,
        'started_at':       time.time(),
        'elapsed':          0,
        'training_metrics': None,
        'training_mode':    training_mode,
    }
    print(f"\n[REQUEST] New build: '{problem[:60]}' job_id={job_id} mode={training_mode}", flush=True)

    def _run():
        def _metrics_cb(step, total, message, **kwargs):
            phase = kwargs.get('phase', (message[:50] if message else 'Training'))
            _jobs[job_id]['training_metrics'] = {
                'phase':            phase,
                'current_epoch':    kwargs.get('epoch', 0),
                'total_epochs':     kwargs.get('total_epochs', 0),
                'current_loss':     kwargs.get('loss'),
                'current_accuracy': kwargs.get('accuracy'),
                'timestamp':        time.time(),
            }
            print(f"  [TRAIN] phase={phase} "
                  f"epoch={kwargs.get('epoch',0)}/{kwargs.get('total_epochs',0)} "
                  f"loss={kwargs.get('loss','?')} acc={kwargs.get('accuracy','?')}",
                  flush=True)

        print(f"[SOLVE] Starting orchestrator for job {job_id}", flush=True)
        try:
            result = orchestrator.solve(problem=problem, progress_callback=_metrics_cb,
                                        training_mode=training_mode)
            _jobs[job_id]['result'] = result
            _jobs[job_id]['status'] = 'complete'
            elapsed = round(time.time() - _jobs[job_id]['started_at'], 1)
            acc     = result.get('avg_accuracy') or result.get('test_accuracy', 0)
            print(f"[COMPLETE] job {job_id} done in {elapsed}s  "
                  f"accuracy={acc}%  from_cache={result.get('from_cache', False)}",
                  flush=True)
        except Exception as e:
            _jobs[job_id]['error']  = str(e)
            _jobs[job_id]['status'] = 'error'
            print(f"[ERROR] job {job_id} failed: {e}", flush=True)
            import traceback
            traceback.print_exc()
        finally:
            _jobs[job_id]['elapsed'] = round(
                time.time() - _jobs[job_id]['started_at'], 1)
            if _jobs[job_id].get('status') == 'running':
                _jobs[job_id]['status'] = 'error'
                _jobs[job_id].setdefault('error', 'Job ended without completing')
                print(f"[ERROR] job {job_id} forced out of running state in finally",
                      flush=True)

    threading.Thread(target=_run, daemon=True).start()
    return jsonify({'job_id': job_id})


@app.route('/api/status/<job_id>', methods=['GET'])
def job_status(job_id):
    """Poll job status. Returns result once status === 'complete'."""
    job = _jobs.get(job_id)
    if not job:
        # Return 200 with status='error' (not 404) so the frontend polling
        # loop breaks cleanly via the existing data.status === 'error' handler.
        # 404 leaves data.status undefined and the while(true) loop runs forever.
        return jsonify({
            'status': 'error',
            'error':  'Job expired — the server was restarted. Please submit your request again.',
        })
    elapsed = round(time.time() - job['started_at'], 1)
    status  = job['status']
    if status == 'complete' and not job.get('_status_logged'):
        job['_status_logged'] = True
        acc = (job['result'] or {}).get('avg_accuracy') or \
              (job['result'] or {}).get('test_accuracy', 0)
        print(f"[STATUS] job {job_id} complete result returned to frontend  "
              f"elapsed={elapsed}s accuracy={acc}%", flush=True)
    return jsonify({
        'status':           status,
        'elapsed':          elapsed,
        'result':           job['result'],
        'error':            job['error'],
        'training_metrics': job.get('training_metrics'),
    })


if __name__ == '__main__':
    port       = int(os.environ.get("PORT", 7860))
    debug_mode = os.environ.get("FLASK_DEBUG", "false").lower() == "true"

    print(f"[STARTUP] AutoArchitect AI is running!")
    print(f"[STARTUP] Binding to 0.0.0.0:{port}  debug={debug_mode}")

    app.run(
        host="0.0.0.0",
        port=port,
        debug=debug_mode,
        use_reloader=False,
    )