#!/usr/bin/env python3
"""AFIP v3 - Acoustic Flatulence Intelligence Platform"""
import os, time, json, hashlib, struct, sqlite3, tempfile, wave
from pathlib import Path
from datetime import datetime, timezone
from typing import List, Tuple
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware

DB = Path(os.getenv("DATABASE_PATH", "./afip.db"))
OUTPUT = Path("./output"); OUTPUT.mkdir(exist_ok=True)
B58 = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"

def b58(v: bytes) -> str:
    n = int.from_bytes(v, "big")
    if n == 0: return B58[0]
    s = ""
    while n: n, r = divmod(n, 58); s = B58[r] + s
    return s

# ── SMF MIDI Writer ──
class SMF:
    def __init__(self, tpq=480):
        self.tpq = tpq; self.tracks: List[List[Tuple[int, bytes]]] = []
    def add(self) -> int: self.tracks.append([]); return len(self.tracks)-1
    def _vlq(self, v: int) -> bytes:
        buf = [v & 0x7F]; v >>= 7
        while v: buf.append((v & 0x7F) | 0x80); v >>= 7
        return bytes(reversed(buf))
    def meta(self, t: int, d: bytes = b"") -> bytes: return bytes([0xFF, t, len(d)]) + d
    def tempo(self, trk: int, tus: int = 500000): self.tracks[trk].append((0, self.meta(0x51, struct.pack(">I", tus)[1:])))
    def prog(self, trk: int, ch: int, p: int): self.tracks[trk].append((0, bytes([0xC0 | (ch & 0x0F), p & 0x7F])))
    def on(self, trk: int, ch: int, n: int, v: int, dt: int = 0): self.tracks[trk].append((dt, bytes([0x90 | (ch & 0x0F), n & 0x7F, v & 0x7F])))
    def off(self, trk: int, ch: int, n: int, v: int = 0, dt: int = 0): self.tracks[trk].append((dt, bytes([0x80 | (ch & 0x0F), n & 0x7F, v & 0x7F])))
    def eot(self, trk: int, dt: int = 0): self.tracks[trk].append((dt, self.meta(0x2F)))
    def save(self, path):
        with open(path, "wb") as f:
            f.write(b"MThd"); f.write(struct.pack(">I", 6)); f.write(struct.pack(">H", 1))
            f.write(struct.pack(">H", len(self.tracks))); f.write(struct.pack(">H", self.tpq))
            for evts in self.tracks:
                data = b"".join(self._vlq(dt) + msg for dt, msg in evts)
                f.write(b"MTrk"); f.write(struct.pack(">I", len(data))); f.write(data)

# ── YIN Pitch Detector ──
class YIN:
    def __init__(self, sr=16000, frame_ms=46):
        self.sr = sr; self.fs = int(sr * frame_ms / 1000)
        self.hs = max(1, self.fs // 4); self.th = 0.15
    def _diff(self, x):
        n, mt = len(x), len(x)//2; d = np.zeros(mt)
        for tau in range(1, mt): d[tau] = np.sum((x[:n-tau] - x[tau:n]) ** 2)
        return d
    def _cmdf(self, df):
        cm = np.ones(len(df)); rs = 0.0
        for tau in range(1, len(df)): rs += df[tau]; cm[tau] = df[tau]/(rs/tau) if rs else 1.0
        return cm
    def pitch(self, frame):
        frame = (frame[:self.fs] if len(frame) >= self.fs else np.pad(frame, (0, self.fs - len(frame)))) * np.hanning(self.fs)
        df = self._diff(frame); cm = self._cmdf(df); est = None
        for tau in range(2, len(cm)):
            if cm[tau] < self.th:
                while tau+1 < len(cm) and cm[tau+1] < cm[tau]: tau += 1
                est = tau; break
        if est is None: est = int(np.argmin(cm[2:])) + 2
        if 1 <= est < len(cm)-1:
            p = 0.5 * (cm[est-1] - cm[est+1]) / (cm[est-1] - 2*cm[est] + cm[est+1])
            est += p
        return self.sr / est if est > 0 else None
    def detect(self, y):
        return [(i / self.sr, self.pitch(y[i:i + self.fs])) for i in range(0, len(y) - self.fs, self.hs)]

# ── WAV I/O ──
def wav_read(path):
    with wave.open(str(path), "rb") as w:
        ch, sw, sr, nf = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes()
        raw = np.frombuffer(w.readframes(nf), dtype=np.int16)
        if ch == 2: raw = ((raw[0::2] + raw[1::2]) / 2).astype(np.int16)
        return raw.astype(np.float32) / 32768.0, sr

def wav_write(path, y, sr):
    y = np.clip(y * 32767, -32767, 32767).astype(np.int16)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr); w.writeframes(y.tobytes())

