"""Buy the first observed Pump.fun launch, hold 3 seconds, then sell.

SolTrench builds unsigned transactions. This script signs locally, arms
tracking, and only then broadcasts through the caller's own Solana RPC.
"""

import asyncio
import base64
import json
import os
import uuid
from pathlib import Path

import requests
import websockets
from solders.keypair import Keypair
from solders.transaction import VersionedTransaction


API_HOST = os.getenv("SOLTRENCH_API_HOST", "https://api.soltrench.io")
WS_HOST = os.getenv("SOLTRENCH_WS_HOST", "wss://api.soltrench.io")
API_KEY = os.environ["SOLTRENCH_API_KEY"]
RPC_URL = os.environ["SOLANA_RPC_URL"]
BUY_SOL = float(os.getenv("BUY_SOL", "0.0001"))
HOLD_SECONDS = float(os.getenv("HOLD_SECONDS", "3"))
BUY_PRIORITY_FEE_SOL = float(os.getenv("BUY_PRIORITY_FEE_SOL", "0.00001"))
SELL_PRIORITY_FEE_SOL = float(os.getenv("SELL_PRIORITY_FEE_SOL", "0.00001"))
SOL_MINT = "So11111111111111111111111111111111111111112"

if os.getenv("ENABLE_LIVE_TRADING") != "yes":
    raise SystemExit(
        "Set ENABLE_LIVE_TRADING=yes only after reviewing the code and inputs."
    )
if BUY_SOL <= 0 or HOLD_SECONDS < 0:
    raise SystemExit("BUY_SOL must be positive and HOLD_SECONDS cannot be negative.")
if not all(0 <= fee <= 0.05 for fee in (BUY_PRIORITY_FEE_SOL, SELL_PRIORITY_FEE_SOL)):
    raise SystemExit("Priority fees must be between 0 and 0.05 SOL.")


def load_keypair() -> Keypair:
    keypair_path = Path(os.environ["SOLANA_KEYPAIR"]).expanduser()
    secret = bytes(json.loads(keypair_path.read_text(encoding="utf-8")))
    return Keypair.from_bytes(secret)


KEYPAIR = load_keypair()
WALLET = str(KEYPAIR.pubkey())


def api_post(path: str, body: dict | None = None, *, request_id: str | None = None) -> dict:
    headers = {"X-API-Key": API_KEY}
    if body is not None:
        headers["Content-Type"] = "application/json"
    if request_id:
        headers["Idempotency-Key"] = request_id
    response = requests.post(
        API_HOST + path,
        headers=headers,
        json=body,
        timeout=10,
    )
    response.raise_for_status()
    return response.json()


def build_trade(action: str, mint: str) -> dict:
    body = {
        "publicKey": WALLET,
        "action": action,
        "mint": mint,
        "amount": BUY_SOL if action == "buy" else "100%",
        "denominatedInSol": action == "buy",
        "mode": "standard",
        "slippageBps": 800 if action == "buy" else 1500,
        "priorityFee": (
            BUY_PRIORITY_FEE_SOL if action == "buy" else SELL_PRIORITY_FEE_SOL
        ),
        "pool": "pump",
        "confirm": True,
        "computeUnitLimit": 400000 if action == "buy" else 300000,
        "closeAccount": False,
    }
    built = api_post(
        "/v1/trade-local",
        body,
        request_id=f"{action}-{uuid.uuid4().hex}",
    )
    if built.get("signature_state") != "unsigned":
        raise RuntimeError("builder did not return an unsigned transaction")
    if built.get("required_signers") != [WALLET]:
        raise RuntimeError("builder returned an unexpected required signer")
    return built


def sign_locally(built: dict) -> tuple[VersionedTransaction, str]:
    unsigned = VersionedTransaction.from_bytes(
        base64.b64decode(built["transaction_base64"])
    )
    signed = VersionedTransaction(unsigned.message, [KEYPAIR])
    return signed, str(signed.signatures[0])


