# main.py — Timetable SLM API (HuggingFace Spaces deployment)

import os, re, json, time
from typing import List
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# ── Globals ───────────────────────────────────────────────────────────────
MODEL     = None
TOKENIZER = None
DEVICE    = None

# These are set as HuggingFace Space Secrets
HF_REPO   = os.environ.get('HF_REPO', 'vishwasmsme/timetable-slm')
HF_TOKEN  = os.environ.get('HF_TOKEN', '')
MODEL_DIR = os.environ.get('MODEL_DIR', '/tmp/timetable_slm_model')

MAX_INPUT     = 96
MAX_OUTPUT    = 200
PROMPT_PREFIX = (
    'Convert the following timetable instruction into a structured JSON constraint. '
    'Only return valid JSON, no explanation.\n\nInstruction: '
)

CONSTRAINT_SCHEMA = {
    'FACULTY_UNAVAILABLE':        {'faculty_id':None,'days':None,'slots':None},
    'FACULTY_PREFERRED_TIME':     {'faculty_id':None,'period':None,'days':None,'priority':'SOFT'},
    'FACULTY_MAX_DAILY_HOURS':    {'faculty_id':None,'max_hours':None},
    'FACULTY_NO_CONSECUTIVE':     {'faculty_id':None,'max_consecutive':None},
    'FACULTY_FREE_DAY':           {'faculty_id':None,'preferred_day':None},
    'SUBJECT_PREFERRED_TIME':     {'subject_code':None,'period':None,'section_id':None,'priority':'SOFT'},
    'SUBJECT_SPACING':            {'subject_code':None,'section_id':None,'min_gap_days':None},
    'SUBJECT_AVOID_SLOT':         {'subject_code':None,'section_id':None,'avoid_first':False,'avoid_last':False},
    'LAB_MUST_CONSECUTIVE':       {'subject_code':None,'section_id':None,'num_slots':2},
    'HEAVY_SUBJECT_MORNING':      {'section_id':None},
    'NO_BACK_TO_BACK_SUBJECTS':   {'subject_codes':None,'section_id':None},
    'SECTION_PREFERRED_TIME':     {'section_id':None,'period':None,'priority':'SOFT'},
    'SECTION_FREE_SLOT':          {'section_id':None,'slot':None,'days':None},
    'SECTION_MAX_DAILY_SUBJECTS': {'section_id':None,'max_subjects':None},
    'LUNCH_BREAK':                {'slot':None,'sections':None},
    'ROOM_PREFERENCE':            {'subject_code':None,'faculty_id':None,'section_id':None,'room_id':None},
    'ELECTIVE_SAME_SLOT':         {'elective_group':None,'days':None},
    'DISTRIBUTE_SUBJECTS_EVENLY': {'subject_code':None,'section_id':None},
    'WORKING_DAYS':               {'days':None},
}
DAY_LOOKUP = {
    'monday':'MON','mon':'MON','tuesday':'TUE','tue':'TUE',
    'wednesday':'WED','wed':'WED','thursday':'THU','thu':'THU',
    'friday':'FRI','fri':'FRI','saturday':'SAT','sat':'SAT',
}
PERIOD_LOOKUP = {
    'morning':'MORNING','early':'MORNING','first half':'MORNING',
    'afternoon':'AFTERNOON','post-lunch':'AFTERNOON','second half':'AFTERNOON',
}

# ── Startup ───────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    global MODEL, TOKENIZER, DEVICE
    try:
        import torch
        from transformers import T5ForConditionalGeneration, T5Tokenizer
        from huggingface_hub import snapshot_download

        print(f'HF_REPO   = {HF_REPO}')
        print(f'HF_TOKEN  = {"SET ✅" if HF_TOKEN else "NOT SET ❌"}')
        print(f'MODEL_DIR = {MODEL_DIR}')

        if not os.path.exists(os.path.join(MODEL_DIR, 'config.json')):
            print('Downloading model from HuggingFace...')
            os.makedirs(MODEL_DIR, exist_ok=True)
            snapshot_download(
                repo_id=HF_REPO,
                local_dir=MODEL_DIR,
                token=HF_TOKEN if HF_TOKEN else None,
                ignore_patterns=['*.git*'],
            )
            print('✅ Model downloaded')
        else:
            print('✅ Model already cached')

        TOKENIZER = T5Tokenizer.from_pretrained(MODEL_DIR)
        MODEL     = T5ForConditionalGeneration.from_pretrained(MODEL_DIR)
        DEVICE    = 'cuda' if torch.cuda.is_available() else 'cpu'
        MODEL.to(DEVICE).eval()
        print(f'✅ Model ready on {DEVICE}')

    except Exception as e:
        print(f'❌ STARTUP ERROR: {e}')
        import traceback
        traceback.print_exc()
        raise
    yield

# ── App ───────────────────────────────────────────────────────────────────
app = FastAPI(
    title='Timetable SLM API',
    description='Converts natural language timetable instructions into JSON constraints.',
    version='1.0.0',
    lifespan=lifespan,
)
app.add_middleware(CORSMiddleware, allow_origins=['*'],
                   allow_methods=['*'], allow_headers=['*'])

# ── Request models ────────────────────────────────────────────────────────
class SingleRequest(BaseModel):
    instruction: str

