from __future__ import annotations

from dataclasses import dataclass, asdict
from datetime import datetime, timezone, timedelta
from typing import Any
import os, uuid
from .. import db

TIER_RANK={"Free":0,"Plus":1,"Pro":2,"Ultra":3}
DEFAULT_LIMITS={
    "Free":{"daily_total":500,"daily_plus":None,"daily_pro":None,"daily_ultra":None},
    "Plus":{"daily_total":2000,"daily_plus":None,"daily_pro":None,"daily_ultra":None},
    "Pro":{"daily_total":4000,"daily_plus":None,"daily_pro":None,"daily_ultra":None},
    "Ultra":{"daily_total":5000,"daily_plus":None,"daily_pro":None,"daily_ultra":None},
}

@dataclass(frozen=True)
class ModelSpec:
    key:str; label:str; minimum_tier:str; quota_class:str; provider:str; provider_model:str; style_note:str

MODELS={
    "qwen3-free":ModelSpec("qwen3-free","Qwen3.5 4B","Free","free","cloudflare","@cf/qwen/qwen3-30b-a3b-fp8","Local CPU Free model · compact everyday RP"),
    "deepseek-v4-flash":ModelSpec("deepseek-v4-flash","DeepSeek V4 Flash","Plus","plus","openrouter","deepseek/deepseek-v4-flash","Smooth, descriptive, softer RP"),
    "deepseek-v3.1":ModelSpec("deepseek-v3.1","DeepSeek V3.1","Plus","plus","openrouter","deepseek/deepseek-chat-v3.1","Intense, funny, dynamic RP"),
    "qwen36-pro":ModelSpec("qwen36-pro","Qwen3.6 35B","Pro","pro","modal-pro",os.getenv("MELUMIYACHAT_MODAL_PRO_MODEL","qwen3.6-35b-a3b-fp8"),"Higher-quality persistent RP"),
    "qwen35-122b-ultra":ModelSpec("qwen35-122b-ultra","Qwen3.5 122B","Ultra","ultra","modal-ultra",os.getenv("MELUMIYACHAT_MODAL_ULTRA_MODEL","qwen3.5-122b-a10b-fp8"),"Largest launch RP model"),
}
DEFAULT_MODEL={"Free":"qwen3-free","Plus":"deepseek-v4-flash","Pro":"qwen36-pro","Ultra":"qwen35-122b-ultra"}

def tier(v:str|None)->str:
    x=str(v or 'Free').title(); return x if x in TIER_RANK else 'Free'
def today()->str:return datetime.now(timezone.utc).date().isoformat()
def tomorrow()->str:return (datetime.now(timezone.utc).date()+timedelta(days=1)).isoformat()

def ensure_limits()->None:
    for t,v in DEFAULT_LIMITS.items():
        db.execute('''INSERT OR IGNORE INTO generation_limit_settings(plan_tier,daily_total,daily_plus,daily_pro,daily_ultra) VALUES (?,?,?,?,?)''',(t,v['daily_total'],v['daily_plus'],v['daily_pro'],v['daily_ultra']))

def promote_due()->None:
    ensure_limits()
    rows=db.rows("SELECT * FROM generation_limit_settings WHERE pending_effective_date IS NOT NULL AND pending_effective_date<=?",(today(),))
    for r in rows:
        db.execute('''UPDATE generation_limit_settings SET daily_total=COALESCE(pending_daily_total,daily_total),daily_plus=pending_daily_plus,daily_pro=pending_daily_pro,daily_ultra=pending_daily_ultra,pending_daily_total=NULL,pending_daily_plus=NULL,pending_daily_pro=NULL,pending_daily_ultra=NULL,pending_effective_date=NULL,updated_at=CURRENT_TIMESTAMP WHERE plan_tier=?''',(r['plan_tier'],))

