# worker.py — FastAPI generation worker for HuggingFace Docker Space
# Receives job requests, streams from OpenRouter, calls back PythonAnywhere when done.

import os
import re
import json
import time
import logging
import threading
import requests

from datetime import datetime
from fastapi import FastAPI, BackgroundTasks, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

# ============================================================
#  LOGGING
# ============================================================
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s'
)
logger = logging.getLogger(__name__)

app = FastAPI()

# ============================================================
#  CONFIG — set these as HF Space secrets
# ============================================================
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
OPENROUTER_BASE    = "https://openrouter.ai/api/v1"
OPENROUTER_MODEL = os.environ.get(
    "OPENROUTER_MODEL",
    "meta-llama/llama-3-8b-instruct:free"  # confirmed free tier
)

# PythonAnywhere callback — HF Space POSTs results here when done
# Set this as a secret: e.g. https://yourname.pythonanywhere.com/api/generation-callback
CALLBACK_BASE = os.environ.get("CALLBACK_BASE", "https://smilystudy.pythonanywhere.com")

# Shared secret so PA knows the callback is legit
CALLBACK_SECRET = os.environ.get("CALLBACK_SECRET", "change-me-please")

MAX_TOKENS   = 16000
TEMPERATURE  = 0.55
TOP_P        = 0.9

COURSE_SYSTEM_PROMPT = """You are SupaNova, an expert AI curriculum designer built by SmilyAI Labs.

A user will tell you a topic they want to learn. Your job:

1. Inside <think> tags, reason step by step about: the right scope for a 
   beginner-to-intermediate learner, a logical order of lessons, and what 
   each lesson should cover. Be thorough but efficient.

2. After closing </think>, output ONLY a single valid JSON object — nothing 
   else. No markdown code fences, no commentary, no text before or after it.
   The JSON must exactly follow this structure:

{
  "title": "Course title",
  "description": "A one or two sentence description of the course",
  "lessons": [
    {
      "title": "Lesson title",
      "content": "Full lesson content in markdown, at least 300 words, 
                  using ## headers, bullet points, and concrete examples."
    }
  ]
}

Include between 5 and 10 lessons. Make every lesson genuinely educational,
well organised, and self-contained. Output valid JSON only after </think>."""

# ============================================================
#  IN-MEMORY JOB STATE  (HF Spaces don't restart randomly like PA)
# ============================================================
# job_id -> {"status": "running"|"done"|"error", "result": ..., "error": ...}
_jobs: dict[str, dict] = {}
_jobs_lock = threading.Lock()


def set_job(job_id: str, data: dict):
    with _jobs_lock:
        _jobs[job_id] = data


def get_job(job_id: str) -> dict | None:
    with _jobs_lock:
        return _jobs.get(job_id)


def update_job(job_id: str, **kwargs):
    with _jobs_lock:
        if job_id in _jobs:
            _jobs[job_id].update(kwargs)


# ============================================================
#  HELPERS
# ============================================================

def openrouter_headers() -> dict:
    return {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type":  "application/json",
        "HTTP-Referer":  "https://smilystudy.app",
        "X-Title":       "SmilyStudy",
    }


class ThinkFilterStream:
    """Strips <think>…</think> blocks and returns only the final JSON text."""

    def __init__(self):
        self.buffer   = ""
        self.in_think = False

    def feed(self, new_chars: str):
        self.buffer += new_chars
        yield from self._process()

    def _process(self):
        while self.buffer:
            if self.in_think:
                end = self.buffer.find("</think>")
                if end == -1:
                    # Keep a safe tail in case the tag is split across chunks
                    safe = self.buffer.rfind("<")
                    if safe > 0:
                        self.buffer = self.buffer[safe:]
                    return
                self.in_think = False
                self.buffer   = self.buffer[end + len("</think>"):]
            else:
                start = self.buffer.find("<think>")
                if start == -1:
                    safe = self.buffer.rfind("<")
                    if safe == -1:
                        yield self.buffer
                        self.buffer = ""
                    else:
                        yield self.buffer[:safe]
                        self.buffer = self.buffer[safe:]
                    return
                yield self.buffer[:start]
                self.in_think = True
                self.buffer   = self.buffer[start + len("<think>"):]

    def flush(self):
        if self.buffer and not self.in_think:
            yield self.buffer
        self.buffer = ""


