import json
import os
import time
import threading

import torch

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer


# ============================================================
# Configuration
# ============================================================

MODEL_ID = os.getenv(
    "MODEL_ID",
    "mjpsm/activity-generation-model-v1",
)

MAX_NEW_TOKENS = int(
    os.getenv("MAX_NEW_TOKENS", "300")
)

# Optional pricing.
#
# Example:
# INPUT_PRICE_PER_1K_TOKENS=0.001
# OUTPUT_PRICE_PER_1K_TOKENS=0.002
#
# Keep these at 0 until you decide on pricing.

INPUT_PRICE_PER_1K_TOKENS = float(
    os.getenv("INPUT_PRICE_PER_1K_TOKENS", "0.001")
)

OUTPUT_PRICE_PER_1K_TOKENS = float(
    os.getenv("OUTPUT_PRICE_PER_1K_TOKENS", "0.005")
)


# ============================================================
# FastAPI
# ============================================================

app = FastAPI(
    title="MyVillage Activity Generation API",
    description=(
        "Generate a student's next learning activity from their "
        "village goal, previous activity, and knowledge submission."
    ),
    version="1.0.0",
)


# ============================================================
# System Prompt
# ============================================================

SYSTEM_PROMPT = """You are an educational activity generator for MyVillage.

Your job is to create exactly one logical next learning activity for a student.

You will receive:

1. The goal of the student's village.
2. The title of the student's previous activity.
3. The student's knowledge submission describing what they learned or completed.

Create a new activity that:

- directly builds on the student's knowledge submission;
- moves the student toward the village goal;
- does not simply repeat the previous activity;
- is specific and actionable;
- uses clear student-facing language;
- includes a concrete task or deliverable.

Return valid JSON only.

Return exactly these fields:

{
  "title": "activity title",
  "description": "activity description",
  "instructions": "activity instructions"
}

Do not include markdown.
Do not include commentary.
Do not include additional fields.
"""


# ============================================================
# Load model ONCE
# ============================================================

print(f"Loading model: {MODEL_ID}")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
)

model.to("cpu")
model.eval()

print("Model loaded successfully.")


# Prevent multiple CPU generations from competing for memory.
generation_lock = threading.Lock()


# ============================================================
# Request / Response Models
# ============================================================

class ActivityRequest(BaseModel):
    village_goal: str = Field(
        ...,
        min_length=1,
        description="The overall goal of the student's village.",
    )

    previous_activity_title: str = Field(
        ...,
        min_length=1,
        description="The title of the student's previous activity.",
    )

    knowledge_submission: str = Field(
        ...,
        min_length=1,
        description="What the student learned or completed.",
    )


class Activity(BaseModel):
    title: str
    description: str
    instructions: str


class TokenUsage(BaseModel):
    input_tokens: int
    output_tokens: int
    total_tokens: int


class CostEstimate(BaseModel):
    input_cost: float
    output_cost: float
    total_cost: float
    currency: str = "USD"


class ActivityResponse(BaseModel):
    activity: Activity
    usage: TokenUsage
    estimated_cost: CostEstimate
    generation_time_seconds: float


# ============================================================
# Helper Functions
# ============================================================

def build_prompt(request: ActivityRequest):

    user_message = f"""Village goal:
{request.village_goal}

Previous activity:
{request.previous_activity_title}

Knowledge submission:
{request.knowledge_submission}

Create the student's next activity."""

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": user_message,
        },
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    return prompt


def calculate_cost(
    input_tokens: int,
    output_tokens: int,
):

    input_cost = (
        input_tokens / 1000
    ) * INPUT_PRICE_PER_1K_TOKENS

    output_cost = (
        output_tokens / 1000
    ) * OUTPUT_PRICE_PER_1K_TOKENS

    total_cost = input_cost + output_cost

    return {
        "input_cost": round(input_cost, 8),
        "output_cost": round(output_cost, 8),
        "total_cost": round(total_cost, 8),
        "currency": "USD",
    }


# ============================================================
# Routes
# ============================================================

@app.get("/")
def root():

    return {
        "name": "MyVillage Activity Generation API",
        "model": MODEL_ID,
        "status": "running",
        "docs": "/docs",
    }


@app.get("/health")
def health():

    return {
        "status": "healthy",
        "model": MODEL_ID,
        "model_loaded": True,
    }


@app.post(
    "/generate",
    response_model=ActivityResponse,
)
def generate_activity(
    request: ActivityRequest,
):

    start_time = time.perf_counter()

    prompt = build_prompt(request)

    # --------------------------------------------------------
    # Tokenize input
    # --------------------------------------------------------

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    )

    input_tokens = inputs["input_ids"].shape[1]

    # --------------------------------------------------------
    # Generate
    # --------------------------------------------------------

    try:

        with generation_lock:

            with torch.inference_mode():

                outputs = model.generate(
                    **inputs,
                    max_new_tokens=MAX_NEW_TOKENS,
                    do_sample=False,
                    repetition_penalty=1.05,
                    pad_token_id=tokenizer.pad_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                )

    except Exception as error:

        raise HTTPException(
            status_code=500,
            detail=f"Model generation failed: {str(error)}",
        )

    # --------------------------------------------------------
    # Separate output from prompt
    # --------------------------------------------------------

    generated_tokens = outputs[
        0,
        input_tokens:
    ]

    output_tokens = generated_tokens.shape[0]

    total_tokens = (
        input_tokens
        + output_tokens
    )

    # --------------------------------------------------------
    # Decode model response
    # --------------------------------------------------------

    response_text = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    ).strip()

    # --------------------------------------------------------
    # Parse JSON
    # --------------------------------------------------------

    try:

        activity_data = json.loads(
            response_text
        )

    except json.JSONDecodeError:

        raise HTTPException(
            status_code=500,
            detail={
                "message": (
                    "Model did not return valid JSON."
                ),
                "raw_output": response_text,
                "usage": {
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "total_tokens": total_tokens,
                },
            },
        )

    required_fields = {
        "title",
        "description",
        "instructions",
    }

    if set(activity_data.keys()) != required_fields:

        raise HTTPException(
            status_code=500,
            detail={
                "message": (
                    "Model returned an invalid schema."
                ),
                "raw_output": activity_data,
            },
        )

    # --------------------------------------------------------
    # Pricing
    # --------------------------------------------------------

    cost = calculate_cost(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )

    generation_time = (
        time.perf_counter()
        - start_time
    )

    # --------------------------------------------------------
    # Response
    # --------------------------------------------------------

    return {
        "activity": activity_data,

        "usage": {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
        },

        "estimated_cost": cost,

        "generation_time_seconds": round(
            generation_time,
            3,
        ),
    }