#!/usr/bin/env python3
"""llama.cpp command generator - main Flask application."""

import os
import sys
import json
import re
import math
import functools
import urllib.request
from flask import Flask, request, jsonify, send_from_directory, session, redirect
from flask_cors import CORS
from db import get_db, init_db, DB_PATH

app = Flask(__name__, static_folder='static', static_url_path='')
app.secret_key = 'llama-cmd-gen-secret-key-2026'
CORS(app)


# ==================== Admin Auth ====================

def admin_login_required(f):
    @functools.wraps(f)
    def wrapped(*args, **kwargs):
        if not session.get('admin_logged_in'):
            return jsonify({'error': '未登录或权限不足'}), 401
        return f(*args, **kwargs)
    return wrapped


# ==================== Helper Functions ====================

def parse_param_value(param_type, value):
    """Parse parameter value based on its type."""
    if value is None or value == '':
        return None
    if param_type == 'number':
        try:
            if '.' in str(value):
                return float(value)
            return int(value)
        except (ValueError, TypeError):
            return value
    elif param_type == 'boolean':
        return str(value).lower() in ('true', '1', 'yes', 'on')
    return value


def format_param_value(param_type, value):
    """Format parameter value for display."""
    if param_type == 'number':
        try:
            f = float(value)
            if f == int(f):
                return str(int(f))
            return str(f)
        except (ValueError, TypeError):
            return str(value)
    return str(value)


def get_flag_for_param(param):
    """Get the flag string for a parameter (prefers short flag)."""
    if param['short_flag']:
        return param['short_flag']
    return param['long_flag']


def get_kv_cache_bytes_per_element(cache_type):
    """Get bytes per element for KV cache type."""
    cache_type_map = {
        'f32': 4,
        'f16': 2,
        'bf16': 2,
        'q8_0': 1,
        'q4_0': 0.5,
        'q4_1': 0.5625,
        'iq4_nl': 0.5,
        'q5_0': 0.625,
        'q5_1': 0.6875,
    }
    return cache_type_map.get(cache_type, 2)


