"""
ARCHAI Adaptive Assessment API
==============================
FastAPI backend that plugs into the archai frontend.

Endpoints:
  POST /api/v1/session/start      → Initialize assessment
  POST /api/v1/session/answer   → Submit answer, get next question
  GET  /api/v1/session/{id}      → Get current state/results
  POST /api/v1/path/generate     → Generate learning path with day/week/month actionables
  GET  /api/v1/questions         → Get full question bank (for offline study)
  GET  /api/v1/health            → Health check
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from adaptive_engine import AdaptiveAssessmentEngine, engine, Dimension, DIMENSION_LABELS, DIMENSION_COLORS

app = FastAPI(
    title="ARCHAI Adaptive Assessment Engine",
    description="IRT-based adaptive AI readiness assessment with LLM-powered learning paths",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

# CORS for your Netlify frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, restrict to your domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================================
# REQUEST/RESPONSE SCHEMAS
# ============================================================================

class StartSessionResponse(BaseModel):
    session_id: str
    question: Optional[Dict[str, Any]]
    progress: Dict[str, Any]
    status: str

class SubmitAnswerRequest(BaseModel):
    session_id: str
    question_id: str
    option_index: int = Field(ge=0, le=3)

class SubmitAnswerResponse(BaseModel):
    session_id: str
    question: Optional[Dict[str, Any]]
    progress: Dict[str, Any]
    interim_scores: Optional[Dict[str, int]]
    status: str

class GeneratePathRequest(BaseModel):
    session_id: str
    persona_id: str
    hours_per_week: int = Field(ge=1, le=40)
    budget_usd: int = Field(ge=0)
    hardware_id: Optional[str] = None
    preference: Optional[str] = None  # "local", "api", "both"

class GeneratePathResponse(BaseModel):
    session_id: str
    overall_score: int
    stage: Dict[str, Any]
    archetype: Dict[str, Any]
    dimension_scores: Dict[str, int]
    gaps: List[Dict[str, Any]]
    strengths: List[Dict[str, Any]]
    learning_path: Dict[str, List[Any]]
    projections: Dict[str, Any]
    meta: Dict[str, Any]

# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.get("/api/v1/health")
def health():
    return {
        "status": "healthy",
        "engine": "IRT-2PL adaptive",
        "version": "2.0.0",
        "features": ["adaptive_selection", "bayesian_knowledge_tracing", "structured_learning_paths"],
    }

@app.post("/api/v1/session/start", response_model=StartSessionResponse)
def start_session():
    """Start a new adaptive assessment session."""
    result = engine.start_session()
    return result

@app.post("/api/v1/session/answer", response_model=SubmitAnswerResponse)
def submit_answer(req: SubmitAnswerRequest):
    """Submit an answer and get the next adaptive question."""
    result = engine.submit_answer(req.session_id, req.question_id, req.option_index)
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return result

@app.get("/api/v1/session/{session_id}")
def get_session(session_id: str):
    """Get current session state or final results."""
    state = engine.sessions.get(session_id)
    if not state:
        raise HTTPException(status_code=404, detail="Session not found")
    
    # If complete, return results
    if len(state.asked_questions) > 0 and engine.selector.should_stop(state):
        return engine._finalize(state)
    
    # Otherwise return current state
    dim_coverage = set()
    for qid in state.asked_questions:
        q = next((qq for qq in engine.question_bank if qq.id == qid), None)
        if q:
            dim_coverage.add(q.dimension.value)
    
    return {
        "session_id": session_id,
        "status": "in_progress",
        "progress": {
            "asked": len(state.asked_questions),
            "total": 12,
            "dimensions_covered": list(dim_coverage),
        },
        "interim_scores": engine.tracer.get_dimension_scores(state),
        "latent_abilities": {d.value: round(t, 2) for d, t in state.theta.items()},
    }

@app.post("/api/v1/path/generate", response_model=GeneratePathResponse)
def generate_path(req: GeneratePathRequest):
    """Generate a structured learning path with day/week/month actionables."""
    result = engine.generate_path(
        req.session_id,
        req.persona_id,
        req.hours_per_week,
        req.budget_usd,
        req.hardware_id,
        req.preference,
    )
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return result

@app.get("/api/v1/questions")
def get_questions():
    """Get the full calibrated question bank for offline study."""
    return {
        "questions": [
            {
                "id": q.id,
                "dimension": q.dimension.value,
                "dimension_label": DIMENSION_LABELS.get(q.dimension, q.dimension.value),
                "text": q.text,
                "options": q.options,
                "difficulty": round(q.difficulty, 2),
                "discrimination": round(q.discrimination, 2),
                "concept_tags": q.concept_tags,
            }
            for q in engine.question_bank
        ],
        "dimensions": [
            {"id": d.value, "label": DIMENSION_LABELS.get(d, d.value), "color": DIMENSION_COLORS.get(d.value, "#14B8A6")}
            for d in Dimension
        ],
    }

@app.get("/api/v1/dimensions")
def get_dimensions():
    """Get dimension metadata for UI rendering."""
    return {
        "dimensions": [
            {"id": d.value, "label": DIMENSION_LABELS.get(d, d.value), "color": DIMENSION_COLORS.get(d.value, "#14B8A6")}
            for d in Dimension
        ]
    }

# ============================================================================
# COMPATIBILITY ENDPOINTS — archai v1 API mapping
# ============================================================================

@app.post("/api/v1/assessment/start")
def assessment_start_compat():
    """Backward-compatible endpoint name."""
    return start_session()

@app.post("/api/v1/assessment/answer")
def assessment_answer_compat(req: SubmitAnswerRequest):
    """Backward-compatible endpoint name."""
    return submit_answer(req)

@app.get("/api/v1/assessment/results/{session_id}")
def assessment_results_compat(session_id: str):
    """Backward-compatible endpoint for getting results."""
    return get_session(session_id)

# ============================================================================
# Root
# ============================================================================

@app.get("/")
def root():
    return {
        "service": "ARCHAI Adaptive Assessment Engine",
        "version": "2.0.0",
        "docs": "/docs",
        "endpoints": {
            "start": "POST /api/v1/session/start",
            "answer": "POST /api/v1/session/answer",
            "results": "GET /api/v1/session/{session_id}",
            "path": "POST /api/v1/path/generate",
            "questions": "GET /api/v1/questions",
            "health": "GET /api/v1/health",
        },
        "adaptive_features": {
            "irt_model": "2PL (Two-Parameter Logistic)",
            "selection": "Fisher Information Maximization",
            "tracing": "Bayesian Knowledge Updating",
            "stopping": "Precision-based early stopping",
            "question_count": "Adaptive 6-12 questions",
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)