#!/usr/bin/env python3
"""
IC Generate FastAPI — Single-file app combining:
  - Featherless API proxy with key rotation
  - Concurrent UI question & solution generator
  - SQLite database with crash recovery
  - REST API for management (Admin-protected)
  - API Key management routes
  - [ENGINE] Question generation engine with randomized prompt templates
  - [ENGINE] Runtime prompt management endpoints
Persistent storage: /data
  - /data/ic_data.db            — SQLite database
  - /data/keys.txt              — API keys (one per line)
  - /data/exported_code/        — File exports
  - /data/prompt_engine_state.json — [ENGINE] Prompt engine state
"""

# === Imports ===
import os
import re
import json
import time
import random
import sqlite3
import logging
import signal
import threading
import itertools
from typing import Optional, List, Dict, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI, Request, HTTPException, Query, Depends, Security, status
from fastapi.security import APIKeyHeader
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from pydantic import BaseModel
from dotenv import load_dotenv
from openai import OpenAI, APIError, APIConnectionError, RateLimitError

import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# NOW your imports will work:
from prompt_engine import QuestionGenerationEngine, PromptConfig

load_dotenv()

# === Configuration ===
DATA_DIR = os.environ.get("DATA_DIR", "/data")
os.makedirs(DATA_DIR, exist_ok=True)

DB_PATH = os.path.join(DATA_DIR, "ic_data.db")
KEYS_FILE = os.path.join(DATA_DIR, "keys.txt")
EXPORT_DIR = os.path.join(DATA_DIR, "exported_code")

# [ENGINE] State file for prompt engine persistence
PROMPT_STATE_FILE = os.path.join(DATA_DIR, "prompt_engine_state.json")

FEATHERLESS_API_BASE = os.environ.get(
    "FEATHERLESS_API_BASE", "https://api.featherless.ai/v1"
)
PORT = int(os.environ.get("PORT", "7860"))
STACK = "HTML/CSS/JS"

# [ENGINE] Initialize the global prompt engine
prompt_engine = QuestionGenerationEngine(
    state_file=PROMPT_STATE_FILE,
    min_temperature=float(os.environ.get("PROMPT_MIN_TEMP", "0.85")),
    max_temperature=float(os.environ.get("PROMPT_MAX_TEMP", "1.15")),
    max_tokens=4096,
)

# === Admin Authentication ===
ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "change-me-default-admin-key")
api_key_header_auth = APIKeyHeader(name="X-Admin-Key", auto_error=False)


async def verify_admin_key(api_key: str = Security(api_key_header_auth)):
    if api_key == ADMIN_API_KEY:
        return api_key
    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Invalid or missing Admin API Key. Provide it in the 'X-Admin-Key' header.",
    )

# === Logging ===
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [%(threadName)s]: %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

# === Shutdown Event ===
_shutdown = threading.Event()


def _handle_signal(sig, frame):
    logger.warning("⚠️  Shutdown requested. Finishing in-flight requests...")
    _shutdown.set()


signal.signal(signal.SIGINT, _handle_signal)
signal.signal(signal.SIGTERM, _handle_signal)

# === Database Layer ===
DB_LOCK = threading.Lock()


def init_db(db_path: str = DB_PATH):
    """Create all tables including pending_jobs for crash recovery."""
    os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
    conn = sqlite3.connect(db_path, timeout=60)
    conn.execute("PRAGMA foreign_keys = ON")
    conn.execute("PRAGMA journal_mode = DELETE")
    conn.execute("PRAGMA busy_timeout = 30000")
    c = conn.cursor()

    c.execute(
        """CREATE TABLE IF NOT EXISTS questions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        question_text TEXT NOT NULL,
        thinking_trace TEXT,
        full_response TEXT,
        model TEXT,
        finish_reason TEXT,
        usage_json TEXT,
        raw_chunks_json TEXT,
        chunk_count INTEGER,
        generation_time_s REAL,
        prompt_metadata TEXT,
        created_at TEXT NOT NULL DEFAULT (datetime('now'))
    )"""
    )

    c.execute(
        """CREATE TABLE IF NOT EXISTS solutions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        question_id INTEGER NOT NULL,
        stack TEXT NOT NULL DEFAULT 'HTML/CSS/JS',
        solution_code TEXT NOT NULL,
        thinking_trace TEXT,
        full_response TEXT,
        model TEXT,
        finish_reason TEXT,
        usage_json TEXT,
        raw_chunks_json TEXT,
        chunk_count INTEGER,
        generation_time_s REAL,
        created_at TEXT NOT NULL DEFAULT (datetime('now')),
        FOREIGN KEY(question_id) REFERENCES questions(id) ON DELETE CASCADE
    )"""
    )

    c.execute(
        """CREATE TABLE IF NOT EXISTS pending_jobs (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        question_id INTEGER NOT NULL,
        status TEXT NOT NULL DEFAULT 'pending',
        attempts INTEGER NOT NULL DEFAULT 0,
        last_error TEXT,
        started_at TEXT,
        completed_at TEXT,
        created_at TEXT NOT NULL DEFAULT (datetime('now')),
        FOREIGN KEY(question_id) REFERENCES questions(id) ON DELETE CASCADE
    )"""
    )

    c.execute(
        """CREATE TABLE IF NOT EXISTS run_log (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        started_at TEXT NOT NULL DEFAULT (datetime('now')),
        finished_at TEXT,
        q_model TEXT, s_model TEXT,
        requested INTEGER, completed INTEGER, failed INTEGER, pending INTEGER,
        config_json TEXT
    )"""
    )

    # Migrations
    _migrate_add_columns(
        c,
        "questions",
        {
            "thinking_trace": "TEXT",
            "full_response": "TEXT",
            "generation_time_s": "REAL",
            "prompt_metadata": "TEXT",  # [ENGINE] New column
        },
    )
    _migrate_add_columns(
        c,
        "solutions",
        {"thinking_trace": "TEXT", "full_response": "TEXT", "generation_time_s": "REAL"},
    )
    _migrate_add_columns(
        c,
        "pending_jobs",
        {
            "attempts": "INTEGER DEFAULT 0",
            "last_error": "TEXT",
            "started_at": "TEXT",
            "completed_at": "TEXT",
        },
    )

    conn.commit()
    conn.close()
    logger.info(f"Database ready: {db_path}")


def _migrate_add_columns(cursor, table, columns):
    for col, col_type in columns.items():
        try:
            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {col} {col_type}")
        except sqlite3.OperationalError:
            pass


def _safe_db_execute(db_path, operations, max_retries=5):
    for attempt in range(max_retries):
        try:
            with DB_LOCK:
                conn = sqlite3.connect(db_path, timeout=60)
                conn.execute("PRAGMA journal_mode = DELETE")
                conn.execute("PRAGMA busy_timeout = 30000")
                cur = conn.cursor()
                try:
                    result = operations(cur)
                    conn.commit()
                    return result
                except Exception:
                    conn.rollback()
                    raise
                finally:
                    conn.close()
        except sqlite3.OperationalError as e:
            if "locked" in str(e).lower() and attempt < max_retries - 1:
                wait = 0.5 * (2**attempt)
                logger.warning(
                    f"DB locked (attempt {attempt+1}/{max_retries}), retrying in {wait:.1f}s..."
                )
                time.sleep(wait)
            else:
                raise