def estimate_vram(params_dict, gpus, version_params):
    """
    Estimate VRAM usage based on parameters.
    Returns dict with breakdown.
    """
    # Model parameters (user configurable hints)
    model_size_gb = float(params_dict.get('_model_size_gb', 0) or 0)
    model_layers = int(params_dict.get('_model_layers', 0) or 0)
    model_embd = int(params_dict.get('_model_embd', 0) or 0)
    model_heads = int(params_dict.get('_model_heads', 0) or 0)
    model_kv_heads = int(params_dict.get('_model_kv_heads', 0) or 0)
    model_head_dim = int(params_dict.get('_model_head_dim', 0) or 0)

    # If model_layers not specified, try to estimate from model size
    if model_layers == 0 and model_size_gb > 0:
        # Rough estimate: ~0.3GB per layer for 7B, scale accordingly
        model_layers = max(1, int(model_size_gb / 0.4))

    # n_gpu_layers
    ngl_raw = params_dict.get('n_gpu_layers', 'auto')
    if ngl_raw in ('auto', 'all', '-1'):
        ngl = model_layers if model_layers > 0 else 32
        if ngl_raw == '0':
            ngl = 0
    else:
        try:
            ngl = int(ngl_raw)
        except (ValueError, TypeError):
            ngl = 0

    # Context size
    ctx_size = int(params_dict.get('ctx_size', 0) or 0)
    if ctx_size == 0:
        ctx_size = 4096  # default assumption

    # Batch size
    batch_size = int(params_dict.get('batch_size', 2048) or 2048)
    ubatch_size = int(params_dict.get('ubatch_size', 512) or 512)

    # KV cache types
    ctk = params_dict.get('cache_type_k', 'f16')
    ctv = params_dict.get('cache_type_v', 'f16')
    k_bytes = get_kv_cache_bytes_per_element(ctk)
    v_bytes = get_kv_cache_bytes_per_element(ctv)

    # Parallel slots
    parallel = int(params_dict.get('parallel', 1) or 1)
    if parallel < 1:
        parallel = 1

    # Flash attention
    fa = params_dict.get('flash_attn', 'auto')
    fa_enabled = fa in ('on', 'auto')

    # === Calculate VRAM components ===
    total_vram = sum(g['vram_mb'] for g in gpus) if gpus else 0
    max_single_vram = max((g['vram_mb'] for g in gpus), default=0)

    # 1. Model weights in VRAM
    if model_size_gb > 0 and model_layers > 0:
        ratio = min(1.0, ngl / model_layers) if ngl > 0 else 0
        weights_vram_mb = model_size_gb * 1024 * ratio
    else:
        # Fallback estimate: ~1GB per 1B params at Q4
        weights_vram_mb = 0

    # 2. KV cache
    # KV cache per layer = 2 (K and V) * n_kv_heads * head_dim * ctx_size * bytes_per_element
    if model_kv_heads > 0 and model_head_dim > 0:
        kv_per_layer_bytes = 2 * model_kv_heads * model_head_dim * ctx_size
        kv_total_bytes = kv_per_layer_bytes * model_layers * parallel
    elif model_embd > 0:
        # Fallback: use embedding dim
        kv_per_layer_bytes = 2 * model_embd * ctx_size
        kv_total_bytes = kv_per_layer_bytes * model_layers * parallel
    else:
        # Very rough estimate: ~0.5MB per layer per 1K context at f16
        kv_total_bytes = 0
        if model_layers > 0:
            kv_total_bytes = int(0.5 * 1024 * 1024 * model_layers * (ctx_size / 1024) * parallel)

    kv_vram_mb = (kv_total_bytes * (k_bytes + v_bytes) / 2) / (1024 * 1024)

    # If ngl < model_layers, only ngl layers' KV is on GPU
    if model_layers > 0 and ngl < model_layers:
        kv_vram_mb = kv_vram_mb * (ngl / model_layers)

    # 3. Compute buffer / overhead
    # Rough: batch_size * model_embd * 4 bytes * some factor
    compute_mb = 0
    if model_embd > 0:
        compute_mb = (batch_size * model_embd * 4 * 2) / (1024 * 1024)  # logits buffer
    compute_mb = max(compute_mb, 100)  # minimum overhead

    # CUDA context overhead (~300-500MB per GPU)
    cuda_overhead_mb = len(gpus) * 400 if gpus else 0

    # Total
    total_estimate = weights_vram_mb + kv_vram_mb + compute_mb + cuda_overhead_mb

    return {
        'total_mb': round(total_estimate, 1),
        'total_gb': round(total_estimate / 1024, 2),
        'weights_mb': round(weights_vram_mb, 1),
        'weights_gb': round(weights_vram_mb / 1024, 2),
        'kv_cache_mb': round(kv_vram_mb, 1),
        'kv_cache_gb': round(kv_vram_mb / 1024, 2),
        'compute_mb': round(compute_mb, 1),
        'cuda_overhead_mb': round(cuda_overhead_mb, 1),
        'total_vram_available_mb': total_vram,
        'total_vram_available_gb': round(total_vram / 1024, 2) if total_vram else 0,
        'max_single_vram_mb': max_single_vram,
        'fits': total_estimate <= total_vram if total_vram > 0 else None,
        'usage_percent': round(total_estimate / total_vram * 100, 1) if total_vram > 0 else None,
        # For CPU+GPU mode
        'cpu_weights_mb': round(model_size_gb * 1024 * max(0, 1 - (ngl / model_layers if model_layers > 0 else 0)), 1) if model_size_gb > 0 else 0,
    }