class BatchRequest(BaseModel):
    instructions: List[str]

# ── Helpers ───────────────────────────────────────────────────────────────
def _infer(prompt):
    import torch
    inp = TOKENIZER(PROMPT_PREFIX + prompt, return_tensors='pt',
                    max_length=MAX_INPUT, truncation=True).to(DEVICE)
    with torch.no_grad():
        out = MODEL.generate(
            **inp, max_new_tokens=MAX_OUTPUT, num_beams=4,
            no_repeat_ngram_size=4, repetition_penalty=1.5,
            early_stopping=True,
            decoder_start_token_id=TOKENIZER.pad_token_id,
            eos_token_id=TOKENIZER.eos_token_id,
            pad_token_id=TOKENIZER.pad_token_id,
        )
    return TOKENIZER.decode(out[0], skip_special_tokens=True)

def _parse(text):
    text = text.strip()
    text = re.sub(r'^```json\s*|^```\s*|```$', '', text, flags=re.MULTILINE).strip()
    try: return json.loads(text)
    except: pass
    t = text
    if not t.startswith('{'): t = '{' + t
    if not t.endswith('}'): t = t + '}'
    try: return json.loads(t)
    except: pass
    m = re.search(r'\{.*\}', text, re.DOTALL)
    if m:
        try: return json.loads(m.group())
        except: pass
    ctype = re.search(r'"type"\s*:\s*"([^"]+)"', text)
    if ctype:
        obj = {'type': ctype.group(1)}
        for f in ['faculty_id','subject_code','section_id','elective_group','room_id']:
            m2 = re.search(f'"{f}"\\s*:\\s*"([^"]+)"', text)
            if m2: obj[f] = m2.group(1)
        for f in ['max_hours','min_gap_days','num_slots','slot','max_subjects','max_consecutive']:
            m2 = re.search(f'"{f}"\\s*:\\s*(\\d+)', text)
            if m2: obj[f] = int(m2.group(1))
        return {'constraints': [obj]}
    return None

def _fix(c, prompt):
    c = dict(c)
    ctype = c.get('type','').upper()
    c['type'] = ctype
    if 'DAYS' in c and 'days' not in c: c['days'] = c.pop('DAYS')
    if ctype in ('FACULTY_PREFERRED_TIME','SUBJECT_PREFERRED_TIME','SECTION_PREFERRED_TIME'):
        if not c.get('priority'): c['priority'] = 'SOFT'
        if not c.get('period'):
            p = prompt.lower()
            for w, v in PERIOD_LOOKUP.items():
                if w in p: c['period'] = v; break
    if ctype in ('FACULTY_UNAVAILABLE','FACULTY_PREFERRED_TIME') and not c.get('days'):
        found = []
        for w, code in DAY_LOOKUP.items():
            if w in prompt.lower() and code not in found: found.append(code)
        if found: c['days'] = found
    if ctype == 'FACULTY_FREE_DAY' and not c.get('preferred_day'):
        for w, code in DAY_LOOKUP.items():
            if w in prompt.lower(): c['preferred_day'] = code; break
    if ctype == 'SUBJECT_AVOID_SLOT':
        if c.get('avoid_first') is None: c['avoid_first'] = False
        if c.get('avoid_last')  is None: c['avoid_last']  = False
    if ctype == 'LAB_MUST_CONSECUTIVE' and not c.get('num_slots'):
        c['num_slots'] = 2
    for field, default in CONSTRAINT_SCHEMA.get(ctype, {}).items():
        if field not in c: c[field] = default
    return c

def _process(instruction):
    t0  = time.time()
    raw = _infer(instruction)
    parsed = _parse(raw)
    ms = round((time.time() - t0) * 1000)
    if not parsed:
        return {'success':False,'instruction':instruction,
                'error':'parse_failed','raw':raw,'latency_ms':ms}
    fixed = [_fix(c, instruction) for c in parsed.get('constraints', [])]
    return {'success':True,'instruction':instruction,
            'constraints':fixed,'raw':raw,'latency_ms':ms}

# ── Endpoints ─────────────────────────────────────────────────────────────
@app.get('/')
def root():
    return {
        'service': 'Timetable SLM API',
        'status': 'running',
        'model': HF_REPO,
        'device': DEVICE,
        'docs': 'Visit /docs for interactive API documentation',
    }

@app.get('/health')
def health():
    return {'status':'ok','model_loaded':MODEL is not None,'device':DEVICE}

@app.post('/constraint')
def single(req: SingleRequest):
    if MODEL is None: raise HTTPException(503, 'Model not loaded')
    if not req.instruction.strip(): raise HTTPException(400, 'Empty instruction')
    return _process(req.instruction)

@app.post('/constraints/batch')
def batch(req: BatchRequest):
    if MODEL is None: raise HTTPException(503, 'Model not loaded')
    if not req.instructions: raise HTTPException(400, 'Empty list')
    results = [_process(i) for i in req.instructions]
    all_c   = [c for r in results if r['success'] for c in r['constraints']]
    failed  = [r['instruction'] for r in results if not r['success']]
    return {
        'success': len(failed) == 0,
        'total': len(results),
        'parsed': len(results) - len(failed),
        'failed': len(failed),
        'failed_list': failed,
        'all_constraints': all_c,
        'results': results,
    }