def save_question(
    db_path, q_text, q_thinking, q_full, q_model, q_finish, q_usage, q_chunks, gen_time,
    prompt_metadata=None,  # [ENGINE] New parameter
):
    def ops(cur):
        cur.execute(
            """INSERT INTO questions
            (question_text, thinking_trace, full_response, model, finish_reason,
             usage_json, raw_chunks_json, chunk_count, generation_time_s, prompt_metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                q_text,
                q_thinking,
                q_full,
                q_model,
                q_finish,
                json.dumps(q_usage, ensure_ascii=False) if q_usage else None,
                json.dumps(q_chunks, ensure_ascii=False, default=str),
                len(q_chunks),
                gen_time,
                json.dumps(prompt_metadata, ensure_ascii=False) if prompt_metadata else None,  # [ENGINE]
            ),
        )
        question_id = cur.lastrowid
        cur.execute(
            "INSERT INTO pending_jobs (question_id, status) VALUES (?, 'pending')",
            (question_id,),
        )
        job_id = cur.lastrowid
        return job_id, question_id

    return _safe_db_execute(db_path, ops)


def save_solution(
    db_path,
    job_id,
    question_id,
    s_code,
    s_thinking,
    s_full,
    s_model,
    s_finish,
    s_usage,
    s_chunks,
    gen_time,
):
    def ops(cur):
        cur.execute(
            """INSERT INTO solutions
            (question_id, stack, solution_code, thinking_trace, full_response,
             model, finish_reason, usage_json, raw_chunks_json, chunk_count, generation_time_s)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                question_id,
                STACK,
                s_code,
                s_thinking,
                s_full,
                s_model,
                s_finish,
                json.dumps(s_usage, ensure_ascii=False) if s_usage else None,
                json.dumps(s_chunks, ensure_ascii=False, default=str),
                len(s_chunks),
                gen_time,
            ),
        )
        cur.execute(
            "UPDATE pending_jobs SET status='done', completed_at=datetime('now') WHERE id=?",
            (job_id,),
        )

    _safe_db_execute(db_path, ops)


def mark_job_started(db_path, job_id):
    def ops(cur):
        cur.execute(
            "UPDATE pending_jobs SET started_at=datetime('now'), attempts=attempts+1 WHERE id=?",
            (job_id,),
        )

    _safe_db_execute(db_path, ops)


def mark_job_failed(db_path, job_id, error_msg):
    def ops(cur):
        cur.execute(
            "UPDATE pending_jobs SET last_error=?, status='pending' WHERE id=?",
            (str(error_msg)[:500], job_id),
        )

    _safe_db_execute(db_path, ops)


def get_pending_jobs(db_path):
    conn = sqlite3.connect(db_path, timeout=30)
    cur = conn.cursor()
    cur.execute(
        """SELECT pj.id, pj.question_id, q.question_text, pj.attempts
                   FROM pending_jobs pj
                   JOIN questions q ON pj.question_id = q.id
                   WHERE pj.status = 'pending'
                   ORDER BY pj.id"""
    )
    rows = cur.fetchall()
    conn.close()
    return [(r[0], r[1], r[2], r[3]) for r in rows]


def log_run(
    db_path, q_model, s_model, requested, completed, failed, pending, config
):
    def ops(cur):
        cur.execute(
            """INSERT INTO run_log
            (q_model, s_model, requested, completed, failed, pending,
             finished_at, config_json)
            VALUES (?, ?, ?, ?, ?, ?, datetime('now'), ?)""",
            (
                q_model,
                s_model,
                requested,
                completed,
                failed,
                pending,
                json.dumps(config, ensure_ascii=False),
            ),
        )

    try:
        _safe_db_execute(db_path, ops)
    except Exception as e:
        logger.warning(f"Failed to log run: {e}")


def delete_question_db(db_path, question_id):
    def ops(cur):
        cur.execute("DELETE FROM questions WHERE id=?", (question_id,))
        return cur.rowcount > 0

    return _safe_db_execute(db_path, ops)


def get_stats_dict(db_path):
    if not os.path.exists(db_path):
        return {"error": f"Database {db_path} does not exist."}

    with sqlite3.connect(db_path) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM questions")
        q_count = cursor.fetchone()[0]
        cursor.execute("SELECT COUNT(*) FROM solutions")
        s_count = cursor.fetchone()[0]
        cursor.execute("SELECT COALESCE(SUM(chunk_count), 0) FROM questions")
        q_chunks = cursor.fetchone()[0]
        cursor.execute("SELECT COALESCE(SUM(chunk_count), 0) FROM solutions")
        s_chunks = cursor.fetchone()[0]
        cursor.execute(
            "SELECT COUNT(*) FROM questions WHERE thinking_trace IS NOT NULL"
        )
        q_with_thinking = cursor.fetchone()[0]
        cursor.execute(
            "SELECT COUNT(*) FROM solutions WHERE thinking_trace IS NOT NULL"
        )
        s_with_thinking = cursor.fetchone()[0]
        cursor.execute("SELECT COUNT(*) FROM pending_jobs WHERE status='pending'")
        pending = cursor.fetchone()[0]
        cursor.execute(
            "SELECT COALESCE(AVG(generation_time_s), 0) FROM questions WHERE generation_time_s IS NOT NULL"
        )
        avg_q_time = cursor.fetchone()[0]
        cursor.execute(
            "SELECT COALESCE(AVG(generation_time_s), 0) FROM solutions WHERE generation_time_s IS NOT NULL"
        )
        avg_s_time = cursor.fetchone()[0]

        cursor.execute(
            "SELECT id, question_text, model, created_at FROM questions ORDER BY id DESC LIMIT 5"
        )
        recent = cursor.fetchall()

        runs = []
        try:
            cursor.execute(
                "SELECT started_at, completed, failed, pending, q_model, s_model FROM run_log ORDER BY id DESC LIMIT 3"
            )
            runs = cursor.fetchall()
        except sqlite3.OperationalError:
            pass

    return {
        "questions": q_count,
        "questions_with_thinking": q_with_thinking,
        "solutions": s_count,
        "solutions_with_thinking": s_with_thinking,
        "pending_jobs": pending,
        "total_chunks": q_chunks + s_chunks,
        "question_chunks": q_chunks,
        "solution_chunks": s_chunks,
        "avg_question_time_s": round(avg_q_time, 1),
        "avg_solution_time_s": round(avg_s_time, 1),
        "recent_questions": [
            {"id": qid, "model": qm, "created_at": qc, "preview": qt[:80]}
            for qid, qt, qm, qc in reversed(recent)
        ],
        "recent_runs": [
            {
                "started_at": r[0],
                "completed": r[1],
                "failed": r[2],
                "pending": r[3],
                "q_model": r[4],
                "s_model": r[5],
            }
            for r in reversed(runs)
        ],
        # [ENGINE] Include prompt engine stats
        "prompt_engine_stats": prompt_engine.get_stats(),
    }