def parse_natural_language(text):
    """
    Parse natural language description into parameter values.
    Supports Chinese and English keywords.
    """
    result = {}
    text_lower = text.lower()

    # GPU selection
    gpu_patterns = [
        (r'(?:rtx\s*)?3090', 'RTX 3090'),
        (r'(?:rtx\s*)?4090', 'RTX 4090'),
        (r'(?:rtx\s*)?4080', 'RTX 4080'),
        (r'(?:rtx\s*)?3080', 'RTX 3080'),
        (r'(?:rtx\s*)?5090', 'RTX 5090'),
        (r'(?:rtx\s*)?4070\s*ti', 'RTX 4070 Ti'),
        (r'(?:rtx\s*)?4060', 'RTX 4060'),
        (r'a100\s*80', 'A100 80GB'),
        (r'a100\s*40', 'A100 40GB'),
        (r'h100', 'H100 80GB'),
        (r'v100', 'V100 32GB'),
        (r'a6000', 'RTX A6000'),
        (r'a5000', 'RTX A5000'),
        (r'7900\s*xtx', 'RX 7900 XTX'),
    ]
    for pattern, gpu_name in gpu_patterns:
        if re.search(pattern, text_lower):
            result['_gpu_name'] = gpu_name
            break

    # GPU count
    gpu_count_match = re.search(r'(\d+)\s*(?:张|块|个)?\s*(?:gpu|显卡|卡)', text_lower)
    if gpu_count_match:
        result['_gpu_count'] = int(gpu_count_match.group(1))

    # Context size
    ctx_patterns = [
        r'(?:上下文|context|ctx)[\s大小为]*[:：\s]*(\d+)',
        r'(\d+)\s*(?:k|K)\s*(?:上下文|context|ctx)',
    ]
    for pattern in ctx_patterns:
        m = re.search(pattern, text_lower)
        if m:
            val = int(m.group(1))
            if val < 100:  # like "8k context"
                val = val * 1024
            result['ctx_size'] = str(val)
            break

    # GPU layers
    ngl_patterns = [
        r'(?:gpu\s*层|gpu\s*layers?|ngl|offload)[\s:：]*(\d+)',
        r'(\d+)\s*(?:层|layers?)\s*(?:gpu|offload)',
        r'全部(?:offload|卸载|gpu)|all\s*(?:gpu|layers?)',
    ]
    for pattern in ngl_patterns:
        m = re.search(pattern, text_lower)
        if m:
            val = m.group(1) if m.lastindex else '-1'
            result['n_gpu_layers'] = val
            break

    # Temperature
    temp_match = re.search(r'(?:温度|temp|temperature)[\s:：]*(\d+(?:\.\d+)?)', text_lower)
    if temp_match:
        result['temperature'] = temp_match.group(1)

    # Top-k
    topk_match = re.search(r'(?:top[\s-]*k)[\s:：]*(\d+)', text_lower)
    if topk_match:
        result['top_k'] = topk_match.group(1)

    # Top-p
    topp_match = re.search(r'(?:top[\s-]*p)[\s:：]*(\d+(?:\.\d+)?)', text_lower)
    if topp_match:
        result['top_p'] = topp_match.group(1)

    # Threads
    threads_match = re.search(r'(?:线程|threads?|cpu\s*线程)[\s:：]*(\d+)', text_lower)
    if threads_match:
        result['threads'] = threads_match.group(1)

    # Batch size
    batch_match = re.search(r'(?:batch|批处理|批次)[\s大小:：]*(\d+)', text_lower)
    if batch_match:
        result['batch_size'] = batch_match.group(1)

    # Port
    port_match = re.search(r'(?:端口|port)[\s:：]*(\d+)', text_lower)
    if port_match:
        result['port'] = port_match.group(1)

    # Model path
    model_match = re.search(r'(?:模型|model)[\s路径]*[:：\s]+([^\s,，]+\.gguf)', text_lower)
    if model_match:
        result['model'] = model_match.group(1)

    # HF repo
    hf_match = re.search(r'(?:hf|hugging\s*face|仓库)[\s:：]+([^\s,，]+)', text_lower)
    if hf_match and not model_match:
        result['hf_repo'] = hf_match.group(1)

    # Flash attention
    if re.search(r'flash\s*atten|flash\s*attn|fa', text_lower):
        result['flash_attn'] = 'on'

    # Mode
    if re.search(r'cpu\s*\+\s*gpu|gpu\s*\+\s*cpu|混合', text_lower):
        result['_mode'] = 'gpu_cpu'
    elif re.search(r'纯\s*gpu|gpu\s*only|仅\s*gpu', text_lower):
        result['_mode'] = 'gpu'

    # Parallel
    parallel_match = re.search(r'(?:并行|parallel|slots?)[\s:：]*(\d+)', text_lower)
    if parallel_match:
        result['parallel'] = parallel_match.group(1)

    # Split mode
    if re.search(r'row\s*split|行分割|按行分割', text_lower):
        result['split_mode'] = 'row'
    elif re.search(r'tensor\s*split|张量分割|按张量分割', text_lower):
        result['split_mode'] = 'tensor'

    # mlock
    if re.search(r'mlock|锁[\s定]*内存|内存锁', text_lower):
        result['mlock'] = 'true'

    # numa
    numa_match = re.search(r'numa[\s:：]*(distribute|isolate|numactl)', text_lower)
    if numa_match:
        result['numa'] = numa_match.group(1)

    # Reasoning budget
    reasoning_match = re.search(r'(?:reasoning|推理|thinking)[\s预算budget]*[:：\s]*(\d+)', text_lower)
    if reasoning_match:
        result['reasoning_budget'] = reasoning_match.group(1)

    return result