# ── 36-Model Registry ──
REGISTRY = {
    "ml": [
        {"id": "MIT/ast-finetuned-audioset-10-10-0.4593", "task": "audio-classification", "role": "Primary Acoustic Classifier"},
        {"id": "facebook/wav2vec2-base-960h", "task": "asr", "role": "Spectral Transcription"},
        {"id": "microsoft/wavlm-base", "task": "feature-extraction", "role": "Embedding Extractor"},
        {"id": "facebook/hubert-base-ls960", "task": "feature-extraction", "role": "Hidden-Unit BERT"},
        {"id": "google/yamnet", "task": "audio-classification", "role": "Mobile Audio Tagger"},
        {"id": "espnet/owsm_ctc", "task": "asr", "role": "Open Whisper CTC"},
        {"id": "patrickvonplaten/whisper-large-v2", "task": "asr", "role": "Multilingual Whisper"},
        {"id": "openai/whisper-base", "task": "asr", "role": "Baseline Whisper"},
        {"id": "spotify/basic-pitch", "task": "audio-to-audio", "role": "Fundamental Freq Tracker"},
        {"id": "facebook/encodec_24khz", "task": "audio-to-audio", "role": "Neural Codec"},
        {"id": "speechbrain/sepformer-wsj02mix", "task": "audio-to-audio", "role": "Source Separation"},
        {"id": "m3hrdadfi/wav2vec2-base-100k-gtzan-music-genre", "task": "audio-classification", "role": "Genre Classifier"},
        {"id": "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition", "task": "audio-classification", "role": "Emotion Detector"},
        {"id": "superb/wav2vec2-base-superb-er", "task": "audio-classification", "role": "SUPERB Emotion"},
        {"id": "alefiury/wav2vec2-base-960h-gender-recognition-libri", "task": "audio-classification", "role": "Gender Profile"},
        {"id": "facebook/wav2vec2-xlsr-53", "task": "feature-extraction", "role": "XLS-R Encoder"},
        {"id": "jonatasgrosman/wav2vec2-large-xlsr-53-english", "task": "asr", "role": "English ASR"},
        {"id": "facebook/s2t-small-librispeech-asr", "task": "asr", "role": "Speech-to-Text S2T"},
        {"id": "speechbrain/emotion-recognition-wav2vec2-IEMOCAP", "task": "audio-classification", "role": "IEMOCAP Baseline"},
        {"id": "sentence-transformers/all-MiniLM-L6-v2", "task": "feature-extraction", "role": "Semantic Embedding"},
        {"id": "sentence-transformers/all-mpnet-base-v2", "task": "feature-extraction", "role": "MPNet Encoder"},
        {"id": "facebook/bart-base", "task": "feature-extraction", "role": "BART Feature"},
        {"id": "facebook/roberta-base", "task": "feature-extraction", "role": "RoBERTa Context"},
        {"id": "cardiffnlp/twitter-roberta-base-emotion", "task": "text-classification", "role": "Twitter Emotion"},
        {"id": "distilbert-base-uncased-finetuned-sst-2-english", "task": "text-classification", "role": "SST-2 Sentiment"},
        {"id": "dslim/bert-base-NER", "task": "token-classification", "role": "NER Tagger"},
        {"id": "huggingface-course/audio-transformers", "task": "audio-classification", "role": "Course Ref"},
        {"id": "sanchit-gandhi/whisper-medium-finetuned-common-voice-13", "task": "asr", "role": "CV-13 Whisper"},
        {"id": "jonatasgrosman/wavlm-large-xtreme-s", "task": "audio-classification", "role": "XTreme Emotion"},
        {"id": "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition", "task": "audio-classification", "role": "Emotion Re-Classifier"},
    ],
    "llm": [
        {"id": "mistralai/Mistral-7B-Instruct-v0.1", "task": "text-generation", "role": "Poetry Engine"},
        {"id": "meta-llama/Llama-2-7b-chat-hf", "task": "text-generation", "role": "Scientific Abstract"},
        {"id": "google/gemma-7b-it", "task": "text-generation", "role": "Naming Conventions"},
        {"id": "HuggingFaceH4/zephyr-7b-beta", "task": "text-generation", "role": "Roast & Critique"},
        {"id": "microsoft/Phi-3-mini-4k-instruct", "task": "text-generation", "role": "Shakespearean Xlator"},
        {"id": "tiiuae/falcon-7b-instruct", "task": "text-generation", "role": "Tokenomics Architect"},
    ],
}
ALL_MODELS = REGISTRY["ml"] + REGISTRY["llm"]

