from contextlib import asynccontextmanager
from pathlib import Path
import logging
from fastapi import Depends, FastAPI, Header, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, Field

import security
import site_memory
from trainer import trainer
from text_utils import load_training_lines

BASE = Path(__file__).resolve().parent
log = logging.getLogger('hooshmark')

MSG_UNAVAILABLE = 'دستیار فعلاً آمادهٔ پاسخ‌گویی نیست. لطفاً کمی بعد دوباره تلاش کنید.'
MSG_NO_ANSWER = 'پاسخی برای این پیام پیدا نشد. لطفاً آن را به شکل دیگری بنویسید.'
MODEL_NAME = 'tiny-transformer-from-scratch'


@asynccontextmanager
async def lifespan(_app):
    # بدون کلید API معتبر سرور بالا نمی‌آید (پیام خطا راه ساخت و دادن کلید را می‌گوید).
    security.configured_key()
    yield


app = FastAPI(title='HooshMark Native AI', version='0.4.0', lifespan=lifespan)
app.mount('/static', StaticFiles(directory=BASE / 'static'), name='static')
templates = Jinja2Templates(directory=str(BASE / 'templates'))


class ApiError(Exception):
    def __init__(self, status: int, code: str, message: str, headers: dict | None = None):
        self.status, self.code, self.message, self.headers = status, code, message, headers


@app.exception_handler(ApiError)
async def api_error_handler(_request: Request, exc: ApiError):
    return JSONResponse({'error': exc.message, 'code': exc.code}, status_code=exc.status, headers=exc.headers)


def require_api_key(authorization: str | None = Header(default=None)):
    """هدر `Authorization: Bearer <key>` باید با HOOSHMARK_API_KEY یکی باشد."""
    try:
        expected = security.configured_key()
    except security.ApiKeyConfigError:
        log.error('%s تنظیم نشده یا کوتاه است؛ درخواست رد شد.', security.ENV_NAME)
        raise ApiError(503, 'API_KEY_NOT_CONFIGURED', 'کلید API روی سرویس تنظیم نشده است.')
    if not security.is_valid(security.extract_bearer(authorization), expected):
        raise ApiError(401, 'INVALID_API_KEY', 'کلید API نامعتبر یا ارسال‌نشده است.', {'WWW-Authenticate': 'Bearer'})


class ChatInput(BaseModel):
    message: str = Field(min_length=1, max_length=4000)


class MemoryEntry(BaseModel):
    # طول را اینجا محدود نمی‌کنیم: یک عبارت بلند نباید کل همگام‌سازی را رد کند؛
    # site_memory عبارت‌های بلندتر از BLOCK را کنار می‌گذارد و در گزارش (tooLong) می‌شمارد.
    question: str
    answer: str


class MemoryInput(BaseModel):
    entries: list[MemoryEntry] = Field(max_length=200000)

MIN_TRAIN_LINES = 5  # هم‌اندازهٔ حداقل مورد انتظار train.py


@app.get('/', response_class=HTMLResponse)
def home(request: Request):
    return templates.TemplateResponse(request=request, name='index.html', context={})


@app.get('/chat', response_class=HTMLResponse)
def chat_page(request: Request):
    return templates.TemplateResponse(request=request, name='chat.html', context={})


@app.get('/api/health')
def health():
    from model_runtime import is_trained
    return {'ok': True, 'model': MODEL_NAME, 'trained': is_trained()}


@app.get('/api/verify', dependencies=[Depends(require_api_key)])
def verify():
    """برای بررسی اتصال و درستی کلید؛ با کلید درست ۲۰۰ و با کلید نادرست ۴۰۱."""
    from model_runtime import is_trained
    return {'ok': True, 'model': MODEL_NAME, 'trained': is_trained()}


@app.get('/api/memory', dependencies=[Depends(require_api_key)])
def memory_info():
    """گزارش آخرین همگام‌سازی حافظهٔ سایت."""
    return site_memory.info()


@app.put('/api/memory', dependencies=[Depends(require_api_key)])
def memory_replace(body: MemoryInput):
    """حافظهٔ سایت را کامل جایگزین می‌کند (برای اثر روی مدل، بعدش /api/train لازم است)."""
    return site_memory.save([e.model_dump() for e in body.entries])


@app.post('/api/train', dependencies=[Depends(require_api_key)], status_code=202)
def train_start():
    """آموزش را در پس‌زمینه شروع می‌کند؛ پیشرفت از /api/train/status."""
    if trainer.is_running():
        raise ApiError(409, 'TRAINING_IN_PROGRESS', 'آموزش در حال انجام است.')
    if len(load_training_lines()) < MIN_TRAIN_LINES:
        raise ApiError(400, 'NOT_ENOUGH_DATA', f'برای آموزش دست‌کم {MIN_TRAIN_LINES} خط داده لازم است.')
    trainer.start()
    return trainer.status()


@app.get('/api/train/status', dependencies=[Depends(require_api_key)])
def train_status():
    from model_runtime import is_trained
    return {**trainer.status(), 'trained': is_trained()}


@app.post('/api/chat', dependencies=[Depends(require_api_key)])
def chat(body: ChatInput):
    from model_runtime import is_trained
    if not is_trained():
        log.warning('مدل آموزش ندیده است؛ ابتدا python train.py را اجرا کنید.')
        return {'reply': MSG_UNAVAILABLE, 'mode': 'untrained'}
    try:
        from model_runtime import reply
        answer = reply(body.message)
    except Exception:
        log.exception('خطا در تولید پاسخ')
        return {'reply': MSG_UNAVAILABLE, 'mode': 'error'}
    if not answer:
        return {'reply': MSG_NO_ANSWER, 'mode': 'empty'}
    return {'reply': answer, 'mode': 'model'}