def limits_for(plan:str)->dict[str,Any]:
    promote_due(); p=tier(plan); r=db.row("SELECT * FROM generation_limit_settings WHERE plan_tier=?",(p,)) or {}
    return {k:r.get(k) for k in ('daily_total','daily_plus','daily_pro','daily_ultra','emergency_paused','pending_daily_total','pending_daily_plus','pending_daily_pro','pending_daily_ultra','pending_effective_date')}

def public_models(plan:str)->list[dict[str,Any]]:
    p=tier(plan); out=[]
    for m in MODELS.values():
        if TIER_RANK[p]>=TIER_RANK[m.minimum_tier]:
            out.append({'key':m.key,'label':m.label,'minimum_tier':m.minimum_tier,'quota_class':m.quota_class,'style_note':m.style_note})
    return out

def resolve(plan:str,key:str|None)->ModelSpec:
    p=tier(plan); m=MODELS.get(str(key or DEFAULT_MODEL[p]));
    if not m: raise PermissionError('That model is not available.')
    if TIER_RANK[p]<TIER_RANK[m.minimum_tier]: raise PermissionError(f'{m.label} requires {m.minimum_tier}.')
    return m

def usage(user_id:int,plan:str)->dict[str,Any]:
    p=tier(plan); lim=limits_for(p)
    db.execute('''INSERT OR IGNORE INTO generation_usage_daily(user_id,usage_date) VALUES (?,?)''',(user_id,today()))
    r=db.row("SELECT * FROM generation_usage_daily WHERE user_id=? AND usage_date=?",(user_id,today())) or {}
    def rem(limit,used): return None if limit is None else max(0,int(limit)-int(used or 0))
    return {'date_utc':today(),'reset':'00:00 UTC','plan_tier':p,'total_used':int(r.get('total_generations') or 0),'total_limit':lim['daily_total'],'total_remaining':rem(lim['daily_total'],r.get('total_generations')),'plus_used':int(r.get('plus_generations') or 0),'plus_limit':lim['daily_plus'],'plus_remaining':rem(lim['daily_plus'],r.get('plus_generations')),'pro_used':int(r.get('pro_generations') or 0),'pro_limit':lim['daily_pro'],'pro_remaining':rem(lim['daily_pro'],r.get('pro_generations')),'ultra_used':int(r.get('ultra_generations') or 0),'ultra_limit':lim['daily_ultra'],'ultra_remaining':rem(lim['daily_ultra'],r.get('ultra_generations'))}

def reserve(user_id:int,plan:str,key:str|None)->dict[str,Any]:
    p=tier(plan); m=resolve(p,key); lim=limits_for(p)
    if lim.get('emergency_paused'): raise PermissionError(f'{p} generation is temporarily paused by the Owner.')
    u=usage(user_id,p)
    if u['total_remaining']<=0: raise PermissionError(f'Daily AI reply limit reached for {p}. Resets at 00:00 UTC.')
    rid=str(uuid.uuid4()); col=m.quota_class+'_generations'
    db.execute(f'''UPDATE generation_usage_daily SET total_generations=total_generations+1,{col}={col}+1,updated_at=CURRENT_TIMESTAMP WHERE user_id=? AND usage_date=?''',(user_id,today()))
    db.execute('''INSERT INTO generation_reservations(request_id,user_id,usage_date,model_key,quota_class,status) VALUES (?,?,?,?,?,'reserved')''',(rid,user_id,today(),m.key,m.quota_class))
    return {'request_id':rid,'model':m,'usage_after':usage(user_id,p)}

def release(request_id:str)->None:
    r=db.row("SELECT * FROM generation_reservations WHERE request_id=?",(request_id,));
    if not r or r.get('status')!='reserved': return
    col=r['quota_class']+'_generations'
    db.execute(f'''UPDATE generation_usage_daily SET total_generations=MAX(0,total_generations-1),{col}=MAX(0,{col}-1),updated_at=CURRENT_TIMESTAMP WHERE user_id=? AND usage_date=?''',(r['user_id'],r['usage_date']))
    db.execute("UPDATE generation_reservations SET status='released',updated_at=CURRENT_TIMESTAMP WHERE request_id=?",(request_id,))