def extract_json(text: str) -> dict | None:
    text = text.strip()
    # Strip markdown fences if any
    text = re.sub(r"^```(?:json)?", "", text, flags=re.IGNORECASE).strip()
    text = re.sub(r"```$", "", text).strip()

    # Direct parse
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Find first { … last }
    start = text.find("{")
    end   = text.rfind("}")
    if start != -1 and end > start:
        try:
            return json.loads(text[start:end + 1])
        except json.JSONDecodeError as e:
            logger.error(f"extract_json fallback failed: {e}")

    return None


def _send_callback(job_id: str, payload: dict):
    """POST result back to PythonAnywhere. Fire-and-forget with retries."""
    if not CALLBACK_BASE:
        logger.warning("CALLBACK_BASE not set — skipping callback")
        return

    url = f"{CALLBACK_BASE}/api/generation-callback"
    headers = {
        "Content-Type":    "application/json",
        "X-Callback-Secret": CALLBACK_SECRET,
    }

    for attempt in range(5):
        try:
            resp = requests.post(url, json=payload, headers=headers, timeout=30)
            if resp.status_code == 200:
                logger.info(f"[{job_id}] Callback delivered on attempt {attempt+1}")
                return
            logger.warning(
                f"[{job_id}] Callback attempt {attempt+1} got {resp.status_code}"
            )
        except Exception as e:
            logger.warning(f"[{job_id}] Callback attempt {attempt+1} error: {e}")

        time.sleep(2 ** attempt)   # exponential back-off: 1s, 2s, 4s, 8s, 16s

    logger.error(f"[{job_id}] All callback attempts failed")


# ============================================================
#  CORE GENERATION  (runs in a daemon thread inside HF Space)
# ============================================================

def _run_generation(job_id: str, user_id: int, topic: str):
    logger.info(f"[{job_id}] === GENERATION START === topic={topic!r}")

    def finish_error(message: str):
        update_job(job_id, status="error", error=message)
        _send_callback(job_id, {
            "job_id":  job_id,
            "user_id": user_id,
            "status":  "error",
            "error":   message,
        })

    def finish_ok(course_data: dict):
        update_job(job_id, status="done", result=course_data)
        _send_callback(job_id, {
            "job_id":      job_id,
            "user_id":     user_id,
            "status":      "done",
            "course_data": course_data,
        })

    try:
        if not OPENROUTER_API_KEY:
            finish_error("OpenRouter API key not configured on worker.")
            return

        payload = {
            "model":    OPENROUTER_MODEL,
            "stream":   False,  # ← NO STREAMING
            "messages": [
                {"role": "system", "content": COURSE_SYSTEM_PROMPT},
                {"role": "user",   "content": f"I want to learn: {topic}"},
            ],
            "max_tokens":  MAX_TOKENS,
            "temperature": TEMPERATURE,
            "top_p":       TOP_P,
        }

        logger.info(f"[{job_id}] Calling OpenRouter (no stream)...")

        resp = requests.post(
            f"{OPENROUTER_BASE}/chat/completions",
            headers=openrouter_headers(),
            json=payload,
            timeout=600,  # 10 min timeout, one shot
        )

        logger.info(f"[{job_id}] HTTP {resp.status_code}")

        if resp.status_code != 200:
            finish_error(f"API error {resp.status_code}: {resp.text[:300]}")
            return

        data = resp.json()

        # Check for API-level error
        if "error" in data:
            err = data["error"]
            msg = err.get("message", str(err)) if isinstance(err, dict) else str(err)
            logger.error(f"[{job_id}] API error: {msg}")
            finish_error(f"API error: {msg}")
            return

        # Extract content from response
        choices = data.get("choices") or []
        if not choices:
            finish_error("No response from model.")
            return

        content = choices[0].get("message", {}).get("content", "")
        if not content:
            finish_error("Model returned empty content.")
            return

        logger.info(f"[{job_id}] Got {len(content)} chars from model")
        logger.debug(f"[{job_id}] First 500 chars: {content[:500]}")

        # Parse JSON
        course_data = extract_json(content)
        if not course_data or "lessons" not in course_data:
            logger.error(f"[{job_id}] Failed to parse JSON")
            logger.debug(f"[{job_id}] Full content:\n{content}")
            finish_error("Model returned invalid JSON. Please try again.")
            return

        if not course_data.get("lessons"):
            finish_error("No lessons were generated.")
            return

        logger.info(f"[{job_id}] ✓ {len(course_data['lessons'])} lessons parsed")
        finish_ok(course_data)

    except requests.exceptions.Timeout:
        logger.error(f"[{job_id}] Request timeout (10 min)")
        finish_error("Generation timed out. Please try again.")
    except Exception as e:
        logger.exception(f"[{job_id}] Unhandled exception")
        finish_error(f"Worker error: {e}")
    finally:
        logger.info(f"[{job_id}] === THREAD EXIT ===")