def broadcast(signed: VersionedTransaction, expected_signature: str) -> None:
    encoded = base64.b64encode(bytes(signed)).decode("ascii")
    response = requests.post(
        RPC_URL,
        json={
            "jsonrpc": "2.0",
            "id": 1,
            "method": "sendTransaction",
            "params": [
                encoded,
                {"encoding": "base64", "skipPreflight": False, "maxRetries": 3},
            ],
        },
        timeout=15,
    )
    response.raise_for_status()
    payload = response.json()
    if payload.get("error"):
        raise RuntimeError(payload["error"])
    if payload.get("result") != expected_signature:
        raise RuntimeError("RPC returned a different transaction signature")


async def receive_for_intent(ws, intent_id: str, wanted: set[str]) -> dict:
    while True:
        event = json.loads(await ws.recv())
        if event.get("type") == "error":
            raise RuntimeError(event)
        if event.get("intentId") == intent_id and event.get("status") in wanted:
            return event


async def track_broadcast_confirm(
    ws,
    built: dict,
    signed: VersionedTransaction,
    signature: str,
    action: str,
    mint: str,
) -> None:
    await ws.send(
        json.dumps(
            {
                "method": "trackTransaction",
                "intentId": built["intentId"],
                "signature": signature,
                "wallet": WALLET,
                "mint": mint,
                "action": action,
                "source": "standard",
                "timeoutMs": 120000,
            }
        )
    )
    await receive_for_intent(ws, built["intentId"], {"tracking"})
    await asyncio.to_thread(broadcast, signed, signature)
    final = await receive_for_intent(
        ws,
        built["intentId"],
        {"confirmed", "failed", "timeout"},
    )
    if final["status"] != "confirmed":
        raise RuntimeError(f"{action} did not confirm: {final}")


async def main() -> None:
    print(
        f"LIVE: buy={BUY_SOL:g} SOL, hold={HOLD_SECONDS:g}s, "
        f"requested priority fees={BUY_PRIORITY_FEE_SOL:g}+{SELL_PRIORITY_FEE_SOL:g} SOL; "
        "base fees, rent, and trading costs are additional"
    )
    ws_token = await asyncio.to_thread(api_post, "/v1/data/token")
    async with websockets.connect(
        f"{WS_HOST}/v1/data?token={ws_token['token']}",
        ping_interval=20,
    ) as ws:
        ready = json.loads(await ws.recv())
        if ready.get("type") != "ready":
            raise RuntimeError(f"websocket did not become ready: {ready}")

        await ws.send(
            json.dumps({"method": "subscribeNewToken", "source": "standard"})
        )
        while True:
            event = json.loads(await ws.recv())
            if event.get("type") == "error":
                raise RuntimeError(event)
            if event.get("type") != "newToken":
                continue
            data = event.get("data") or {}
            mint = str(data.get("mint") or "")
            if not mint or data.get("quoteMint") != SOL_MINT:
                continue
            print(f"launch {mint}")
            break

        await ws.send(
            json.dumps({"method": "unsubscribeNewToken", "source": "standard"})
        )
        while True:
            control = json.loads(await ws.recv())
            if control.get("type") == "error":
                raise RuntimeError(control)
            if (
                control.get("type") == "unsubscribed"
                and control.get("method") == "unsubscribeNewToken"
            ):
                break

        buy = await asyncio.to_thread(build_trade, "buy", mint)
        signed_buy, buy_signature = sign_locally(buy)
        await track_broadcast_confirm(
            ws,
            buy,
            signed_buy,
            buy_signature,
            "buy",
            mint,
        )
        print(f"buy confirmed {buy_signature}; holding {HOLD_SECONDS:g}s")

        await asyncio.sleep(HOLD_SECONDS)

        sell = await asyncio.to_thread(build_trade, "sell", mint)
        signed_sell, sell_signature = sign_locally(sell)
        await track_broadcast_confirm(
            ws,
            sell,
            signed_sell,
            sell_signature,
            "sell",
            mint,
        )
        print(f"sell confirmed {sell_signature}")


if __name__ == "__main__":
    asyncio.run(main())