import json
from typing import Optional
from datetime import datetime, timezone
from pydantic import BaseModel, Field
from pydantic.json_schema import SkipJsonSchema

vector_scaling = json.load(open("assets/vector_scaling.json", "r"))

CONFIG = {
    "max_new_tokens": 2400,
    "top_p": 0.95,
    "temperature": 0.6,
    "vec_scale": 1.5,
    "layer": 25
}

class UserRequest(BaseModel):
    session_id: str
    prompt: str = None
    steering: bool = True
    coeff: float = -1.0
    max_new_tokens: int = Field(CONFIG["max_new_tokens"], le=3048)
    top_p: float = Field(CONFIG["top_p"], ge=0.0, le=1.0)
    temperature: float = Field(CONFIG["temperature"], ge=0.0, le=1.0)
    vec_scale: float = Field(CONFIG["vec_scale"])
    k: Optional[float] = None
    layer: int = Field(CONFIG["layer"])

    def model_post_init(self, __context):
        if self.coeff < 0:
            self.k = self.vec_scale * vector_scaling[self.layer]["k_pos"]
        else:
            self.k = self.vec_scale * vector_scaling[self.layer]["k_neg"]

    def generation_config(self):
        return {
            "max_new_tokens": self.max_new_tokens,
            "top_p": self.top_p,
            "temperature": self.temperature
        }


class SteeringOutput(UserRequest):
    max_new_tokens: SkipJsonSchema[int] = Field(exclude=True)
    vec_scale: SkipJsonSchema[float] = Field(exclude=True)
    reasoning: str = None
    answer: Optional[str] = None
    upvote: Optional[bool] = None
    timestamp: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())