# ==================== API Routes ====================

@app.route('/')
def index():
    return send_from_directory('static', 'index.html')


@app.route('/admin')
def admin():
    # Serve admin page; JS handles login check
    return send_from_directory('static', 'admin.html')


# ----- Admin Login/Logout -----
@app.route('/api/admin/login', methods=['POST'])
def admin_login():
    data = request.json
    password = data.get('password', '')
    db = get_db()
    setting = db.execute('SELECT value FROM settings WHERE key = ?', ('admin_password',)).fetchone()
    db.close()
    stored_password = setting['value'] if setting else 'admin123'
    if password == stored_password:
        session['admin_logged_in'] = True
        return jsonify({'status': 'ok'})
    return jsonify({'error': '密码错误'}), 401

@app.route('/api/admin/logout', methods=['POST'])
def admin_logout():
    session.pop('admin_logged_in', None)
    return jsonify({'status': 'ok'})

@app.route('/api/admin/check')
def admin_check():
    return jsonify({'logged_in': session.get('admin_logged_in', False)})


# ----- Versions -----
@app.route('/api/versions')
def get_versions():
    db = get_db()
    versions = db.execute('SELECT * FROM llama_versions WHERE is_active = 1 ORDER BY sort_order').fetchall()
    result = [dict(v) for v in versions]
    db.close()
    return jsonify(result)


@app.route('/api/versions/<int:vid>/params')
def get_version_params(vid):
    db = get_db()
    params = db.execute(
        'SELECT * FROM params WHERE version_id = ? ORDER BY is_important DESC, sort_order',
        (vid,)
    ).fetchall()
    result = []
    for p in params:
        d = dict(p)
        if d.get('options'):
            try:
                d['options'] = json.loads(d['options'])
            except (json.JSONDecodeError, TypeError):
                pass
        result.append(d)
    db.close()
    return jsonify(result)


# ----- GPUs -----
@app.route('/api/gpus')
def get_gpus():
    db = get_db()
    gpus = db.execute('SELECT * FROM gpus ORDER BY sort_order, name').fetchall()
    result = [dict(g) for g in gpus]
    db.close()
    return jsonify(result)