def export_data_json(db_path):
    if not os.path.exists(db_path):
        return []
    with sqlite3.connect(db_path) as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT q.id, q.question_text, q.thinking_trace, q.full_response,
                   q.model AS q_model, q.created_at, q.generation_time_s,
                   q.prompt_metadata,
                   s.id AS s_id, s.stack, s.solution_code, s.thinking_trace AS s_thinking,
                   s.full_response AS s_full,
                   s.model AS s_model, s.finish_reason, s.chunk_count,
                   s.created_at AS s_created, s.generation_time_s AS s_gen_time
            FROM questions q
            LEFT JOIN solutions s ON q.id = s.question_id
            ORDER BY q.id ASC
        """
        )
        rows = cursor.fetchall()

    questions_map = {}
    for (
        q_id,
        q_text,
        q_thinking,
        q_full,
        q_model,
        q_created,
        q_gen_time,
        q_prompt_meta,  # [ENGINE]
        s_id,
        stack,
        code,
        s_thinking,
        s_full,
        s_model,
        s_finish,
        s_chunks,
        s_created,
        s_gen_time,
    ) in rows:
        if q_id not in questions_map:
            questions_map[q_id] = {
                "id": q_id,
                "question": q_text,
                "thinking_trace": q_thinking,
                "full_response": q_full,
                "model": q_model,
                "created_at": q_created,
                "generation_time_s": q_gen_time,
                "prompt_metadata": json.loads(q_prompt_meta) if q_prompt_meta else None,  # [ENGINE]
                "solutions": [],
            }
        if s_id and code:
            questions_map[q_id]["solutions"].append(
                {
                    "id": s_id,
                    "stack": stack,
                    "code": code,
                    "thinking_trace": s_thinking,
                    "full_response": s_full,
                    "model": s_model,
                    "finish_reason": s_finish,
                    "chunk_count": s_chunks,
                    "created_at": s_created,
                    "generation_time_s": s_gen_time,
                }
            )
    return list(questions_map.values())


def export_files(db_path, output_dir):
    if not os.path.exists(db_path):
        return {"error": "Database does not exist."}
    os.makedirs(output_dir, exist_ok=True)
    with sqlite3.connect(db_path) as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT q.id, q.question_text, q.thinking_trace, q.full_response,
                   q.model AS q_model, q.created_at, q.generation_time_s,
                   q.prompt_metadata,
                   s.id AS s_id, s.stack, s.solution_code, s.thinking_trace AS s_thinking,
                   s.full_response AS s_full,
                   s.model AS s_model, s.finish_reason, s.chunk_count,
                   s.created_at AS s_created, s.generation_time_s AS s_gen_time
            FROM questions q
            LEFT JOIN solutions s ON q.id = s.question_id
            ORDER BY q.id ASC
        """
        )
        rows = cursor.fetchall()

    exported = 0
    for row in rows:
        (
            q_id,
            q_text,
            q_thinking,
            q_full,
            q_model,
            q_created,
            q_gen_time,
            q_prompt_meta,  # [ENGINE]
            s_id,
            stack,
            code,
            s_thinking,
            s_full,
            s_model,
            s_finish,
            s_chunks,
            s_created,
            s_gen_time,
        ) = row
        if not s_id or not code:
            continue
        safe_q = re.sub(r"[^\w\-_\. ]", "_", q_text)[:50].strip()
        q_folder = os.path.join(output_dir, f"Q{q_id}_{safe_q}")
        os.makedirs(q_folder, exist_ok=True)
        with open(os.path.join(q_folder, f"solution_{s_id}.html"), "w") as f:
            f.write(code)
        if q_thinking:
            with open(os.path.join(q_folder, f"q_thinking_{q_id}.txt"), "w") as f:
                f.write(q_thinking)
        if s_thinking:
            with open(os.path.join(q_folder, f"s_thinking_{s_id}.txt"), "w") as f:
                f.write(s_thinking)
        if q_full:
            with open(os.path.join(q_folder, f"q_full_response_{q_id}.txt"), "w") as f:
                f.write(q_full)
        if s_full:
            with open(os.path.join(q_folder, f"s_full_response_{s_id}.txt"), "w") as f:
                f.write(s_full)
        meta = {
            "question_id": q_id,
            "question": q_text,
            "q_model": q_model,
            "has_q_thinking": q_thinking is not None,
            "q_generation_time_s": q_gen_time,
            "prompt_metadata": json.loads(q_prompt_meta) if q_prompt_meta else None,  # [ENGINE]
            "solution_id": s_id,
            "stack": stack,
            "s_model": s_model,
            "has_s_thinking": s_thinking is not None,
            "finish_reason": s_finish,
            "chunk_count": s_chunks,
            "s_generation_time_s": s_gen_time,
            "q_created": q_created,
            "s_created": s_created,
        }
        with open(os.path.join(q_folder, f"meta_{s_id}.json"), "w") as f:
            json.dump(meta, f, indent=2, ensure_ascii=False)
        exported += 1

    return {"exported": exported, "output_dir": output_dir}


# === Key Rotation ===


def load_keys():
    """Load API keys from /data/keys.txt or environment variable."""
    keys = []
    if os.path.exists(KEYS_FILE):
        with open(KEYS_FILE, "r") as f:
            keys = [line.strip() for line in f if line.strip()]
    env_keys = os.environ.get("FEATHERLESS_API_KEYS", "")
    if env_keys:
        keys.extend([k.strip() for k in env_keys.split(",") if k.strip()])
    seen = set()
    unique_keys = []
    for k in keys:
        if k not in seen:
            seen.add(k)
            unique_keys.append(k)
    if not unique_keys:
        logger.warning(f"No keys found in {KEYS_FILE} or env. Using dummy key.")
        unique_keys = ["dummy_key"]
    return unique_keys


_keys = load_keys()
_key_cycle = itertools.cycle(_keys)
_key_lock = threading.Lock()


def get_next_key():
    with _key_lock:
        return next(_key_cycle)


def reload_keys():
    global _keys, _key_cycle
    _keys = load_keys()
    _key_cycle = itertools.cycle(_keys)
    return len(_keys)


# === Retry Logic ===


def retry_api_call(func, max_retries=5, delay=2):
    last_err = None
    for attempt in range(max_retries):
        if _shutdown.is_set():
            raise InterruptedError("Shutdown requested")
        try:
            return func()
        except RateLimitError as e:
            last_err = e
            backoff = 120.0 + random.uniform(0, 10)
            logger.warning(
                f"Rate limit (429). Waiting {backoff:.0f}s... (attempt {attempt+1}/{max_retries})"
            )
            time.sleep(backoff)
        except APIConnectionError as e:
            last_err = e
            backoff = min(60, delay * (2**attempt)) + random.uniform(0, 2)
            logger.warning(
                f"Connection error: {e}. Retry {attempt+1}/{max_retries} in {backoff:.1f}s..."
            )
            time.sleep(backoff)
        except APIError as e:
            last_err = e
            if hasattr(e, "status_code") and e.status_code == 429:
                backoff = 120.0 + random.uniform(0, 10)
                logger.warning(f"Rate limit (via APIError 429). Waiting {backoff:.0f}s...")
            elif hasattr(e, "status_code") and e.status_code in (500, 502, 503, 504):
                backoff = min(90, delay * (2**attempt)) + random.uniform(0, 5)
                logger.warning(
                    f"Server error {e.status_code}: {e}. Retry in {backoff:.1f}s..."
                )
            else:
                raise
            time.sleep(backoff)
        except Exception:
            raise
    raise last_err


# === Streaming Completion ===