def stub_infer(m: dict, seed: int) -> dict:
    rng = np.random.default_rng(seed)
    t = m["task"]
    if t == "audio-classification":
        return {"label": str(rng.choice(["toot","brap","poot","squeak","rumble","whistle","plop","thunder"])), "score": round(float(rng.random()*0.4+0.5),4)}
    if t == "asr":
        return {"text": str(rng.choice(["brrrraaaaap","pfffffttt","prrrrrrrt","squeeeeeak","thunderclap"]))}
    if t == "feature-extraction":
        return {"dims": 768, "preview": [round(float(x),6) for x in rng.random(4)]}
    if t == "audio-to-audio":
        return {"output": "synthetic_reconstruction.wav", "quality": round(float(rng.random()),4)}
    if t == "text-classification":
        return {"label": str(rng.choice(["POSITIVE","NEGATIVE","NEUTRAL"])), "score": round(float(rng.random()),4)}
    if t == "token-classification":
        return {"entities": [{"word": "fart", "label": "B-FART", "score": 0.99}]}
    if t == "text-generation":
        return {"generated_text": f"[stub] {m['role']} says: beep boop"}
    return {"stub": True}

# ── DB ──
def init_db():
    with sqlite3.connect(DB, check_same_thread=False) as c:
        c.execute("PRAGMA journal_mode=WAL")
        c.execute("CREATE TABLE IF NOT EXISTS farts (id TEXT PRIMARY KEY, ts TEXT, audio_hash TEXT, fingerprint TEXT, fartscore INTEGER, midi_path TEXT, note_count INTEGER, duration REAL, report JSON, prev_hash TEXT, receipt_hash TEXT)")
        c.execute("CREATE TABLE IF NOT EXISTS analyses (id INTEGER PRIMARY KEY, fart_id TEXT, model_id TEXT, task TEXT, role TEXT, result JSON, latency_ms REAL, ts TEXT)")
        c.execute("CREATE INDEX IF NOT EXISTS idx_farts_ts ON farts(ts)")
        c.execute("CREATE INDEX IF NOT EXISTS idx_analyses_fart ON analyses(fart_id)")

def db_conn():
    c = sqlite3.connect(DB, check_same_thread=False)
    c.row_factory = sqlite3.Row
    return c

def latest_receipt() -> str:
    with db_conn() as c:
        r = c.execute("SELECT receipt_hash FROM farts ORDER BY ts DESC LIMIT 1").fetchone()
        return r["receipt_hash"] if r else ""