# ----- Generate Command -----
@app.route('/api/generate', methods=['POST'])
def generate_command():
    data = request.json
    version_id = data.get('version_id')
    params = data.get('params', {})
    mode = data.get('mode', 'gpu')
    gpu_selections = data.get('gpu_selections', [])
    binary = data.get('binary', 'llama-server')
    system_memory_gb = data.get('system_memory_gb', 0)

    db = get_db()

    # Get param definitions
    param_defs = db.execute(
        'SELECT * FROM params WHERE version_id = ?',
        (version_id,)
    ).fetchall()
    param_def_map = {p['param_key']: dict(p) for p in param_defs}

    db.close()

    # Build command
    cmd_parts = [binary]

    # GPU layers
    if mode == 'gpu':
        # Add GPU-related args
        if gpu_selections:
            gpu_names = [g.get('name', '') for g in gpu_selections]
            device_str = ','.join(str(i) for i in range(len(gpu_selections)))

            # n_gpu_layers
            ngl = params.get('n_gpu_layers', 'auto')
            if ngl and ngl != param_def_map.get('n_gpu_layers', {}).get('default_value', 'auto'):
                flag = param_def_map.get('n_gpu_layers', {}).get('short_flag') or param_def_map.get('n_gpu_layers', {}).get('long_flag', '-ngl')
                cmd_parts.append(f'{flag} {ngl}')

            # tensor_split for multiple GPUs
            if len(gpu_selections) > 1:
                ts = params.get('tensor_split', '')
                if ts and ts != param_def_map.get('tensor_split', {}).get('default_value', ''):
                    flag = param_def_map.get('tensor_split', {}).get('short_flag', '-ts')
                    cmd_parts.append(f'{flag} {ts}')
                elif not ts:
                    # Auto-generate tensor split based on VRAM ratio
                    total_vram = sum(g.get('vram_mb', 0) for g in gpu_selections)
                    if total_vram > 0:
                        ratios = [str(round(g.get('vram_mb', 0) / total_vram, 2)) for g in gpu_selections]
                        cmd_parts.append(f'-ts {",".join(ratios)}')

    elif mode == 'gpu_cpu':
        # GPU+CPU mode
        if gpu_selections:
            ngl = params.get('n_gpu_layers', 'auto')
            if ngl and ngl != param_def_map.get('n_gpu_layers', {}).get('default_value', 'auto'):
                flag = param_def_map.get('n_gpu_layers', {}).get('short_flag') or param_def_map.get('n_gpu_layers', {}).get('long_flag', '-ngl')
                cmd_parts.append(f'{flag} {ngl}')

    # Iterate through params and add non-default ones
    skip_params = {'n_gpu_layers', 'tensor_split'}
    for key, value in params.items():
        if key.startswith('_'):
            continue
        if key in skip_params:
            continue

        p_def = param_def_map.get(key)
        if not p_def:
            continue

        default_val = p_def['default_value']
        # Skip if value equals default
        if str(value) == str(default_val):
            continue

        # Skip empty values
        if value is None or value == '' or value == 'false':
            if str(default_val).lower() == 'false' and str(value).lower() == 'false':
                continue
            if value == '' or value is None:
                continue

        # Boolean params: only add if true (and default is false)
        if p_def['param_type'] == 'boolean':
            if str(value).lower() == 'true' and str(default_val).lower() != 'true':
                flag = p_def['short_flag'] or p_def['long_flag']
                cmd_parts.append(flag)
            elif str(value).lower() == 'false' and str(default_val).lower() == 'true':
                # Add --no- variant
                flag = p_def['long_flag']
                cmd_parts.append(f'--no-{flag.lstrip("--")}')
            continue

        flag = p_def['short_flag'] or p_def['long_flag']
        cmd_parts.append(f'{flag} {value}')

    command = ' '.join(cmd_parts)
    return jsonify({'command': command, 'mode': mode, 'binary': binary})