def stream_completion(client, model, messages, temperature=0.2, max_tokens=16384):
    raw_chunks = []
    content_parts = []
    thinking_parts = []
    tool_call_parts = []
    function_call_parts = []
    finish_reason = None
    usage = None
    actual_model = model
    t0 = time.monotonic()

    def make_call():
        return client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
            stream_options={"include_usage": True},
        )

    stream = retry_api_call(make_call)

    for chunk in stream:
        if _shutdown.is_set():
            break
        try:
            chunk_dict = chunk.model_dump()
        except AttributeError:
            try:
                chunk_dict = json.loads(chunk.json())
            except Exception:
                chunk_dict = {"_raw": str(chunk)}
        raw_chunks.append(chunk_dict)

        if chunk.choices:
            choice = chunk.choices[0]
            delta = choice.delta
            if delta:
                if delta.content:
                    content_parts.append(delta.content)
                for attr in (
                    "reasoning_content",
                    "reasoning",
                    "thinking",
                    "thought",
                    "internal_thoughts",
                ):
                    val = getattr(delta, attr, None)
                    if val:
                        thinking_parts.append(val)
                        break
                if hasattr(delta, "tool_calls") and delta.tool_calls:
                    for tc in delta.tool_calls:
                        try:
                            tool_call_parts.append(tc.model_dump())
                        except Exception:
                            tool_call_parts.append(str(tc))
                if hasattr(delta, "function_call") and delta.function_call:
                    try:
                        function_call_parts.append(delta.function_call.model_dump())
                    except Exception:
                        function_call_parts.append(str(delta.function_call))
            if choice.finish_reason:
                finish_reason = choice.finish_reason
        if hasattr(chunk, "model") and chunk.model:
            actual_model = chunk.model
        if hasattr(chunk, "usage") and chunk.usage:
            try:
                usage = chunk.usage.model_dump()
            except AttributeError:
                usage = {
                    "prompt_tokens": getattr(chunk.usage, "prompt_tokens", None),
                    "completion_tokens": getattr(chunk.usage, "completion_tokens", None),
                    "total_tokens": getattr(chunk.usage, "total_tokens", None),
                }

    gen_time = time.monotonic() - t0
    content = "".join(content_parts)
    thinking = "".join(thinking_parts) or None

    full_parts = []
    if thinking:
        full_parts.append(f"<thinking>\n{thinking}\n</thinking>\n\n")
    if content:
        full_parts.append(content)
    if tool_call_parts:
        full_parts.append(
            f"\n\n<tool_calls>\n{json.dumps(tool_call_parts, indent=2)}\n</tool_calls>"
        )
    if function_call_parts:
        full_parts.append(
            f"\n\n<function_calls>\n{json.dumps(function_call_parts, indent=2)}\n</function_calls>"
        )
    full_response = "".join(full_parts) or None

    return (
        content,
        thinking,
        full_response,
        raw_chunks,
        usage,
        finish_reason,
        actual_model,
        gen_time,
    )


# === Generation Functions ===


# [ENGINE] Modified to use the prompt engine with randomized templates and high temperature
def generate_question(client, model):
    """
    Generate a question using the prompt engine with randomized templates.
    Returns 9 values: (content, thinking, full, chunks, usage, finish, model, gen_time, config)
    """
    try:
        config = prompt_engine.generate()
    except Exception as e:
        logger.warning(f"Prompt engine error, using fallback: {e}")
        config = PromptConfig(
            system_message=(
                "You are an expert frontend developer and technical interviewer. "
                "Generate a unique, creative UI coding problem in English. "
                "Mid-difficulty, under 1000 words."
            ),
            user_prompt=(
                "Generate 1 unique, creative, and detailed question about building "
                "a User Interface (UI) using HTML, CSS, and JavaScript. The question "
                "must be in English. The problem should be specific enough that a "
                "developer can write a complete, self-contained solution.\n\n"
                "Output ONLY the problem description. Do not include greetings, "
                "numbering, or any other formatting."
            ),
            temperature=1.0,
            max_tokens=4096,
            metadata={"fallback": True, "error": str(e)},
        )

    messages = config.to_messages()
    topic_preview = config.metadata.get("topic", "unknown")[:60]
    logger.info(
        f"📝 Streaming question generation "
        f"(temp={config.temperature}, tpl={config.metadata.get('template_id', '?')}, "
        f"topic={topic_preview}...)..."
    )
    result = stream_completion(
        client, model, messages,
        temperature=config.temperature,
        max_tokens=config.max_tokens,
    )
    # Return 9 values (original 8 + config)
    return (*result, config)


def generate_solution(client, model, question):
    prompt = f"""You are an expert frontend developer. I will give you a UI problem.
You must solve it using ONLY: {STACK}
Return ONLY the raw code. Do not include any explanations, greetings, or commentary.
If multiple files are required, use markdown code blocks and CLEARLY indicate the filename before each block (e.g., `**index.html**`).
CRITICAL INSTRUCTIONS FOR REASONING/THINKING:
1. Keep your thinking/reasoning process brief (under 2000 words).
2. DO NOT write long draft, or preview any long code inside your thinking block.
Problem: {question}
"""
    messages = [{"role": "user", "content": prompt}]
    logger.info("🔧 Streaming solution generation...")
    return stream_completion(client, model, messages, temperature=0.2, max_tokens=16384)


def health_check(client, model):
    try:
        resp = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": "Say 'ok'"}],
            max_tokens=5,
            temperature=0,
        )
        text = resp.choices[0].message.content if resp.choices else ""
        logger.info(
            f"✅ Health check passed (model={resp.model}, response='{text.strip()}')"
        )
        return True
    except Exception as e:
        logger.error(f"❌ Health check failed: {e}")
        return False


def process_single_job(
    job_id, question_id, q_text, s_client, s_model, db_path, max_attempts=3
):
    if _shutdown.is_set():
        return None
    mark_job_started(db_path, job_id)
    last_error = None
    for attempt in range(max_attempts):
        if _shutdown.is_set():
            return None
        try:
            (
                s_code,
                s_thinking,
                s_full,
                s_chunks,
                s_usage,
                s_finish,
                s_model_actual,
                gen_time,
            ) = generate_solution(s_client, s_model, q_text)
        except InterruptedError:
            logger.info(f"  Job {job_id} interrupted (shutdown).")
            return None
        except Exception as e:
            last_error = e
            if attempt < max_attempts - 1:
                wait = 5 * (2**attempt) + random.uniform(0, 3)
                logger.warning(
                    f"  Job {job_id} (Q{question_id}) attempt {attempt+1} failed: {e}. "
                    f"Retrying in {wait:.0f}s..."
                )
                time.sleep(wait)
                continue
            else:
                logger.error(
                    f"  Job {job_id} (Q{question_id}) failed after {max_attempts} attempts: {e}"
                )
                mark_job_failed(db_path, job_id, str(e))
                return None

        if not s_code and not s_full:
            last_error = "Empty response (no content or full_response)"
            if attempt < max_attempts - 1:
                logger.warning(
                    f"  Job {job_id} (Q{question_id}): empty solution, retrying..."
                )
                time.sleep(3)
                continue
            else:
                logger.warning(
                    f"  Job {job_id} (Q{question_id}): empty solution after {max_attempts} attempts."
                )
                mark_job_failed(db_path, job_id, last_error)
                return None

        save_solution(
            db_path,
            job_id,
            question_id,
            s_code or "(empty content — see full_response)",
            s_thinking,
            s_full,
            s_model_actual,
            s_finish,
            s_usage,
            s_chunks,
            gen_time,
        )
        thinking_info = f", thinking={len(s_thinking)}ch" if s_thinking else ""
        logger.info(
            f"  ✓ Job {job_id} (Q{question_id}) done "
            f"({len(s_chunks)} chunks, {gen_time:.1f}s, finish={s_finish}{thinking_info})"
        )
        return question_id

    mark_job_failed(db_path, job_id, str(last_error))
    return None


