#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# /// script
# dependencies = ["torch","transformers","huggingface_hub","hf_xet","numpy"]
# ///

import argparse, math, json
from dataclasses import dataclass
from contextlib import nullcontext

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2TokenizerFast
from huggingface_hub import hf_hub_download


@dataclass
class GPTConfig:
    vocab_size:int
    block_size:int
    n_layer:int
    n_head:int
    n_embd:int
    dropout:float=0.0
    bias:bool=False
    pcs_a:float=0.8309193524478643
    pcs_b:float=0.0


class PCS(nn.Module):
    def __init__(self,a=0.8309193524478643,b=0.0):
        super().__init__()
        self.a=float(a)
        self.b=float(b)
    def forward(self,x):
        return x*torch.sin(self.a*x)+self.b*torch.cos(x)


class Attn(nn.Module):
    def __init__(self,c):
        super().__init__()
        assert c.n_embd % c.n_head == 0
        self.n_head=c.n_head
        self.head_dim=c.n_embd//c.n_head
        self.dropout=c.dropout
        self.qkv=nn.Linear(c.n_embd,3*c.n_embd,bias=c.bias)
        self.proj=nn.Linear(c.n_embd,c.n_embd,bias=c.bias)
        self.drop=nn.Dropout(c.dropout)
    def forward(self,x):
        B,T,C=x.shape
        q,k,v=self.qkv(x).split(C,dim=2)
        q=q.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        k=k.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        v=v.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        y=F.scaled_dot_product_attention(q,k,v,dropout_p=self.dropout if self.training else 0.0,is_causal=True)
        y=y.transpose(1,2).contiguous().view(B,T,C)
        return self.drop(self.proj(y))