# ----- Estimate VRAM -----
@app.route('/api/estimate', methods=['POST'])
def estimate():
    data = request.json
    params = data.get('params', {})
    gpu_selections = data.get('gpu_selections', [])
    mode = data.get('mode', 'gpu')
    system_memory_gb = data.get('system_memory_gb', 0)

    # Prepare GPU list with VRAM
    gpus = []
    for gs in gpu_selections:
        gpus.append({'vram_mb': gs.get('vram_mb', 0), 'name': gs.get('name', '')})

    result = estimate_vram(params, gpus, [])

    # Add system memory estimation for GPU+CPU mode
    if mode == 'gpu_cpu':
        model_size_gb = float(params.get('_model_size_gb', 0) or 0)
        model_layers = int(params.get('_model_layers', 0) or 0)
        ngl_raw = params.get('n_gpu_layers', 'auto')

        if ngl_raw in ('auto', 'all', '-1'):
            ngl = model_layers if model_layers > 0 else 32
        else:
            try:
                ngl = int(ngl_raw)
            except (ValueError, TypeError):
                ngl = 0

        # CPU portion of weights
        if model_layers > 0 and ngl < model_layers:
            cpu_weights_gb = model_size_gb * (1 - ngl / model_layers)
        else:
            cpu_weights_gb = 0

        # KV cache on CPU
        ctx_size = int(params.get('ctx_size', 0) or 0)
        if ctx_size == 0:
            ctx_size = 4096
        parallel = int(params.get('parallel', 1) or 1)
        if parallel < 1:
            parallel = 1

        # Estimate KV cache on CPU (layers not on GPU)
        model_kv_heads = int(params.get('_model_kv_heads', 0) or 0)
        model_head_dim = int(params.get('_model_head_dim', 0) or 0)
        model_embd = int(params.get('_model_embd', 0) or 0)

        if model_layers > 0 and ngl < model_layers:
            remaining_layers = model_layers - ngl
            if model_kv_heads > 0 and model_head_dim > 0:
                kv_cpu_bytes = 2 * model_kv_heads * model_head_dim * ctx_size * remaining_layers * parallel
            elif model_embd > 0:
                kv_cpu_bytes = 2 * model_embd * ctx_size * remaining_layers * parallel
            else:
                kv_cpu_bytes = int(0.5 * 1024 * 1024 * remaining_layers * (ctx_size / 1024) * parallel) * 2
            kv_cpu_mb = kv_cpu_bytes * 2 / (1024 * 1024)  # Assuming f16
        else:
            kv_cpu_mb = 0

        total_cpu_mb = cpu_weights_gb * 1024 + kv_cpu_mb + 500  # overhead
        result['cpu_total_mb'] = round(total_cpu_mb, 1)
        result['cpu_total_gb'] = round(total_cpu_mb / 1024, 2)
        result['cpu_weights_gb'] = round(cpu_weights_gb, 2)
        result['cpu_kv_cache_mb'] = round(kv_cpu_mb, 1)
        result['system_memory_gb'] = system_memory_gb
        if system_memory_gb > 0:
            result['cpu_usage_percent'] = round(total_cpu_mb / (system_memory_gb * 1024) * 100, 1)

    return jsonify(result)


# ----- Models (public read) -----
@app.route('/api/models')
def get_models():
    db = get_db()
    models = db.execute('SELECT * FROM models ORDER BY sort_order, name').fetchall()
    result = [dict(m) for m in models]
    db.close()
    return jsonify(result)

@app.route('/api/models/grouped')
def get_models_grouped():
    db = get_db()
    models = db.execute('SELECT * FROM models ORDER BY base_model, sort_order, name').fetchall()
    db.close()
    grouped = {}
    for m in models:
        d = dict(m)
        base = d['base_model']
        if base not in grouped:
            grouped[base] = []
        grouped[base].append(d)
    # Get default quant from settings
    db2 = get_db()
    dq = db2.execute("SELECT value FROM settings WHERE key = 'default_quant'").fetchone()
    db2.close()
    default_quant = dq['value'] if dq else 'Q4_K_M'
    return jsonify({'models': grouped, 'default_quant': default_quant})


# ----- Parse Natural Language -----
@app.route('/api/parse-nl', methods=['POST'])
def parse_nl():
    data = request.json
    text = data.get('text', '')
    # Try LLM API first if enabled
    db = get_db()
    llm_enabled = db.execute("SELECT value FROM settings WHERE key = 'llm_enabled'").fetchone()
    if llm_enabled and llm_enabled['value'] == 'true':
        llm_url = db.execute("SELECT value FROM settings WHERE key = 'llm_api_url'").fetchone()
        llm_key = db.execute("SELECT value FROM settings WHERE key = 'llm_api_key'").fetchone()
        llm_model = db.execute("SELECT value FROM settings WHERE key = 'llm_api_model'").fetchone()
        llm_prompt = db.execute("SELECT value FROM settings WHERE key = 'llm_system_prompt'").fetchone()
        db.close()
        url = llm_url['value'] if llm_url else ''
        key = llm_key['value'] if llm_key else ''
        model = llm_model['value'] if llm_model else ''
        system_prompt = llm_prompt['value'] if llm_prompt else ''
        if url:
            try:
                result = call_llm_for_parsing(url, key, model, system_prompt, text)
                if result:
                    return jsonify(result)
            except Exception as e:
                print(f'LLM parse failed: {e}', file=sys.stderr)
    else:
        db.close()
    # Fallback to regex parsing
    result = parse_natural_language(text)
    return jsonify(result)