def insert_fart(fid, ah, fp, fscore, midi, notes, dur, report, prev, receipt):
    with db_conn() as c:
        c.execute("INSERT INTO farts VALUES (?,?,?,?,?,?,?,?,?,?,?)", (fid, datetime.now(timezone.utc).isoformat(), ah, fp, fscore, midi, notes, dur, json.dumps(report), prev, receipt))
        c.commit()

def insert_analysis(fid, m, res, lat):
    with db_conn() as c:
        c.execute("INSERT INTO analyses (fart_id, model_id, task, role, result, latency_ms, ts) VALUES (?,?,?,?,?,?,?)", (fid, m["id"], m["task"], m["role"], json.dumps(res), lat, datetime.now(timezone.utc).isoformat()))
        c.commit()

def list_farts(limit=50):
    with db_conn() as c:
        return [dict(r) for r in c.execute("SELECT * FROM farts ORDER BY ts DESC LIMIT ?", (limit,)).fetchall()]

def get_fart(fid):
    with db_conn() as c:
        r = c.execute("SELECT * FROM farts WHERE id=?", (fid,)).fetchone()
        return dict(r) if r else None

def leaderboard():
    with db_conn() as c:
        return [dict(r) for r in c.execute("SELECT fingerprint, fartscore, ts, note_count FROM farts ORDER BY fartscore DESC LIMIT 20").fetchall()]

# ── Ledger ──
def fingerprint(audio_bytes: bytes) -> tuple:
    h = hashlib.sha256(audio_bytes).digest()
    score = int(hashlib.sha256(h).hexdigest(), 16) % 101
    return "Fart" + b58(h)[:38], score

def receipt(fid: str, ah: str, fp: str, fscore: int, prev: str) -> str:
    return hashlib.sha256(f"{fid}:{ah}:{fp}:{fscore}:{prev}".encode()).hexdigest()

# ── Note Segmentation ──
def segment_notes(pitches):
    notes = []; active = False; nstart = 0.0; cur = None
    for t, p in pitches:
        if p and 40 <= p <= 2000:
            mn = max(0, min(127, int(69 + 12 * np.log2(p / 440))))
            vel = min(127, max(30, int(70 + np.random.randn() * 20)))
            if not active: active, nstart, cur = True, t, mn
            elif abs(mn - cur) > 2:
                if t - nstart >= 0.05: notes.append((nstart, t - nstart, cur, vel))
                nstart, cur = t, mn
        else:
            if active and t - nstart >= 0.05: notes.append((nstart, t - nstart, cur, vel))
            active = False
    if active and pitches and pitches[-1][0] - nstart >= 0.05:
        notes.append((nstart, pitches[-1][0] - nstart, cur, vel))
    return notes

def build_midi(notes, path):
    if not notes: return None
    w = SMF(); t = w.add(); w.tempo(t); w.prog(t, 0, 58)
    notes = sorted(notes, key=lambda x: x[0])
    tps = 480 * (120 / 60.0); last = 0
    for s, d, n, v in notes:
        on, dur = int(s * tps), max(1, int(d * tps))
        w.on(t, 0, n, v, on - last); w.off(t, 0, n, 0, dur); last = on + dur
    w.eot(t, 0); w.save(path); return path