class MLP(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.fc1=nn.Linear(c.n_embd,4*c.n_embd,bias=c.bias)
        self.act=PCS(c.pcs_a,c.pcs_b)
        self.fc2=nn.Linear(4*c.n_embd,c.n_embd,bias=c.bias)
        self.drop=nn.Dropout(c.dropout)
    def forward(self,x):
        return self.drop(self.fc2(self.act(self.fc1(x))))


class Block(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.ln1=nn.LayerNorm(c.n_embd)
        self.attn=Attn(c)
        self.ln2=nn.LayerNorm(c.n_embd)
        self.mlp=MLP(c)
    def forward(self,x):
        x=x+self.attn(self.ln1(x))
        x=x+self.mlp(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.cfg=c
        self.tok_emb=nn.Embedding(c.vocab_size,c.n_embd)
        self.pos_emb=nn.Embedding(c.block_size,c.n_embd)
        self.drop=nn.Dropout(c.dropout)
        self.blocks=nn.ModuleList([Block(c) for _ in range(c.n_layer)])
        self.ln_f=nn.LayerNorm(c.n_embd)
        self.lm_head=nn.Linear(c.n_embd,c.vocab_size,bias=False)
        self.tok_emb.weight=self.lm_head.weight
    def forward(self,idx):
        B,T=idx.shape
        if T>self.cfg.block_size:
            idx=idx[:,-self.cfg.block_size:]
            B,T=idx.shape
        pos=torch.arange(T,device=idx.device)
        x=self.drop(self.tok_emb(idx)+self.pos_emb(pos))
        for b in self.blocks:
            x=b(x)
        return self.lm_head(self.ln_f(x))


def dtype_of(dev):
    if dev=="cuda" and torch.cuda.is_bf16_supported():
        return "bfloat16"
    if dev=="cuda":
        return "float16"
    return "float32"


def ac(dev,dt):
    if dev!="cuda" or dt=="float32":
        return nullcontext()
    return torch.amp.autocast("cuda", dtype=torch.bfloat16 if dt=="bfloat16" else torch.float16)


@torch.no_grad()
def generate(model,tok,prompt,dev,dt,max_new=220,temperature=0.05,top_p=0.80):
    model.eval()
    ids=torch.tensor([tok.encode(prompt)],dtype=torch.long,device=dev)
    for _ in range(max_new):
        x=ids[:,-model.cfg.block_size:]
        with ac(dev,dt):
            logits=model(x)
        logits=logits[:,-1,:].float()
        if temperature<=0:
            nxt=torch.argmax(logits,dim=-1,keepdim=True)
        else:
            probs=torch.softmax(logits/max(1e-6,temperature),dim=-1)
            sp,si=torch.sort(probs,descending=True)
            cum=torch.cumsum(sp,dim=-1)
            mask=cum>top_p
            mask[...,1:]=mask[...,:-1].clone()
            mask[...,0]=False
            sp[mask]=0
            sp=sp/sp.sum(dim=-1,keepdim=True)
            nxt=si.gather(-1,torch.multinomial(sp,1))
        ids=torch.cat([ids,nxt],dim=1)
        if int(nxt.item())==tok.eos_token_id:
            break
    txt=tok.decode(ids[0].tolist(),skip_special_tokens=True)
    return txt.split("### Response:",1)[-1].strip() if "### Response:" in txt else txt.strip()


def load_model(repo, ckpt_name, device):
    print("Scarico checkpoint:", repo, ckpt_name, flush=True)
    ckpt_path=hf_hub_download(repo_id=repo, filename=ckpt_name, repo_type="model")
    print("Checkpoint locale:", ckpt_path, flush=True)

    ck=torch.load(ckpt_path,map_location="cpu")
    tok=GPT2TokenizerFast.from_pretrained("gpt2")
    tok.pad_token=tok.eos_token

    cfg=GPTConfig(**ck["config"]) if isinstance(ck,dict) and "config" in ck else GPTConfig(
        vocab_size=tok.vocab_size, block_size=1024, n_layer=24, n_head=16, n_embd=2048
    )
    cfg.dropout=0.0
    model=GPT(cfg)
    sd=ck["model"] if isinstance(ck,dict) and "model" in ck else ck
    if any(k.startswith("module.") for k in sd.keys()):
        sd={k.replace("module.","",1):v for k,v in sd.items()}
    model.load_state_dict(sd,strict=True)
    model.to(device)
    model.eval()
    return model,tok


def P(body):
    return "### Instruction:\n\n"+body.strip()+"\n\n### Response:\n"


PROMPTS=[
"""TOOL_RESULT:
{"tool":"weather.forecast","result":{"city":"Trento","condition":"neve","temperature_c":-2,"wind_kmh":18,"request_id":"abc123"}}

Scrivi una risposta naturale in italiano usando solo i dati utili. Ignora request_id.""",

"""TOOL_RESULT:
{"tool":"finance.quote","result":{"symbol":"MSFT","price":512.34,"currency":"USD","change_percent":-1.72}}

Rispondi in italiano usando solo questi dati.""",

"""TOOL_RESULT:
{"tool":"spotify.current_song","result":{"artist":"Coldplay","title":"Yellow","album":"Parachutes"}}

Trasforma il risultato in una risposta naturale.""",

"""CONTEXT:
Una password robusta deve essere lunga, unica e non riutilizzata.

QUESTION:
Come deve essere una password sicura?

Rispondi solo usando il contesto.""",

"""CONTEXT:
Il documento descrive esclusivamente la regola 3-2-1 dei backup.

QUESTION:
Chi ha fondato Microsoft?

Rispondi solo usando il contesto.""",

"""Trasforma questo JSON in una risposta naturale.

{
 "weather":{"city":"Bari","condition":"sereno","temperature_c":29},
 "mail":{"unread":9,"important":3,"latest_sender":"Giulia"},
 "calendar":{"title":"Riunione","date":"venerdì","time":"14:00"}
}""",

"""TOOL_RESULT:
{"tool":"system.status","result":{"cpu":37,"ram":58,"disk":81}}

Trasforma i dati in una frase naturale.""",

"""Dati disponibili:

- città: Palermo
- meteo: soleggiato
- temperatura: 34
- email non lette: 11
- importanti: 4

Scrivi una risposta naturale senza aggiungere informazioni.""",

"""TOOL_RESULT:
{"tool":"calendar.next_event","result":{"title":"Audit","date":"martedì","time":"09:30"}}

Rispondi in italiano.""",

"""TOOL_RESULT:
{"tool":"home.sensor","result":{"device":"porta ingresso","state":"aperta"}}

Scrivi una risposta naturale.""",

"""### System:
You can call tools when needed.
Use only the available tool names and copy arguments exactly.

Available tools:
- weather.forecast: Get current weather | required: city

### User:
Che tempo fa a Genova?

### Assistant:""",

"""### System:
You can call tools when needed.
Use only the available tool names and copy arguments exactly.

Available tools:
- finance.quote: Get stock quote | required: symbol

### User:
Quanto quota TSLA?

### Assistant:""",

"""### System:
You can call tools when needed.
Use only the available tool names and copy arguments exactly.

Available tools:
- spotify.current_song: Current song

### User:
Che canzone sta suonando?

### Assistant:""",

"""### System:
You can call tools when needed.
Use only the available tool names and copy arguments exactly.

Available tools:
- calendar.next_event: Next calendar event

### User:
Qual è il mio prossimo appuntamento?

### Assistant:""",

"""### System:
You can call tools when needed.
Use only the available tool names and copy arguments exactly.

Available tools:
- unread_mail_count: Count unread emails

### User:
Quante email non lette ho?

### Assistant:""",

"""Trasforma questo JSON mantenendo tutti i valori.

{
 "home":{"garage":"chiuso","porta":"aperta"},
 "weather":{"city":"Aosta","condition":"vento","temperature_c":6},
 "mail":{"unread":2,"important":1}
}""",

"""TOOL_RESULT:
{
 "tool":"weather.forecast",
 "result":{"city":"Ancona","condition":"pioggia","temperature_c":15,"humidity":82,"debug":"ignore"}
}

Ignora debug e usa gli altri dati.""",

"""CONTEXT:
Un backup offline è una copia non sempre collegata alla rete.

QUESTION:
Che cos'è un backup offline?

Rispondi solo usando il contesto.""",

"""TOOL_RESULT:
{"tool":"stocks","result":{"symbol":"AMD","price":165.77,"currency":"USD","change_percent":5.21}}

Scrivi una frase naturale.""",

"""Trasforma in linguaggio naturale.

{
 "weather":{"city":"Lecce","condition":"caldo","temperature_c":36},
 "calendar":{"title":"Dentista","date":"domani","time":"16:30"},
 "mail":{"unread":5}
}"""
]



import re

def simple_router(user):
 t=user.lower();
 import json
 m=None
 if "meteo" in t or "tempo" in t:
  import re;m=re.search(r"(genova|udine|roma|milano|trento)",t);city=(m.group(1).title() if m else "Genova");return {"tool":"weather.forecast","result":{"city":city,"condition":"pioggia","temperature_c":18,"wind_kmh":12}}
 if "tsla" in t or "quota" in t:return {"tool":"finance.quote","result":{"symbol":"TSLA","price":312.45,"currency":"USD","change_percent":0.8}}
 return None

def tool_prompt(r):
 return "### Instruction:\nTOOL_RESULT:\n"+json.dumps(r,ensure_ascii=False)+"\n\nScrivi una risposta naturale in italiano usando solo i dati utili.\n\n### Response:\n"

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--repo-id",default="ProjectScugnizz/scugnizz-1b")
    ap.add_argument("--ckpt",default="training-runs/sft-universal-tool-renderer-1b-v3-agentic-smart-mix/checkpoint_best.pt")
    ap.add_argument("--max-new",type=int,default=220)
    ap.add_argument("--temperature",type=float,default=0.05)
    ap.add_argument("--top-p",type=float,default=0.80)
    ap.add_argument("--prompt",default=None)
    ap.add_argument("--chat",action="store_true")
    a=ap.parse_args()

    device="cuda" if torch.cuda.is_available() else "cpu"
    dt=dtype_of(device)
    print("device",device,"dtype",dt,flush=True)

    model,tok=load_model(a.repo_id,a.ckpt,device)
    if a.prompt:
        r=simple_router(a.prompt)
        p=tool_prompt(r) if r else P(a.prompt)
        print(generate(model,tok,p,device,dt,max_new=a.max_new,temperature=0.05 if r else a.temperature,top_p=a.top_p));return
    if a.chat:
        print("Agent CLI");
        while True:
            q=input("> ")
            if q.lower() in ("exit","quit"):break
            r=simple_router(q)
            p=tool_prompt(r) if r else P(q)
            print(generate(model,tok,p,device,dt,max_new=a.max_new,temperature=0.05 if r else a.temperature,top_p=a.top_p))
        return

    print("\n"+"="*100)
    print("SCUGNIZZ V3 - 20 DOMANDE NUOVE")
    print("="*100+"\n",flush=True)

    for i,p in enumerate(PROMPTS,1):
        out=generate(model,tok,P(p),device,dt,max_new=a.max_new,temperature=a.temperature,top_p=a.top_p)
        print("\n"+"="*100)
        print(f"TEST {i:02d}")
        print("-"*100)
        print("PROMPT:\n"+p)
        print("-"*100)
        print("RISPOSTA:\n"+out)
        print("="*100,flush=True)


if __name__=="__main__":
    main()