def call_llm_for_parsing(url, key, model, system_prompt, user_text):
    """Call LLM API to parse natural language into params."""
    headers = {'Content-Type': 'application/json'}
    if key:
        headers['Authorization'] = f'Bearer {key}'
    body = {
        'model': model,
        'messages': [
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': user_text}
        ],
        'temperature': 0.1,
        'max_tokens': 2000,
    }
    req = urllib.request.Request(url, data=json.dumps(body).encode('utf-8'), headers=headers, method='POST')
    with urllib.request.urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read().decode('utf-8'))
    # OpenAI-compatible response
    content = data['choices'][0]['message']['content']
    # Try to extract JSON from the response
    content = content.strip()
    if content.startswith('```'):
        content = re.sub(r'^```\w*\n?', '', content)
        content = re.sub(r'\n?```$', '', content)
    result = json.loads(content)
    return result


# ==================== Admin API ====================

    # All admin routes below require login

@app.route('/api/admin/gpus', methods=['GET', 'POST'])
@admin_login_required
def admin_gpus():
    db = get_db()
    if request.method == 'GET':
        gpus = db.execute('SELECT * FROM gpus ORDER BY sort_order, name').fetchall()
        result = [dict(g) for g in gpus]
        db.close()
        return jsonify(result)

    elif request.method == 'POST':
        data = request.json
        db.execute(
            'INSERT INTO gpus (name, vram_mb, compute_capability, description, sort_order) VALUES (?, ?, ?, ?, ?)',
            (data['name'], data['vram_mb'], data.get('compute_capability', ''),
             data.get('description', ''), data.get('sort_order', 0))
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/gpus/<int:gid>', methods=['PUT', 'DELETE'])
@admin_login_required
def admin_gpu_edit(gid):
    db = get_db()
    if request.method == 'PUT':
        data = request.json
        db.execute(
            'UPDATE gpus SET name=?, vram_mb=?, compute_capability=?, description=?, sort_order=? WHERE id=?',
            (data['name'], data['vram_mb'], data.get('compute_capability', ''),
             data.get('description', ''), data.get('sort_order', 0), gid)
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})
    elif request.method == 'DELETE':
        db.execute('DELETE FROM gpus WHERE id=?', (gid,))
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/versions', methods=['GET', 'POST'])
@admin_login_required
def admin_versions():
    db = get_db()
    if request.method == 'GET':
        versions = db.execute('SELECT * FROM llama_versions ORDER BY sort_order').fetchall()
        result = [dict(v) for v in versions]
        db.close()
        return jsonify(result)
    elif request.method == 'POST':
        data = request.json
        db.execute(
            'INSERT INTO llama_versions (version_tag, description, release_date, is_active, sort_order) VALUES (?, ?, ?, ?, ?)',
            (data['version_tag'], data.get('description', ''), data.get('release_date', ''),
             data.get('is_active', 1), data.get('sort_order', 0))
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/versions/<int:vid>', methods=['PUT', 'DELETE'])
@admin_login_required
def admin_version_edit(vid):
    db = get_db()
    if request.method == 'PUT':
        data = request.json
        db.execute(
            'UPDATE llama_versions SET version_tag=?, description=?, release_date=?, is_active=?, sort_order=? WHERE id=?',
            (data['version_tag'], data.get('description', ''), data.get('release_date', ''),
             data.get('is_active', 1), data.get('sort_order', 0), vid)
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})
    elif request.method == 'DELETE':
        db.execute('DELETE FROM llama_versions WHERE id=?', (vid,))
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/versions/<int:vid>/params', methods=['GET', 'POST'])
@admin_login_required
def admin_params(vid):
    db = get_db()
    if request.method == 'GET':
        params = db.execute(
            'SELECT * FROM params WHERE version_id = ? ORDER BY is_important DESC, sort_order',
            (vid,)
        ).fetchall()
        result = []
        for p in params:
            d = dict(p)
            if d.get('options'):
                try:
                    d['options'] = json.loads(d['options'])
                except (json.JSONDecodeError, TypeError):
                    pass
            result.append(d)
        db.close()
        return jsonify(result)
    elif request.method == 'POST':
        data = request.json
        options = data.get('options')
        if isinstance(options, list):
            options = json.dumps(options)
        db.execute(
            '''INSERT INTO params
            (version_id, param_key, short_flag, long_flag, description, category, param_type,
             default_value, options, min_value, max_value, step, unit, is_important, affects_vram, sort_order)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
            (vid, data['param_key'], data.get('short_flag', ''), data['long_flag'],
             data.get('description', ''), data.get('category', 'common'),
             data.get('param_type', 'string'), data.get('default_value', ''),
             options, data.get('min_value'), data.get('max_value'),
             data.get('step'), data.get('unit'), data.get('is_important', 0),
             data.get('affects_vram', 0), data.get('sort_order', 0))
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/params/<int:pid>', methods=['PUT', 'DELETE'])
@admin_login_required
def admin_param_edit(pid):
    db = get_db()
    if request.method == 'PUT':
        data = request.json
        options = data.get('options')
        if isinstance(options, list):
            options = json.dumps(options)
        db.execute(
            '''UPDATE params SET
            param_key=?, short_flag=?, long_flag=?, description=?, category=?, param_type=?,
            default_value=?, options=?, min_value=?, max_value=?, step=?, unit=?,
            is_important=?, affects_vram=?, sort_order=? WHERE id=?''',
            (data['param_key'], data.get('short_flag', ''), data['long_flag'],
             data.get('description', ''), data.get('category', 'common'),
             data.get('param_type', 'string'), data.get('default_value', ''),
             options, data.get('min_value'), data.get('max_value'),
             data.get('step'), data.get('unit'), data.get('is_important', 0),
             data.get('affects_vram', 0), data.get('sort_order', 0), pid)
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})
    elif request.method == 'DELETE':
        db.execute('DELETE FROM params WHERE id=?', (pid,))
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


@app.route('/api/admin/settings', methods=['GET', 'PUT'])
@admin_login_required
def admin_settings():
    db = get_db()
    if request.method == 'GET':
        settings = db.execute('SELECT * FROM settings').fetchall()
        result = {s['key']: s['value'] for s in settings}
        db.close()
        return jsonify(result)
    elif request.method == 'PUT':
        data = request.json
        for key, value in data.items():
            db.execute(
                'INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)',
                (key, str(value))
            )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


# ----- Admin Models CRUD -----
@app.route('/api/admin/models', methods=['POST'])
@admin_login_required
def admin_add_model():
    data = request.json
    db = get_db()
    db.execute(
        '''INSERT INTO models (base_model, name, size_gb, layers, embd, kv_heads, head_dim, attention_heads, quant, description, sort_order)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
        (data['base_model'], data['name'], data['size_gb'], data['layers'], data['embd'],
         data['kv_heads'], data['head_dim'], data['attention_heads'],
         data.get('quant', ''), data.get('description', ''), data.get('sort_order', 0))
    )
    db.commit()
    db.close()
    return jsonify({'status': 'ok'})


@app.route('/api/admin/models/<int:mid>', methods=['PUT', 'DELETE'])
@admin_login_required
def admin_model_edit(mid):
    db = get_db()
    if request.method == 'PUT':
        data = request.json
        db.execute(
            '''UPDATE models SET base_model=?, name=?, size_gb=?, layers=?, embd=?, kv_heads=?,
            head_dim=?, attention_heads=?, quant=?, description=?, sort_order=? WHERE id=?''',
            (data['base_model'], data['name'], data['size_gb'], data['layers'], data['embd'],
             data['kv_heads'], data['head_dim'], data['attention_heads'],
             data.get('quant', ''), data.get('description', ''), data.get('sort_order', 0), mid)
        )
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})
    elif request.method == 'DELETE':
        db.execute('DELETE FROM models WHERE id=?', (mid,))
        db.commit()
        db.close()
        return jsonify({'status': 'ok'})


if __name__ == '__main__':
    init_db()
    app.run(host='0.0.0.0', port=16052, debug=False)