import streamlit as st
import json
import hashlib
from web3 import Web3
from google import genai
from google.genai import types

# --------------------------------------------------
# PAGE CONFIG
# --------------------------------------------------
st.set_page_config(
    page_title="Invoice-Flow AI",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# --------------------------------------------------
# PROFESSIONAL FINTECH STYLING
# --------------------------------------------------
st.markdown("""
<style>
:root {
    --primary: #0f2a44;
    --secondary: #1e3a8a;
    --accent: #047857;
    --muted: #475569;
    --bg: #f5f7fb;
}

.stApp {
    background-color: var(--bg);
}

.block-container {
    max-width: 1200px;
    padding-top: 2rem;
}

.card {
    background: #ffffff;
    border-radius: 14px;
    padding: 1.6rem;
    margin-bottom: 1.5rem;
    border: 1px solid #e6eaf0;
    box-shadow: 0 12px 28px rgba(15, 42, 68, 0.06);
}

h1 {
    color: var(--primary);
    font-weight: 700;
    letter-spacing: -0.02em;
}

h2, h3 {
    color: var(--primary);
    font-weight: 600;
}

.caption {
    color: var(--muted);
    font-size: 0.95rem;
}

.badge-live {
    background: #e0f2fe;
    color: #0369a1;
    padding: 0.35rem 0.8rem;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 600;
}

.badge-success {
    background: #ecfdf5;
    color: #047857;
    padding: 0.35rem 0.8rem;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 600;
}

.badge-pending {
    background: #fff7ed;
    color: #9a3412;
    padding: 0.35rem 0.8rem;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 600;
}

.stButton > button {
    background-color: var(--secondary);
    color: white;
    border-radius: 10px;
    font-weight: 600;
    border: none;
}

.stButton > button:hover {
    background-color: #1e40af;
}

.stLinkButton a {
    border-radius: 10px;
    font-weight: 600;
}
</style>
""", unsafe_allow_html=True)

# --------------------------------------------------
# HERO / BRAND HEADER
# --------------------------------------------------
st.markdown("""
<div class="card">
<h1>Invoice-Flow AI</h1>
<p class="caption">
Agentic invoice verification & on-chain trust infrastructure for MSME finance
</p>
<span class="badge-live">HACKATHON LIVE PROTOTYPE</span>
</div>
""", unsafe_allow_html=True)

# --------------------------------------------------
# SECRETS (HF SAFE)
# --------------------------------------------------
try:
    GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
    RPC_URL = st.secrets["RPC_URL"]
    PRIVATE_KEY = st.secrets["PRIVATE_KEY"]
except Exception:
    st.error("Missing secrets. Add them in Hugging Face → Settings → Variables.")
    st.stop()

# --------------------------------------------------
# CLIENT SETUP
# --------------------------------------------------
genai_client = genai.Client(api_key=GEMINI_API_KEY)
w3 = Web3(Web3.HTTPProvider(RPC_URL))
account = w3.eth.account.from_key(PRIVATE_KEY)

CONTRACT_ADDRESS = "0x3af7C33f4D0303E6eb94c2D24fB0bD82B54c6914"
contract = w3.eth.contract(
    address=Web3.to_checksum_address(CONTRACT_ADDRESS),
    abi=json.load(open("abi.json"))
)

# --------------------------------------------------
# SESSION STATE
# --------------------------------------------------
for key, default in {
    "invoice": None,
    "approved": False,
    "minted": False,
    "tx_hash": None
}.items():
    if key not in st.session_state:
        st.session_state[key] = default

# --------------------------------------------------
# FALLBACK
# --------------------------------------------------
def mock_invoice():
    return {
        "seller_name": "TechFlow Solutions Pvt Ltd",
        "seller_gstin": "29AAACF5724R1Z5",
        "invoice_amount": 75000,
        "invoice_number": "INV-2024-001"
    }

# --------------------------------------------------
# MAIN LAYOUT
# --------------------------------------------------
col1, col2 = st.columns([1, 1])

# ---------------- LEFT -----------------
with col1:
    st.markdown("<div class='card'>", unsafe_allow_html=True)
    st.subheader("Invoice Ingestion")
    st.markdown("<p class='caption'>Upload an invoice PDF for automated trust verification</p>", unsafe_allow_html=True)

    uploaded = st.file_uploader("Invoice PDF", type=["pdf"])

    if uploaded:
        pdf_bytes = uploaded.getvalue()

        if st.button("Analyze Invoice"):
            with st.spinner("Extracting structured data..."):
                try:
                    res = genai_client.models.generate_content(
                        model="gemini-1.5-flash",
                        contents=[
                            types.Part.from_bytes(pdf_bytes, "application/pdf"),
                            "Extract invoice as JSON with keys seller_name, seller_gstin, invoice_amount, invoice_number. Return ONLY JSON."
                        ]
                    )
                    clean = res.text.replace("```json", "").replace("```", "").strip()
                    st.session_state.invoice = json.loads(clean)
                except Exception:
                    st.warning("AI extraction failed. Using demo-safe fallback.")
                    st.session_state.invoice = mock_invoice()
    st.markdown("</div>", unsafe_allow_html=True)

# ---------------- RIGHT ----------------
with col2:
    if st.session_state.invoice:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.subheader("Verification Summary")
        st.json(st.session_state.invoice)
        st.markdown("</div>", unsafe_allow_html=True)

        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.subheader("Buyer Consent")

        if not st.session_state.approved:
            st.markdown("<span class='badge-pending'>Awaiting Buyer Confirmation</span>", unsafe_allow_html=True)
            if st.button("Approve Invoice"):
                st.session_state.approved = True
                st.rerun()
        else:
            st.markdown("<span class='badge-success'>Buyer Approved</span>", unsafe_allow_html=True)

        st.markdown("</div>", unsafe_allow_html=True)

        if st.session_state.approved:
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            st.subheader("On-Chain Settlement")

            if not st.session_state.minted:
                if st.button("Mint Verified Receipt"):
                    with st.spinner("Finalizing on blockchain..."):
                        inv = st.session_state.invoice
                        payload = f"{inv['invoice_number']}{inv['seller_gstin']}{inv['invoice_amount']}"
                        receipt_hash = hashlib.sha256(payload.encode()).hexdigest()

                        tx = contract.functions.mintReceipt(
                            Web3.to_bytes(hexstr="0x" + receipt_hash),
                            inv["seller_gstin"],
                            int(inv["invoice_amount"])
                        ).build_transaction({
                            "from": account.address,
                            "nonce": w3.eth.get_transaction_count(account.address),
                            "gas": 250000,
                            "gasPrice": w3.eth.gas_price,
                            "chainId": 80002
                        })

                        signed = w3.eth.account.sign_transaction(tx, PRIVATE_KEY)
                        tx_hash = w3.eth.send_raw_transaction(signed.raw_transaction)
                        st.session_state.tx_hash = w3.to_hex(tx_hash)
                        st.session_state.minted = True
                        st.rerun()

            if st.session_state.minted:
                st.success("Receipt minted successfully")
                st.link_button(
                    "View Transaction on PolygonScan",
                    f"https://amoy.polygonscan.com/tx/{st.session_state.tx_hash}"
                )

            st.markdown("</div>", unsafe_allow_html=True)

# --------------------------------------------------
# FOOTER
# --------------------------------------------------
st.markdown("---")
st.caption(
    "Invoice-Flow AI • Trust Infrastructure for MSME Finance • "
    "Gemini 1.5 Flash • Polygon Amoy"
)