from __future__ import annotations
import json,os,secrets,time
from datetime import datetime,timezone
from pathlib import Path
from typing import Any,Callable,TypeVar
T=TypeVar('T')
class TraceFlow:
 def __init__(self,enabled=False,log_dir=None,trace_id=None):
  self.enabled=enabled;self.started=time.perf_counter();self.trace_id=trace_id or os.getenv('TRACEFLOW_TRACE_ID') or 'tf_'+secrets.token_hex(8);os.environ['TRACEFLOW_TRACE_ID']=self.trace_id;self.dir=Path(log_dir or Path(__file__).resolve().parents[1]/'storage/logs/traceflow');self.dir.mkdir(parents=True,exist_ok=True);self.log=self.dir/(datetime.now().strftime('%Y-%m-%d')+'.jsonl');self.event('request.start')
 @classmethod
 def start(cls,enabled=False,log_dir=None,trace_id=None):return cls(enabled,log_dir,trace_id)
 def event(self,name,context=None,level='info'):self._write('event',name,context or {},level)
 def span(self,name,fn,*a,context=None,**kw):
  if not self.enabled:return fn(*a,**kw)
  s=time.perf_counter();self._write('span.start',name,context or {})
  try:r=fn(*a,**kw);self._write('span.end',name,context or {},duration=(time.perf_counter()-s)*1000);return r
  except Exception as e:self._write('span.error',name,context or {},'error',(time.perf_counter()-s)*1000,e);raise
 def file_read(self,path,**context):self.event('file.read',{'file':str(path),**context})
 def file_write(self,path,**context):self.event('file.write',{'file':str(path),**context})
 def finish(self,exit_code=0):self._write('span.end','request',{'exit_code':exit_code},duration=(time.perf_counter()-self.started)*1000)
 def _write(self,t,n,c,l='info',duration=None,error=None):
  if not self.enabled:return
  r={'timestamp':datetime.now(timezone.utc).astimezone().isoformat(timespec='seconds'),'trace_id':self.trace_id,'event_id':'ev_'+secrets.token_hex(8),'level':l,'type':t,'name':n,'runtime':'python','context':self._redact(c)}
  if duration is not None:r['duration_ms']=round(duration,2)
  if error:r['error']={'class':type(error).__name__,'message':str(error)}
  with self.log.open('a') as f:f.write(json.dumps(r,separators=(',',':'))+'\n')
 @classmethod
 def _redact(cls,v):
  keys={'password','passwd','token','secret','authorization','cookie','session','api_key','apikey'}
  if isinstance(v,dict):return {k:'[redacted]' if str(k).lower() in keys else cls._redact(x) for k,x in v.items()}
  if isinstance(v,list):return [cls._redact(x) for x in v]
  return v