#!/usr/bin/env python3
import os, sys, time, json, logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import VetModel

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
log = logging.getLogger('vetnet')

MODEL = None

class RunRequest(BaseModel):
  text: str
  species: str = '狗'

class ChatRequest(BaseModel):
  question: str
  species: str = '狗'

@asynccontextmanager
async def lifespan(app: FastAPI):
  global MODEL
  log.info('loading model...')
  t0 = time.time()
  api_key = os.getenv('DEEPSEEK_API_KEY', '')
  local_path = os.getenv('VET_LOCAL_MODEL', './vet-qwen3b-merged')
  if api_key and os.path.isdir(local_path):
    MODEL = VetModel(api_key=api_key, local_model_path=local_path)
    log.info(f'mode=pipeline loaded in {int((time.time()-t0)*1000)}ms')
  elif os.path.isdir(local_path):
    MODEL = VetModel(local_model_path=local_path)
    log.info(f'mode=local loaded in {int((time.time()-t0)*1000)}ms')
  elif api_key:
    MODEL = VetModel(api_key=api_key)
    log.info(f'mode=api loaded in {int((time.time()-t0)*1000)}ms')
  else:
    log.error('DEEPSEEK_API_KEY or VET_LOCAL_MODEL required')
    raise RuntimeError('DEEPSEEK_API_KEY or VET_LOCAL_MODEL required')
  yield
  if hasattr(MODEL, 'llm') and hasattr(MODEL.llm, 'unload'):
    MODEL.llm.unload()

app = FastAPI(title='VetNet v1.2', lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=['*'], allow_methods=['*'], allow_headers=['*'])

@app.get('/health')
def health():
  return {'status': 'ok', 'mode': MODEL.mode, 'ready': MODEL.is_ready}

@app.post('/run')
def run(req: RunRequest):
  if not MODEL: raise HTTPException(503, 'model not loaded')
  t0 = time.time()
  try:
    r = MODEL.run(req.text, req.species)
    r['server_ms'] = int((time.time() - t0) * 1000)
    return r
  except Exception as e:
    log.exception('run failed')
    raise HTTPException(500, str(e))

@app.post('/chat')
def chat(req: ChatRequest):
  if not MODEL: raise HTTPException(503, 'model not loaded')
  try:
    return {'answer': MODEL.chat(req.question, req.species)}
  except Exception as e:
    log.exception('chat failed')
    raise HTTPException(500, str(e))

if __name__ == '__main__':
  import uvicorn
  port = int(os.getenv('PORT', '8000'))
  uvicorn.run(app, host='0.0.0.0', port=port)