# ── Full Pipeline ──
def run_pipeline(wav_path: Path):
    y, sr = wav_read(wav_path)
    dur = len(y) / sr
    audio_bytes = open(wav_path, "rb").read()
    ah = hashlib.sha256(audio_bytes).hexdigest()
    fp, fscore = fingerprint(audio_bytes)
    prev = latest_receipt() or ""

    pitches = YIN(sr=sr).detect(y)
    notes = segment_notes(pitches)
    midi_file = OUTPUT / f"{fp[:12]}_{int(time.time())}.mid"
    build_midi(notes, midi_file)

    seed = int(hashlib.md5(audio_bytes[:4096]).hexdigest(), 16) % (2**31)
    ml_out, llm_out = [], []
    for m in REGISTRY["ml"]:
        t0 = time.time()
        res = stub_infer(m, seed)
        lat = (time.time() - t0) * 1000
        ml_out.append({"model": m["id"], "role": m["role"], "task": m["task"], "result": res, "latency_ms": round(lat, 2)})
        insert_analysis(fp[:16], m, res, lat)

    prompts = [
        "Write a haiku about this fart.",
        "Name this fart like a startup.",
        "Write a fake Nature abstract about this acoustic emission.",
        "Roast this fart mercilessly.",
        "Translate this fart into Shakespearean English.",
        "Write Solana memecoin tokenomics for this fart.",
    ]
    for m, pr in zip(REGISTRY["llm"], prompts):
        t0 = time.time()
        res = stub_infer(m, int(hashlib.md5((audio_bytes[:4096] + pr.encode())).hexdigest(), 16) % (2**31))
        lat = (time.time() - t0) * 1000
        llm_out.append({"model": m["id"], "role": m["role"], "prompt": pr, "result": res, "latency_ms": round(lat, 2)})
        insert_analysis(fp[:16], m, res, lat)

    rec = receipt(fp[:16], ah, fp, fscore, prev)
    report = {"fingerprint": fp, "fartscore": fscore, "duration": dur, "note_count": len(notes), "models": ml_out + llm_out}
    insert_fart(fp[:16], ah, fp, fscore, str(midi_file), len(notes), dur, report, prev, rec)
    return {
        "fart_id": fp[:16], "fingerprint": fp, "fartscore": fscore,
        "duration_sec": dur, "note_count": len(notes),
        "midi_url": f"/output/{midi_file.name}" if midi_file.exists() else None,
        "receipt": rec, "prev_receipt": prev,
        "model_outputs": ml_out, "llm_outputs": llm_out,
        "notes": [{"start": s, "duration": d, "midi": n, "velocity": v} for s, d, n, v in notes],
    }

# ── FastAPI ──
app = FastAPI(title="AFIP", version="3.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
init_db()

app.mount("/output", StaticFiles(directory="output"), name="output")
if Path("static").exists():
    app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
def root():
    if Path("static/index.html").exists():
        return FileResponse("static/index.html")
    return {"name": "AFIP", "version": "3.0.0", "models": len(ALL_MODELS)}

@app.get("/health")
def health(): return {"status": "ok", "models_loaded": len(ALL_MODELS), "db": str(DB)}

@app.get("/registry")
def registry(): return {"ml_models": REGISTRY["ml"], "llm_models": REGISTRY["llm"], "total": len(ALL_MODELS)}

@app.get("/history")
def history(limit: int = 50): return list_farts(limit=limit)

@app.get("/leaderboard")
def lb(): return leaderboard()

@app.get("/fart/{fart_id}")
def get_fart_api(fart_id: str):
    r = get_fart(fart_id)
    if not r: raise HTTPException(status_code=404, detail="Fart not found")
    return r

@app.post("/analyze")
async def analyze(file: UploadFile = File(...)):
    tmp = Path(tempfile.gettempdir()) / f"afip_{int(time.time()*1000)}.wav"
    try:
        open(tmp, "wb").write(await file.read())
        return run_pipeline(tmp)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if tmp.exists(): tmp.unlink()

@app.get("/mock")
def mock():
    t = np.linspace(0, 1.5, int(16000 * 1.5))
    y = np.sin(2 * np.pi * 120 * t) * np.exp(-t * 2) + np.sin(2 * np.pi * 85 * t) * 0.5
    y += np.random.randn(len(y)) * 0.02
    tmp = OUTPUT / f"mock_{int(time.time())}.wav"
    wav_write(tmp, y, 16000)
    return run_pipeline(tmp)

@app.get("/download/{fname}")
def download(fname: str):
    p = OUTPUT / fname
    if not p.exists(): raise HTTPException(status_code=404)
    return FileResponse(p, media_type="audio/midi", filename=fname)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8080")))