app/api/training.py

"""``/api/training/*`` — list, generate, evaluate, solve endpoints.

Per-type generation / mapping logic lives in
:mod:`app.services.scenario_service` and
:mod:`app.services.challenge_loader`. Per-type evaluation lives in
:mod:`app.services.evaluator`. This router is a thin HTTP shell.
"""

from typing import Optional
import random

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from pydantic import BaseModel

from app.core.auth import get_current_user
from app.core.constants import CYBER_SECURITY_TOPICS
from app.core.module_router import challenge_type_for_module
from app.services.challenge_loader import (
    map_encryption_row_to_training,
    map_code_fixing_row_to_training,
    map_log_analysis_row_to_training,
    map_vuln_hunter_row_to_training,
    map_steganography_row_to_training,
    map_web_exploit_row_to_training,
)
from app.services.evaluator import (
    evaluate_training,
    evaluate_code_fix,
    evaluate_log_analysis,
    evaluate_vuln_hunter,
    evaluate_web_exploit,
)
from app.services.scenario_service import (
    attach_scenario_metadata,
    handle_background_replacement,
    map_scenario_to_list_item,
)
from app.services.supabase_service import (
    fetch_scenario_by_id,
    fetch_random_scenario_from_supabase,
    fetch_all_scenarios_for_type,
    get_supabase_scenario_count,
    scenario_table,
)
from app.types import (
    EvaluateRequest,
    TrainingRequest,
    SolvedRequest,
    CodeFixEvaluateRequest,
    LogAnalysisEvaluateRequest,
    VulnHunterEvaluateRequest,
)


router = APIRouter()


# --------------------------------------------------------------------------- #
#  /api/training/list — dashboard list                                       #
# --------------------------------------------------------------------------- #