# [ENGINE] Modified to unpack 9 values and pass prompt_metadata to save_question
def generate_questions_batch(q_client, q_model, db_path, count):
    jobs = []
    for i in range(count):
        if _shutdown.is_set():
            break
        logger.info(f"📝 Generating question {i+1}/{count}...")
        try:
            (
                q_text,
                q_thinking,
                q_full,
                q_chunks,
                q_usage,
                q_finish,
                q_model_actual,
                gen_time,
                config,  # [ENGINE] 9th return value
            ) = generate_question(q_client, q_model)
        except InterruptedError:
            logger.info("Shutdown during question generation.")
            break
        except Exception as e:
            logger.error(f"Question {i+1} failed: {e}")
            continue

        if not q_text and q_full:
            q_text = q_full
        if not q_text:
            logger.warning(f"Question {i+1}: completely empty. Skipping.")
            continue

        q_text = q_text.strip()
        thinking_info = f" (thinking={len(q_thinking)}ch)" if q_thinking else ""
        logger.info(
            f"  Q{i+1} [{gen_time:.1f}s{thinking_info}]: "
            f"{q_text[:100]}{'...' if len(q_text) > 100 else ''}"
        )
        job_id, question_id = save_question(
            db_path,
            q_text,
            q_thinking,
            q_full,
            q_model_actual,
            q_finish,
            q_usage,
            q_chunks,
            gen_time,
            prompt_metadata=config.metadata,  # [ENGINE] Pass metadata
        )
        jobs.append((job_id, question_id, q_text, 0))

        if i < count - 1 and not _shutdown.is_set():
            time.sleep(random.uniform(0.3, 1.0))
    return jobs


# === Background Generation Runner ===
generation_lock = threading.Lock()
generation_running = False
generation_thread: Optional[threading.Thread] = None


def run_generation(
    number, workers, q_model, s_model, max_attempts, skip_health_check
):
    """Background generation task."""
    global generation_running
    try:
        proxy_url = f"http://127.0.0.1:{PORT}/v1"
        api_key = "proxy-key"

        http_client = httpx.Client(
            timeout=httpx.Timeout(connect=30.0, read=600.0, write=60.0, pool=30.0)
        )
        q_client = OpenAI(api_key=api_key, base_url=proxy_url, http_client=http_client)
        s_client = OpenAI(api_key=api_key, base_url=proxy_url, http_client=http_client)

        logger.info(f"Q client → {proxy_url} (model: {q_model})")
        logger.info(f"S client → {proxy_url} (model: {s_model})")

        if not skip_health_check:
            logger.info("Running health check...")
            if not health_check(s_client, s_model):
                logger.error("Health check failed. Use skip_health_check=true to skip.")
                http_client.close()
                return

        init_db(DB_PATH)

        pending = get_pending_jobs(DB_PATH)
        if pending:
            logger.info(
                f"🔄 Found {len(pending)} pending jobs from previous crash — resuming those first!"
            )

        new_jobs = generate_questions_batch(q_client, q_model, DB_PATH, number)
        logger.info(f"Generated {len(new_jobs)} new questions.")

        all_jobs = pending + new_jobs
        if not all_jobs:
            logger.info("No jobs to process.")
            http_client.close()
            return

        logger.info(
            f"📦 Processing {len(all_jobs)} jobs with {workers} concurrent workers..."
        )

        completed = 0
        failed = 0
        t_start = time.monotonic()

        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = {}
            for job_id, q_id, q_text, _attempts in all_jobs:
                if _shutdown.is_set():
                    break
                f = pool.submit(
                    process_single_job,
                    job_id,
                    q_id,
                    q_text,
                    s_client,
                    s_model,
                    DB_PATH,
                    max_attempts,
                )
                futures[f] = (job_id, q_id)

            for future in as_completed(futures):
                job_id, q_id = futures[future]
                try:
                    result = future.result()
                    if result is not None:
                        completed += 1
                    else:
                        failed += 1
                except Exception as e:
                    logger.error(f"Job {job_id} (Q{q_id}) crashed: {e}")
                    mark_job_failed(DB_PATH, job_id, str(e))
                    failed += 1

                if _shutdown.is_set():
                    logger.warning(
                        "Shutdown flag set — remaining jobs stay pending for next run."
                    )
                    break

        elapsed = time.monotonic() - t_start
        remaining = get_pending_jobs(DB_PATH)
        log_run(
            DB_PATH,
            q_model,
            s_model,
            len(all_jobs),
            completed,
            failed,
            len(remaining),
            {
                "workers": workers,
                "max_attempts": max_attempts,
                "proxy": proxy_url,
                "elapsed_s": round(elapsed, 1),
            },
        )

        logger.info(
            f"✅ Completed: {completed} | ❌ Failed: {failed} | ⏳ Pending: {len(remaining)}"
        )
        logger.info(f"⏱️ Total time: {elapsed:.1f}s")
        http_client.close()
    except Exception as e:
        logger.error(f"Generation run failed: {e}", exc_info=True)


def run_resume(workers, s_model, max_attempts):
    """Resume pending jobs."""
    global generation_running
    try:
        proxy_url = f"http://127.0.0.1:{PORT}/v1"
        http_client = httpx.Client(
            timeout=httpx.Timeout(connect=30.0, read=600.0, write=60.0, pool=30.0)
        )
        s_client = OpenAI(api_key="proxy-key", base_url=proxy_url, http_client=http_client)

        pending = get_pending_jobs(DB_PATH)
        if not pending:
            logger.info("No pending jobs to resume.")
            http_client.close()
            return

        logger.info(f"📦 Resuming {len(pending)} pending jobs with {workers} workers...")
        completed = 0
        failed = 0
        t_start = time.monotonic()

        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = {}
            for job_id, q_id, q_text, _attempts in pending:
                if _shutdown.is_set():
                    break
                f = pool.submit(
                    process_single_job,
                    job_id,
                    q_id,
                    q_text,
                    s_client,
                    s_model,
                    DB_PATH,
                    max_attempts,
                )
                futures[f] = (job_id, q_id)

            for future in as_completed(futures):
                job_id, q_id = futures[future]
                try:
                    result = future.result()
                    if result is not None:
                        completed += 1
                    else:
                        failed += 1
                except Exception as e:
                    logger.error(f"Job {job_id} (Q{q_id}) crashed: {e}")
                    mark_job_failed(DB_PATH, job_id, str(e))
                    failed += 1

        elapsed = time.monotonic() - t_start
        remaining = get_pending_jobs(DB_PATH)
        log_run(
            DB_PATH,
            "resume",
            s_model,
            len(pending),
            completed,
            failed,
            len(remaining),
            {"workers": workers, "max_attempts": max_attempts, "elapsed_s": round(elapsed, 1)},
        )
        logger.info(f"Resume done: ✅{completed} ❌{failed} ⏳{len(remaining)}")
        http_client.close()
    except Exception as e:
        logger.error(f"Resume failed: {e}", exc_info=True)