def finalize(request_id:str,prompt_tokens:int=0,completion_tokens:int=0,cost_usd:float=0.0)->None:
    r=db.row("SELECT * FROM generation_reservations WHERE request_id=?",(request_id,));
    if not r or r.get('status')!='reserved': return
    db.execute("UPDATE generation_reservations SET status='completed',prompt_tokens=?,completion_tokens=?,provider_cost_usd=?,updated_at=CURRENT_TIMESTAMP WHERE request_id=?",(max(0,int(prompt_tokens or 0)),max(0,int(completion_tokens or 0)),max(0,float(cost_usd or 0)),request_id))
    db.execute("UPDATE generation_usage_daily SET prompt_tokens=prompt_tokens+?,completion_tokens=completion_tokens+?,provider_cost_usd=provider_cost_usd+?,updated_at=CURRENT_TIMESTAMP WHERE user_id=? AND usage_date=?",(max(0,int(prompt_tokens or 0)),max(0,int(completion_tokens or 0)),max(0,float(cost_usd or 0)),r['user_id'],r['usage_date']))

def owner_summary()->dict[str,Any]:
    promote_due(); settings=db.rows("SELECT * FROM generation_limit_settings ORDER BY CASE plan_tier WHEN 'Free' THEN 1 WHEN 'Plus' THEN 2 WHEN 'Pro' THEN 3 ELSE 4 END")
    day=today(); used=db.rows("SELECT * FROM generation_usage_daily WHERE usage_date=?",(day,))
    agg={'total_generations':0,'free_generations':0,'plus_generations':0,'pro_generations':0,'ultra_generations':0,'prompt_tokens':0,'completion_tokens':0,'provider_cost_usd':0.0}
    for r in used:
        for k in ('total_generations','free_generations','plus_generations','pro_generations','ultra_generations','prompt_tokens','completion_tokens'): agg[k]+=int(r.get(k) or 0)
        agg['provider_cost_usd']+=float(r.get('provider_cost_usd') or 0)
    return {'date_utc':day,'reset':'00:00 UTC','settings':settings,'today':agg,'estimated_monthly_cost_usd':round(agg['provider_cost_usd']*30,2),'defaults':DEFAULT_LIMITS}

def owner_update(plan:str,payload:dict[str,Any])->dict[str,Any]:
    p=tier(plan); ensure_limits(); mode=str(payload.get('apply_mode') or 'next_reset')
    total=int(payload.get('daily_total') or 0)
    if total<1: raise ValueError('Daily total must be at least 1.')
    paused=1 if payload.get('emergency_paused') else 0
    if mode=='now':
        db.execute('UPDATE generation_limit_settings SET daily_total=?,daily_plus=NULL,daily_pro=NULL,daily_ultra=NULL,emergency_paused=?,pending_daily_total=NULL,pending_daily_plus=NULL,pending_daily_pro=NULL,pending_daily_ultra=NULL,pending_effective_date=NULL,updated_at=CURRENT_TIMESTAMP WHERE plan_tier=?',(total,paused,p))
    else:
        db.execute('UPDATE generation_limit_settings SET emergency_paused=?,pending_daily_total=?,pending_daily_plus=NULL,pending_daily_pro=NULL,pending_daily_ultra=NULL,pending_effective_date=?,updated_at=CURRENT_TIMESTAMP WHERE plan_tier=?',(paused,total,tomorrow(),p))
    return db.row("SELECT * FROM generation_limit_settings WHERE plan_tier=?",(p,)) or {}

def restore_defaults(apply_mode:str='next_reset')->dict[str,Any]:
    for p,v in DEFAULT_LIMITS.items(): owner_update(p,{**v,'apply_mode':apply_mode,'emergency_paused':False})
    return owner_summary()