# ============================================================
#  FASTAPI ENDPOINTS
# ============================================================

class StartRequest(BaseModel):
    job_id:  str          # PA generates this (e.g. uuid4 hex)
    user_id: int
    topic:   str
    secret:  str          # Must match CALLBACK_SECRET


@app.post("/start")
async def start_job(req: StartRequest, background_tasks: BackgroundTasks):
    """PA calls this to kick off a generation job."""

    if req.secret != CALLBACK_SECRET:
        raise HTTPException(status_code=403, detail="Bad secret")

    if not req.topic.strip():
        raise HTTPException(status_code=400, detail="Topic is empty")

    if len(req.topic) > 300:
        raise HTTPException(status_code=400, detail="Topic too long")

    existing = get_job(req.job_id)
    if existing:
        return JSONResponse({"ok": True, "note": "job already exists"})

    set_job(req.job_id, {
        "status":   "running",
        "user_id":  req.user_id,
        "topic":    req.topic,
        "started":  datetime.utcnow().isoformat(),
    })

    # FastAPI BackgroundTasks = proper async threading, no PA issues
    background_tasks.add_task(_run_generation, req.job_id, req.user_id, req.topic)

    logger.info(f"Job {req.job_id} queued for user {req.user_id}, topic={req.topic!r}")
    return JSONResponse({"ok": True, "job_id": req.job_id})


@app.get("/status/{job_id}")
async def job_status(job_id: str, secret: str = ""):
    """PA can poll this as a fallback if the callback is missed."""

    if secret != CALLBACK_SECRET:
        raise HTTPException(status_code=403, detail="Bad secret")

    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")

    # Don't send the full course_data in status — only signal done/error
    response = {
        "job_id": job_id,
        "status": job["status"],
    }
    if job["status"] == "error":
        response["error"] = job.get("error", "Unknown error")
    if job["status"] == "done":
        response["has_result"] = True

    return JSONResponse(response)


@app.get("/result/{job_id}")
async def job_result(job_id: str, secret: str = ""):
    """PA fetches full course data after receiving a 'done' callback."""

    if secret != CALLBACK_SECRET:
        raise HTTPException(status_code=403, detail="Bad secret")

    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")
    if job["status"] != "done":
        raise HTTPException(status_code=409, detail=f"Job is {job['status']}, not done")

    return JSONResponse({
        "job_id":      job_id,
        "user_id":     job["user_id"],
        "course_data": job["result"],
    })


@app.get("/health")
async def health():
    return {"ok": True, "jobs": len(_jobs)}