# === Pydantic Models ===


class GenerateRequest(BaseModel):
    number: int = 20
    workers: int = 20
    q_model: str = "moonshotai/Kimi-K2.6"
    s_model: str = "zai-org/GLM-5.2"
    max_attempts: int = 3
    skip_health_check: bool = False


class ResumeRequest(BaseModel):
    workers: int = 20
    s_model: str = "zai-org/GLM-5.2"
    max_attempts: int = 3


class AddKeyRequest(BaseModel):
    key: str


# [ENGINE] New pydantic models for prompt management
class AddTemplateRequest(BaseModel):
    template_id: str
    template_text: str


class AddTopicRequest(BaseModel):
    category: str
    topic: str


class AddSystemMessageRequest(BaseModel):
    message: str


class TemperatureRangeRequest(BaseModel):
    min_temp: float = 0.85
    max_temp: float = 1.15


class MaxTokensRequest(BaseModel):
    max_tokens: int = 4096


# === FastAPI App ===

timeout_config = httpx.Timeout(connect=10.0, read=300.0, write=300.0, pool=10.0)
proxy_http_client: Optional[httpx.AsyncClient] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global proxy_http_client
    init_db(DB_PATH)
    proxy_http_client = httpx.AsyncClient(timeout=timeout_config)
    logger.info(f"Server starting on port {PORT}")
    logger.info(f"Data directory: {DATA_DIR}")
    logger.info(f"Database: {DB_PATH}")
    logger.info(f"Keys file: {KEYS_FILE}")
    logger.info(f"Loaded {len(_keys)} API key(s)")
    logger.info(f"Featherless base: {FEATHERLESS_API_BASE}")
    # [ENGINE] Log engine info
    engine_stats = prompt_engine.get_stats()
    logger.info(
        f"Prompt engine: {engine_stats['available_system_messages']} system msgs, "
        f"{engine_stats['available_templates']} templates, "
        f"{engine_stats['available_topics']} topics, "
        f"temp range {engine_stats['temperature_range']}"
    )
    yield
    _shutdown.set()
    if proxy_http_client:
        await proxy_http_client.aclose()
    logger.info("Server shutting down...")


app = FastAPI(title="IC Generate API", lifespan=lifespan)


# === Homepage ===


@app.get("/", response_class=HTMLResponse)
async def homepage():
    return """<!DOCTYPE html>
<html>
<head><title>IC Generate API</title>
<style>
body { font-family: monospace; max-width: 900px; margin: 50px auto; padding: 20px; background: #1a1a2e; color: #e0e0e0; }
h1 { color: #00d4ff; }
h2 { color: #00ffaa; margin-top: 30px; }
h3 { color: #ffaa00; margin-top: 20px; }
code { background: #16213e; padding: 2px 6px; border-radius: 3px; color: #ff6b6b; }
ul { line-height: 1.8; }
a { color: #00d4ff; }
</style>
</head>
<body>
<h1>🚀 IC Generate API</h1>
<p>Combined Featherless proxy + UI question/solution generator with full thinking capture.</p>
<p><b>✨ Now with randomized prompt engine (high temperature, 200+ topics, 8 templates)!</b></p>
<h2>Admin Authentication:</h2>
<p>All management routes require the <code>X-Admin-Key</code> header.<br>
Set the <code>ADMIN_API_KEY</code> environment variable in Hugging Face Spaces Secrets.</p>
<h2>Endpoints (Public):</h2>
<ul>
<li><b>POST /v1/{path}</b> — Proxy to Featherless API (key rotation)</li>
</ul>
<h2>Endpoints (Admin-Protected):</h2>
<h3>Generation</h3>
<ul>
<li><b>POST /generate</b> — Start generation run (background)</li>
<li><b>GET /generate/status</b> — Check if generation is running</li>
<li><b>POST /resume</b> — Resume pending jobs</li>
</ul>
<h3>Data & Stats</h3>
<ul>
<li><b>GET /stats</b> — Database statistics (includes engine stats)</li>
<li><b>GET /questions</b> — List questions (paginated)</li>
<li><b>GET /question/{id}</b> — Get question with solutions</li>
<li><b>GET /export/json</b> — Export all data as JSON</li>
<li><b>GET /export/files</b> — Export to /data/exported_code/</li>
<li><b>GET /pending</b> — List pending jobs</li>
<li><b>DELETE /question/{id}</b> — Delete a question</li>
<li><b>GET /health</b> — Health check</li>
</ul>
<h3>API Keys</h3>
<ul>
<li><b>POST /keys/add</b> — Add a new Featherless API key</li>
<li><b>GET /keys/list</b> — List active Featherless API keys</li>
<li><b>DELETE /keys/{key}</b> — Delete a Featherless API key</li>
<li><b>POST /keys/reload</b> — Reload keys from file</li>
</ul>
<h3>✨ Prompt Engine Management</h3>
<ul>
<li><b>GET /prompts/stats</b> — Engine statistics (generations, unique combos, template usage)</li>
<li><b>GET /prompts/preview</b> — Preview a randomly generated prompt (no API call)</li>
<li><b>GET /prompts/templates</b> — List all prompt templates</li>
<li><b>POST /prompts/templates</b> — Add a custom template</li>
<li><b>DELETE /prompts/templates/{id}</b> — Remove a custom template</li>
<li><b>GET /prompts/topics</b> — List all topics</li>
<li><b>POST /prompts/topics</b> — Add a custom topic</li>
<li><b>DELETE /prompts/topics/{topic}</b> — Remove a custom topic</li>
<li><b>GET /prompts/system-messages</b> — List all system messages</li>
<li><b>POST /prompts/system-messages</b> — Add a custom system message</li>
<li><b>DELETE /prompts/system-messages/{index}</b> — Remove a custom system message</li>
<li><b>POST /prompts/temperature</b> — Set temperature range</li>
<li><b>POST /prompts/max-tokens</b> — Set max tokens</li>
<li><b>POST /prompts/reset</b> — Reset engine deduplication state</li>
</ul>
</body>
</html>"""


# === Proxy Endpoint (Public) ===


@app.api_route(
    "/v1/{path:path}",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
)
async def proxy(request: Request, path: str):
    """Proxy requests to Featherless API with key rotation."""
    url = f"{FEATHERLESS_API_BASE}/{path}"
    api_key = get_next_key()

    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("content-length", None)
    headers.pop("content-encoding", None)
    headers.pop("transfer-encoding", None)
    headers["authorization"] = f"Bearer {api_key}"

    body = await request.body()

    req = proxy_http_client.build_request(
        method=request.method,
        url=url,
        headers=headers,
        content=body if body else None,
        params=request.query_params,
    )

    try:
        response = await proxy_http_client.send(req, stream=True)
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"Proxy error: {str(e)}")

    async def generate():
        try:
            async for chunk in response.aiter_bytes():
                yield chunk
        except httpx.ReadTimeout:
            pass
        finally:
            await response.aclose()

    response_headers = dict(response.headers)
    response_headers.pop("content-encoding", None)
    response_headers.pop("transfer-encoding", None)
    response_headers.pop("content-length", None)

    return StreamingResponse(
        generate(),
        status_code=response.status_code,
        headers=response_headers,
        media_type=response.headers.get("content-type"),
    )