@router.get("/api/training/list")
async def list_challenges(team_role: str = "blue", difficulty: Optional[str] = None, limit: int = 100):
    """List cached challenges for a team.

    For each challenge type the team has access to, fetch all rows
    (capped at ``limit``) and map them to the lightweight list shape
    the dashboard renders. Blue team gets code-fixing / log-analysis
    / vulnerability-hunter; red team gets crypto.

    Note: the previous implementation called
    :func:`fetch_random_scenario_from_supabase` with an empty
    ``module`` argument, which built a ``module=eq.`` PostgREST
    filter that never matched (modules in the DB are never empty).
    That made the dashboard report zero challenges even when the pool
    was full. This version uses
    :func:`fetch_all_scenarios_for_type` instead.
    """
    candidates: list[str] = []
    if team_role == "blue":
        candidates = ["code-fixing", "log-analysis", "vulnerability-hunter"]
    else:
        candidates = ["crypto", "steganography", "web-exploitation"]

    items: list[dict] = []
    per_type_cap = max(1, limit // max(1, len(candidates)))
    for ctype in candidates:
        try:
            rows = await fetch_all_scenarios_for_type(
                team_role, challenge_type=ctype, limit=per_type_cap
            )
            if not rows:
                continue
            for row in rows:
                row["_challenge_type"] = ctype
                items.append(map_scenario_to_list_item(row))
                if len(items) >= limit:
                    break
        except Exception as e:
            print(f"[list_challenges] {ctype}: {e}")
            continue
        if len(items) >= limit:
            break

    return {"items": items, "total": len(items)}


# --------------------------------------------------------------------------- #
#  /api/training/solved — consume one from the pool                            #
# --------------------------------------------------------------------------- #

@router.post("/api/training/solved")
async def solve_challenge(req: SolvedRequest, background_tasks: BackgroundTasks):
    ctype = challenge_type_for_module(req.module or "")
    background_tasks.add_task(
        handle_background_replacement,
        req.scenarioId,
        req.teamRole or "red",
        req.module or "",
        req.path or "cryptography",
        req.category or "encryption",
        req.difficulty or "متوسط",
    )
    return {"status": "consumed", "challenge_type": ctype}


# --------------------------------------------------------------------------- #
#  /api/training/generate — hydrate a cached scenario into a full challenge   #
# --------------------------------------------------------------------------- #

@router.post("/api/training/generate")
async def generate_training(req: TrainingRequest, background_tasks: BackgroundTasks):
    module = req.module
    path = req.path
    category = req.category
    challenge_id = req.challengeId
    team_role = req.teamRole or "red"

    if not module or not path or not category:
        raise HTTPException(status_code=400, detail="module / path / category are required")

    # The ``module`` field on the request is the **challenge type** (one of
    # the five canonical values) — see AGENTS.md "Challenge Type vs Module".
    # We use it to (a) look the row up in the right table and (b) pin
    # training.type to the canonical value.
    challenge_type = _normalize_challenge_type(module)

    # 1) Try to grab a cached scenario for this (type, team). When
    # ``challenge_id`` is given we look in every per-type table until
    # we find the row; otherwise we pick a random row from the table
    # that matches the requested type.
    scenario: Optional[dict] = None
    if challenge_id:
        scenario = await _fetch_scenario_by_id_typed(
            team_role, challenge_id, challenge_type
        )
    if not scenario:
        scenario = await _fetch_random_scenario_typed(
            team_role, challenge_type
        )

    if not scenario:
        raise HTTPException(
            status_code=404,
            detail=(
                f"لا يوجد سيناريو في البركة لهذه الوحدة ({module}). "
                "سيناريو جديد سيُولَّد قريباً."
            ),
        )

    # 2) Project the raw DB row straight into the TrainingData shape.
    #    The row mappers in app.services.challenge_loader are the
    #    single source of truth for "DB row → front-end payload" — they
    #    include vulnerable_code / log_url / htmlPreview / etc. for
    #    each type, so no further AI round-trip is needed.
    training_data = _project_row_to_training(scenario, team_role, challenge_type)

    # 3) For legacy callers / fronts that still expect a ``type``
    #    alias and a stable id, run the metadata attach. It is now a
    #    pure pass-through because the mapper already set every field.
    training_data = attach_scenario_metadata(
        training_data, scenario, challenge_type=challenge_type
    )

    return {"training": training_data}


# --------------------------------------------------------------------------- #
#  Row → TrainingData projection                                              #
# --------------------------------------------------------------------------- #

_TYPE_TO_MAPPER = {
    "crypto":               "map_encryption_row_to_training",
    "code-fixing":          "map_code_fixing_row_to_training",
    "log-analysis":         "map_log_analysis_row_to_training",
    "vulnerability-hunter": "map_vuln_hunter_row_to_training",
    "steganography":        "map_steganography_row_to_training",
    "web-exploitation":     "map_web_exploit_row_to_training",
}


def _project_row_to_training(
    scenario: dict, team_role: str, challenge_type: str
) -> dict:
    """Run the per-type mapper against a raw DB row.

    Falls back to a minimal "best-effort" shape when the challenge
    type is unknown so that the front-end at least gets a
    ``type``/``id`` pair to render a generic error. This should
    never fire in practice because the candidate list in
    ``list_challenges`` only returns the four canonical types.
    """
    mapper_name = _TYPE_TO_MAPPER.get(challenge_type)
    if mapper_name is None:
        return {
            "id": str(scenario.get("id", "")),
            "scenarioId": str(scenario.get("id", "")),
            "title": scenario.get("title", ""),
            "story": scenario.get("story", ""),
            "task": scenario.get("task_outline", ""),
            "type": challenge_type,
            "difficulty": scenario.get("difficulty", "متوسط"),
            "xpReward": scenario.get("xp_reward", 150),
            "hints": scenario.get("hints") or [],
        }
    mapper = globals()[mapper_name]
    return mapper(scenario, team_role)


# --------------------------------------------------------------------------- #
#  Helpers used only by /api/training/generate                                #
# --------------------------------------------------------------------------- #

# Canonical challenge types — these are the ONLY values that may appear
# in ``training["type"]`` and they map 1:1 to a per-type Supabase table.
_CANONICAL_CHALLENGE_TYPES: dict[str, str] = {
    "crypto": "crypto",
    "encryption": "crypto",
    "code-fixing": "code-fixing",
    "code_fixing": "code-fixing",
    "log-analysis": "log-analysis",
    "log_analysis": "log-analysis",
    "vulnerability-hunter": "vulnerability-hunter",
    "vulnerability_hunter": "vulnerability-hunter",
    "steganography": "steganography",
    "web-exploitation": "web-exploitation",
    "web_exploitation": "web-exploitation",
    "web": "web-exploitation",
}


def _normalize_challenge_type(module: str) -> str:
    """Resolve any module alias to its canonical challenge type.

    Unknown values fall through to ``"crypto"`` (the legacy default) so
    that the red-team code path keeps working. The trainer endpoint
    never raises on bad input — it logs and best-efforts.
    """
    key = (module or "").strip().lower()
    return _CANONICAL_CHALLENGE_TYPES.get(key, "crypto")


def _tables_for_team(team_role: str) -> list[str]:
    """All per-type tables that can hold a row for ``team_role``.

    Used to look up a row by id without knowing its type up front. The
    list is short and stable so a sequential scan is fine.
    """
    if team_role == "red":
        return [
            "encryption_challenges",
            "steganography_challenges",
            "web_exploitation_challenges",
        ]
    return [
        "code_fixing_challenges",
        "log_analysis_challenges",
        "vulnerability_hunter_challenges",
    ]


async def _fetch_scenario_by_id_typed(
    team_role: str, scenario_id: str, hint_type: str
) -> Optional[dict]:
    """Look up a scenario by id, scoped to the team.

    Tries the hint table first (single round-trip), then falls back to
    every other table the team owns. This replaces the old behaviour
    where ``fetch_scenario_by_id`` always looked in
    ``encryption_challenges`` — which would silently return ``None``
    for code-fixing / log-analysis / vuln-hunter rows.
    """
    if not scenario_id:
        return None
    tables = [scenario_table(team_role, hint_type)] + [
        t for t in _tables_for_team(team_role)
        if t != scenario_table(team_role, hint_type)
    ]
    for table in tables:
        row = await _fetch_scenario_in_table(team_role, scenario_id, table)
        if row:
            return row
    return None


async def _fetch_scenario_in_table(
    team_role: str, scenario_id: str, table: str
) -> Optional[dict]:
    import httpx
    from app.services.supabase_service import supabase_headers
    from app.core.config import SUPABASE_URL, SUPABASE_ANON_KEY
    if not SUPABASE_ANON_KEY or not SUPABASE_URL:
        return None
    url = (
        f"{SUPABASE_URL}/rest/v1/{table}"
        f"?id=eq.{scenario_id}&team_role=eq.{team_role}&limit=1"
    )
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(url, headers=supabase_headers())
            if resp.status_code == 200:
                rows = resp.json()
                if rows:
                    return rows[0]
    except Exception as e:
        print(f"[generate_training] {table} lookup failed: {e}")
    return None


async def _fetch_random_scenario_typed(
    team_role: str, challenge_type: str
) -> Optional[dict]:
    """Pick a random row from the table that matches ``challenge_type``."""
    from app.services.supabase_service import fetch_all_scenarios_for_type
    rows = await fetch_all_scenarios_for_type(team_role, challenge_type, limit=50)
    if rows:
        return random.choice(rows)
    return None


# --------------------------------------------------------------------------- #
#  /api/training/evaluate — generic / legacy red+blue                         #
# --------------------------------------------------------------------------- #

@router.post("/api/training/evaluate")
async def post_evaluate_training(req: EvaluateRequest, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
    # IDOR FIX: Override userId from JWT
    req.userId = user["user_id"]
    return await evaluate_training(req, background_tasks)


# --------------------------------------------------------------------------- #
#  /api/training/evaluate-code-fix                                            #
# --------------------------------------------------------------------------- #

@router.post("/api/training/evaluate-code-fix")
async def post_evaluate_code_fix(req: CodeFixEvaluateRequest, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
    # IDOR FIX: Override userId from JWT
    req.userId = user["user_id"]
    return await evaluate_code_fix(req, background_tasks)


# --------------------------------------------------------------------------- #
#  /api/training/evaluate-log-analysis                                        #
# --------------------------------------------------------------------------- #

@router.post("/api/training/evaluate-log-analysis")
async def post_evaluate_log_analysis(req: LogAnalysisEvaluateRequest, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
    # IDOR FIX: Override userId from JWT
    req.userId = user["user_id"]
    return await evaluate_log_analysis(req, background_tasks)


# --------------------------------------------------------------------------- #
#  /api/training/evaluate-vulnerability-hunter                                #
# --------------------------------------------------------------------------- #

@router.post("/api/training/evaluate-vulnerability-hunter")
async def post_evaluate_vuln_hunter(req: VulnHunterEvaluateRequest, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
    # IDOR FIX: Override userId from JWT
    req.userId = user["user_id"]
    return await evaluate_vuln_hunter(req, background_tasks)


# --------------------------------------------------------------------------- #
#  /api/training/evaluate-web-exploit                                         #
# --------------------------------------------------------------------------- #

class WebExploitEvaluateRequest(BaseModel):
    userId: str = ""
    teamRole: str = "red"
    challengeId: str = ""
    payload: str = ""


@router.post("/api/training/evaluate-web-exploit")
async def post_evaluate_web_exploit(req: WebExploitEvaluateRequest, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
    # IDOR FIX: Override userId from JWT
    req.userId = user["user_id"]
    return await evaluate_web_exploit(req, background_tasks)


# Re-export the row mappers so 1v1 can keep importing them from a
# stable location without taking a hard dependency on the service layer.
__all__ = [
    "map_encryption_row_to_training",
    "map_code_fixing_row_to_training",
    "map_log_analysis_row_to_training",
    "map_vuln_hunter_row_to_training",
    "map_steganography_row_to_training",
    "map_web_exploit_row_to_training",
]