# === Generate Endpoint (Admin) ===


@app.post("/generate", dependencies=[Depends(verify_admin_key)])
async def generate(req: GenerateRequest):
    """Start a generation run in the background."""
    global generation_running, generation_thread
    with generation_lock:
        if generation_running:
            raise HTTPException(
                status_code=409, detail="A generation run is already in progress."
            )
        generation_running = True

    def task_wrapper():
        global generation_running
        try:
            run_generation(
                req.number,
                req.workers,
                req.q_model,
                req.s_model,
                req.max_attempts,
                req.skip_health_check,
            )
        finally:
            with generation_lock:
                generation_running = False

    generation_thread = threading.Thread(target=task_wrapper, daemon=True)
    generation_thread.start()

    return {
        "status": "started",
        "message": f"Generating {req.number} Q&A pairs with {req.workers} workers.",
        "config": {
            "number": req.number,
            "workers": req.workers,
            "q_model": req.q_model,
            "s_model": req.s_model,
            "max_attempts": req.max_attempts,
            "skip_health_check": req.skip_health_check,
        },
    }


@app.get("/generate/status", dependencies=[Depends(verify_admin_key)])
async def generation_status():
    """Check if a generation run is in progress."""
    return {"running": generation_running}


# === Resume Endpoint (Admin) ===


@app.post("/resume", dependencies=[Depends(verify_admin_key)])
async def resume(req: ResumeRequest):
    """Resume pending jobs without generating new questions."""
    global generation_running, generation_thread
    with generation_lock:
        if generation_running:
            raise HTTPException(
                status_code=409, detail="A generation run is already in progress."
            )
        generation_running = True

    pending = get_pending_jobs(DB_PATH)
    if not pending:
        with generation_lock:
            generation_running = False
        return {"status": "no_pending", "message": "No pending jobs to resume."}

    def task_wrapper():
        global generation_running
        try:
            run_resume(req.workers, req.s_model, req.max_attempts)
        finally:
            with generation_lock:
                generation_running = False

    generation_thread = threading.Thread(target=task_wrapper, daemon=True)
    generation_thread.start()

    return {
        "status": "started",
        "message": f"Resuming {len(pending)} pending jobs with {req.workers} workers.",
    }


@app.post("/fix-missing-solutions", dependencies=[Depends(verify_admin_key)])
async def fix_missing_solutions():
    """Mark questions without solutions as pending so they can be generated."""
    fixed_count = 0
    with sqlite3.connect(DB_PATH) as conn:
        c = conn.cursor()
        
        # Find questions that don't have solutions
        c.execute('''
            SELECT q.id 
            FROM questions q 
            LEFT JOIN solutions s ON q.id = s.question_id 
            WHERE s.id IS NULL
        ''')
        missing_qs = c.fetchall()
        
        for q in missing_qs:
            q_id = q[0]
            # Check if it exists in pending_jobs
            c.execute("SELECT id, status FROM pending_jobs WHERE question_id=?", (q_id,))
            job = c.fetchone()
            if job:
                if job[1] != 'pending':
                    c.execute("UPDATE pending_jobs SET status='pending', attempts=0, last_error=NULL, completed_at=NULL WHERE id=?", (job[0],))
                    fixed_count += 1
            else:
                c.execute("INSERT INTO pending_jobs (question_id, status) VALUES (?, 'pending')", (q_id,))
                fixed_count += 1
                
        conn.commit()
    
    return {
        "status": "success", 
        "message": f"Fixed {fixed_count} pending jobs for questions without solutions."
    }


# === Stats Endpoint (Admin) ===


@app.get("/stats", dependencies=[Depends(verify_admin_key)])
async def stats():
    return get_stats_dict(DB_PATH)


# === Questions Endpoints (Admin) ===


@app.get("/questions", dependencies=[Depends(verify_admin_key)])
async def list_questions(limit: int = Query(20, ge=1, le=100), offset: int = Query(0, ge=0)):
    """List questions with pagination."""
    with sqlite3.connect(DB_PATH) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM questions")
        total = cursor.fetchone()[0]
        cursor.execute(
            """
            SELECT q.id, q.question_text, q.model, q.created_at, q.generation_time_s,
                   (SELECT COUNT(*) FROM solutions WHERE question_id = q.id) AS solution_count
            FROM questions q
            ORDER BY q.id DESC
            LIMIT ? OFFSET ?
        """,
            (limit, offset),
        )
        rows = cursor.fetchall()

    return {
        "total": total,
        "limit": limit,
        "offset": offset,
        "questions": [
            {
                "id": r[0],
                "question": r[1],
                "model": r[2],
                "created_at": r[3],
                "generation_time_s": r[4],
                "solution_count": r[5],
            }
            for r in rows
        ],
    }


@app.get("/question/{question_id}", dependencies=[Depends(verify_admin_key)])
async def get_question(question_id: int):
    """Get a specific question with its solutions."""
    with sqlite3.connect(DB_PATH) as conn:
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM questions WHERE id=?", (question_id,))
        q = cursor.fetchone()
        if not q:
            raise HTTPException(
                status_code=404, detail=f"Question {question_id} not found."
            )
        cursor.execute("SELECT * FROM solutions WHERE question_id=?", (question_id,))
        solutions = cursor.fetchall()

    question = dict(q)
    question.pop("raw_chunks_json", None)
    question["usage"] = json.loads(question.pop("usage_json", "null") or "null")
    # [ENGINE] Parse prompt_metadata
    question["prompt_metadata"] = json.loads(
        question.pop("prompt_metadata", "null") or "null"
    )
    question["solutions"] = []
    for s in solutions:
        sol = dict(s)
        sol.pop("raw_chunks_json", None)
        sol["usage"] = json.loads(sol.pop("usage_json", "null") or "null")
        question["solutions"].append(sol)

    return question


# === Export Endpoints (Admin) ===


@app.get("/export/json", dependencies=[Depends(verify_admin_key)])
async def export_json():
    data = export_data_json(DB_PATH)
    return JSONResponse(content=data)


@app.get("/export/files", dependencies=[Depends(verify_admin_key)])
async def export_files_endpoint():
    result = export_files(DB_PATH, EXPORT_DIR)
    return result


# === Pending Jobs Endpoint (Admin) ===


@app.get("/pending", dependencies=[Depends(verify_admin_key)])
async def pending_jobs():
    jobs = get_pending_jobs(DB_PATH)
    return {
        "count": len(jobs),
        "jobs": [
            {
                "job_id": j[0],
                "question_id": j[1],
                "attempts": j[3],
                "question_preview": j[2][:100] if j[2] else "",
            }
            for j in jobs
        ],
    }


# === Delete Question Endpoint (Admin) ===


@app.delete("/question/{question_id}", dependencies=[Depends(verify_admin_key)])
async def delete_question(question_id: int):
    success = delete_question_db(DB_PATH, question_id)
    if success:
        return {"status": "deleted", "question_id": question_id}
    else:
        raise HTTPException(
            status_code=404, detail=f"Question {question_id} not found."
        )


# === Health Endpoint (Admin) ===


@app.get("/health", dependencies=[Depends(verify_admin_key)])
async def health():
    return {
        "status": "ok",
        "data_dir": DATA_DIR,
        "db_path": DB_PATH,
        "db_exists": os.path.exists(DB_PATH),
        "keys_loaded": len(_keys),
        "keys_file": KEYS_FILE,
        "keys_file_exists": os.path.exists(KEYS_FILE),
        "featherless_base": FEATHERLESS_API_BASE,
        "port": PORT,
        "generation_running": generation_running,
    }


# === Keys Management Endpoints (Admin) ===


@app.get("/keys", dependencies=[Depends(verify_admin_key)])
async def keys_info():
    return {
        "count": len(_keys),
        "keys_file": KEYS_FILE,
        "keys_file_exists": os.path.exists(KEYS_FILE),
    }


@app.get("/keys/list", dependencies=[Depends(verify_admin_key)])
async def list_keys():
    """List all active Featherless API keys."""
    return {"count": len(_keys), "keys": _keys}


@app.post("/keys/add", dependencies=[Depends(verify_admin_key)])
async def add_key(req: AddKeyRequest):
    """Add a new Featherless API key to the rotation."""
    global _keys, _key_cycle
    key = req.key.strip()
    if not key:
        raise HTTPException(status_code=400, detail="Key cannot be empty")

    with open(KEYS_FILE, "a") as f:
        f.write(f"{key}\n")

    if key not in _keys:
        _keys.append(key)
        _key_cycle = itertools.cycle(_keys)

    return {"status": "success", "message": "Key added successfully", "total_keys": len(_keys)}


@app.delete("/keys/{key}", dependencies=[Depends(verify_admin_key)])
async def delete_key(key: str):
    """Remove a Featherless API key from the rotation."""
    global _keys, _key_cycle
    if key in _keys:
        _keys.remove(key)
        _key_cycle = itertools.cycle(_keys)

        with open(KEYS_FILE, "w") as f:
            for k in _keys:
                f.write(f"{k}\n")

        return {"status": "success", "message": "Key deleted", "total_keys": len(_keys)}
    raise HTTPException(status_code=404, detail="Key not found")


@app.post("/keys/reload", dependencies=[Depends(verify_admin_key)])
async def reload_keys_endpoint():
    count = reload_keys()
    return {"status": "reloaded", "count": count}


# ============================================================================
# [ENGINE] Prompt Engine Management Endpoints (Admin)
# ============================================================================


@app.get("/prompts/stats", dependencies=[Depends(verify_admin_key)])
async def prompt_engine_stats():
    """Get prompt engine statistics."""
    return prompt_engine.get_stats()


@app.get("/prompts/preview", dependencies=[Depends(verify_admin_key)])
async def prompt_preview():
    """Preview a randomly generated prompt without consuming dedup state or calling the API."""
    config = prompt_engine.peek()
    return config.to_dict()


@app.get("/prompts/templates", dependencies=[Depends(verify_admin_key)])
async def list_templates():
    """List all available prompt templates (built-in + custom)."""
    templates = prompt_engine.get_all_templates()
    return {"count": len(templates), "templates": templates}


@app.post("/prompts/templates", dependencies=[Depends(verify_admin_key)])
async def add_template(req: AddTemplateRequest):
    """Add a custom prompt template."""
    success = prompt_engine.add_template(req.template_id, req.template_text)
    if success:
        return {"status": "success", "message": f"Template '{req.template_id}' added."}
    raise HTTPException(
        status_code=400,
        detail=f"Failed to add template. ID may be duplicate or fields empty."
    )


@app.delete("/prompts/templates/{template_id}", dependencies=[Depends(verify_admin_key)])
async def remove_template(template_id: str):
    """Remove a custom prompt template by ID."""
    success = prompt_engine.remove_template(template_id)
    if success:
        return {"status": "success", "message": f"Template '{template_id}' removed."}
    raise HTTPException(status_code=404, detail=f"Template '{template_id}' not found in custom templates.")


@app.get("/prompts/topics", dependencies=[Depends(verify_admin_key)])
async def list_topics():
    """List all available topics (built-in + custom)."""
    topics = prompt_engine.get_all_topics()
    return {"count": len(topics), "topics": topics}


@app.post("/prompts/topics", dependencies=[Depends(verify_admin_key)])
async def add_topic(req: AddTopicRequest):
    """Add a custom topic."""
    success = prompt_engine.add_topic(req.category, req.topic)
    if success:
        return {"status": "success", "message": f"Topic added to category '{req.category}'."}
    raise HTTPException(
        status_code=400,
        detail="Failed to add topic. It may be a duplicate or fields are empty."
    )


@app.delete("/prompts/topics/{topic}", dependencies=[Depends(verify_admin_key)])
async def remove_topic(topic: str):
    """Remove a custom topic by topic text."""
    success = prompt_engine.remove_topic(topic)
    if success:
        return {"status": "success", "message": f"Topic '{topic[:50]}' removed."}
    raise HTTPException(status_code=404, detail=f"Topic not found in custom topics.")


@app.get("/prompts/system-messages", dependencies=[Depends(verify_admin_key)])
async def list_system_messages():
    """List all system messages (built-in + custom)."""
    messages = prompt_engine.get_all_system_messages()
    return {"count": len(messages), "system_messages": messages}


@app.post("/prompts/system-messages", dependencies=[Depends(verify_admin_key)])
async def add_system_message(req: AddSystemMessageRequest):
    """Add a custom system message."""
    success = prompt_engine.add_system_message(req.message)
    if success:
        return {"status": "success", "message": "System message added."}
    raise HTTPException(
        status_code=400,
        detail="Failed to add system message. It may be a duplicate or empty."
    )


@app.delete("/prompts/system-messages/{index}", dependencies=[Depends(verify_admin_key)])
async def remove_system_message(index: int):
    """Remove a custom system message by index (custom index only)."""
    success = prompt_engine.remove_system_message(index)
    if success:
        return {"status": "success", "message": f"Custom system message at index {index} removed."}
    raise HTTPException(status_code=404, detail=f"No custom system message at index {index}.")


@app.post("/prompts/temperature", dependencies=[Depends(verify_admin_key)])
async def set_temperature_range(req: TemperatureRangeRequest):
    """Set the temperature range for question generation."""
    prompt_engine.set_temperature_range(req.min_temp, req.max_temp)
    stats = prompt_engine.get_stats()
    return {
        "status": "success",
        "message": f"Temperature range set to [{req.min_temp}, {req.max_temp}].",
        "temperature_range": stats["temperature_range"],
    }


@app.post("/prompts/max-tokens", dependencies=[Depends(verify_admin_key)])
async def set_max_tokens(req: MaxTokensRequest):
    """Set max tokens for question generation."""
    prompt_engine.set_max_tokens(req.max_tokens)
    return {
        "status": "success",
        "message": f"Max tokens set to {req.max_tokens}.",
        "max_tokens": req.max_tokens,
    }


@app.post("/prompts/reset", dependencies=[Depends(verify_admin_key)])
async def reset_prompt_engine():
    """Reset the prompt engine's deduplication state."""
    prompt_engine.reset_state()
    return {"status": "success", "message": "Prompt engine state reset."}


# === Main ===

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=PORT)