app.py

"""Customs Compass - AI assistant for US sales tax, customs duties, and product compliance.

Single-file Streamlit application for Chinese SMEs (hardware, batteries, robotics,
electronics) exporting to the United States.
"""

from __future__ import annotations

import json
import os
import re
from html.parser import HTMLParser
from pathlib import Path
from typing import Optional

import pandas as pd
import requests
import streamlit as st

try:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent / ".env")
except ImportError:
    pass


# =============================================================================
# Section A — Configuration
# =============================================================================

st.set_page_config(
    page_title="Customs Compass — AI Trade Compliance",
    page_icon="🧭",
    layout="wide",
    initial_sidebar_state="expanded",
)

OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.2:3b")
OLLAMA_TIMEOUT = int(os.getenv("OLLAMA_TIMEOUT", "180"))  # 3 min — cold model load can be slow on CPU
OLLAMA_PROBE_TIMEOUT = 2

# OrbitAI — used for premium agent features (market intel, localization,
# GTM roadmap). The main Q&A flow still uses local Ollama for privacy/cost.
ORBITAI_API_KEY = os.getenv("ORBITAI_API_KEY", "")
ORBITAI_BASE_URL = os.getenv("ORBITAI_BASE_URL", "https://api.orbitai.global/v1")
ORBITAI_MODEL = os.getenv("ORBITAI_MODEL", "gpt-5.4")
ORBITAI_TIMEOUT = 60

CBP_NEWSROOM_URL = "https://www.cbp.gov/newsroom"
CBP_TRADE_URL = "https://www.cbp.gov/trade"
NEWS_FETCH_TIMEOUT = 5
NEWS_CACHE_TTL = 3600  # 1 hour

USER_AGENT = "CustomsCompass/1.0 (Educational)"

DATA_DIR = Path(__file__).parent

# US customs fees (FY2026 rates)
MPF_RATE = 0.003464   # Merchandise Processing Fee: 0.3464% ad valorem
MPF_MIN = 32.71       # USD
MPF_MAX = 634.62      # USD
HMF_RATE = 0.00125    # Harbor Maintenance Fee: 0.125% (sea freight only)

# Section 301 List-4A surcharge on most Chinese electronics
SECTION_301_RATE = 0.25  # 25% additional ad valorem

SYSTEM_PROMPT = (
    "You are Customs Compass, an AI assistant specialized in US sales tax, "
    "nexus thresholds, customs duties, and product compliance. Rules: "
    "1) Use only the provided CSV data. "
    "2) If answer not in data, say 'I don't have enough information.' "
    "3) Ask clarifying questions if product category is unclear. "
    "4) Always cite source (e.g., 'nexus_thresholds.csv'). "
    "5) Answer in [English] then [中文]. "
    "6) Provide checklist + risk flags + sources."
)

STATE_ABBREVIATIONS = {
    "AL": "Alabama", "AK": "Alaska", "AZ": "Arizona", "AR": "Arkansas",
    "CA": "California", "CO": "Colorado", "CT": "Connecticut", "DE": "Delaware",
    "FL": "Florida", "GA": "Georgia", "HI": "Hawaii", "ID": "Idaho",
    "IL": "Illinois", "IN": "Indiana", "IA": "Iowa", "KS": "Kansas",
    "KY": "Kentucky", "LA": "Louisiana", "ME": "Maine", "MD": "Maryland",
    "MA": "Massachusetts", "MI": "Michigan", "MN": "Minnesota", "MS": "Mississippi",
    "MO": "Missouri", "MT": "Montana", "NE": "Nebraska", "NV": "Nevada",
    "NH": "New Hampshire", "NJ": "New Jersey", "NM": "New Mexico", "NY": "New York",
    "NC": "North Carolina", "ND": "North Dakota", "OH": "Ohio", "OK": "Oklahoma",
    "OR": "Oregon", "PA": "Pennsylvania", "RI": "Rhode Island", "SC": "South Carolina",
    "SD": "South Dakota", "TN": "Tennessee", "TX": "Texas", "UT": "Utah",
    "VT": "Vermont", "VA": "Virginia", "WA": "Washington", "WV": "West Virginia",
    "WI": "Wisconsin", "WY": "Wyoming", "DC": "District of Columbia",
}

# Keywords mapped to HTS categories. The right side must match product_category in hts_duty_codes.csv.
PRODUCT_KEYWORDS = {
    "battery_with_charger": ["power bank", "power banks", "charger", "rechargeable battery pack"],
    "battery_only": ["battery", "batteries", "lithium", "li-ion", "lipo", "lifepo4", "cell"],
    "robotics_with_radio": ["robot", "robotics", "robotic arm", "agv", "amr"],
    "consumer_electronics": ["consumer electronics", "gadget", "electronics"],
    "medical_device": ["medical", "thermometer", "blood pressure", "pulse oximeter", "diagnostic"],
    "industrial_machinery": ["industrial machine", "cnc", "lathe", "press", "factory equipment"],
    "power_tools": ["power tool", "drill", "saw", "grinder", "impact driver"],
    "led_lighting": ["led", "lighting", "lamp", "bulb", "light fixture"],
    "drones": ["drone", "drones", "uav", "quadcopter"],
    "solar_panels": ["solar panel", "solar panels", "pv module", "photovoltaic"],
    "smart_home_devices": ["smart home", "smart plug", "smart bulb", "smart switch", "iot device"],
    "ev_charger": ["ev charger", "ev charging", "electric vehicle charger", "level 2 charger"],
    "wearables": ["smartwatch", "smart watch", "fitness tracker", "wearable"],
    "audio_equipment": ["speaker", "speakers", "headphone", "headphones", "earbuds", "audio equipment"],
}

NO_TAX_STATES = {"Oregon", "Delaware", "Montana", "New Hampshire", "Alaska"}


# =============================================================================
# Section B — Data Loading (cached)
# =============================================================================

NEXUS_COLS = ["state", "threshold_usd", "transaction_rule", "notes"]
HTS_COLS = ["product_category", "hts_code", "duty_rate", "fcc_needed", "ul_needed", "fda_needed", "notes"]
CBP_ALERTS_COLS = ["category", "title", "summary", "relevant_products", "country_focus", "severity", "action_required", "source_url"]

# Common English stopwords — filtered from RAG queries to avoid noise.
_STOPWORDS = frozenset({
    "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "has", "have",
    "he", "in", "is", "it", "its", "of", "on", "or", "that", "the", "to", "was",
    "were", "will", "with", "this", "these", "those", "we", "us", "our", "you",
    "your", "they", "them", "their", "i", "me", "my", "do", "does", "did", "if",
    "but", "not", "no", "any", "some", "all", "can", "could", "should", "would",
    "may", "might", "must", "shall", "what", "when", "where", "who", "how", "why",
    "which", "about", "into", "than", "then", "there", "here", "also", "more",
    "less", "very", "much", "many", "such", "so", "up", "down", "out", "over",
    "under", "again", "further", "once", "been", "being", "had", "having", "am",
    "now", "just", "only", "own", "same", "too", "ll", "re", "ve", "d", "s", "t",
})


@st.cache_data(show_spinner=False)
def load_nexus_thresholds() -> pd.DataFrame:
    path = DATA_DIR / "nexus_thresholds.csv"
    if not path.exists():
        return pd.DataFrame(columns=NEXUS_COLS)
    df = pd.read_csv(path).fillna("")
    missing = [c for c in NEXUS_COLS if c not in df.columns]
    if missing:
        raise ValueError(f"nexus_thresholds.csv missing columns: {missing}")
    df["threshold_usd"] = pd.to_numeric(df["threshold_usd"], errors="coerce").fillna(0)
    return df


@st.cache_data(show_spinner=False)
def load_hts_duty_codes() -> pd.DataFrame:
    path = DATA_DIR / "hts_duty_codes.csv"
    if not path.exists():
        return pd.DataFrame(columns=HTS_COLS)
    df = pd.read_csv(path).fillna("")
    missing = [c for c in HTS_COLS if c not in df.columns]
    if missing:
        raise ValueError(f"hts_duty_codes.csv missing columns: {missing}")
    return df


@st.cache_data(show_spinner=False)
def load_tax_rates() -> dict:
    path = DATA_DIR / "tax_rates_by_state.json"
    if not path.exists():
        return {}
    return json.loads(path.read_text(encoding="utf-8"))


@st.cache_data(show_spinner=False)
def load_cbp_alerts() -> pd.DataFrame:
    """Static CBP enforcement & tariff alerts (Section 301, UFLPA, AD/CVD, etc.)."""
    path = DATA_DIR / "cbp_alerts.csv"
    if not path.exists():
        return pd.DataFrame(columns=CBP_ALERTS_COLS)
    df = pd.read_csv(path).fillna("")
    missing = [c for c in CBP_ALERTS_COLS if c not in df.columns]
    if missing:
        raise ValueError(f"cbp_alerts.csv missing columns: {missing}")
    return df


# ---------- CBP RAG (chunks from cbp_chunks.jsonl) -----------------------

import math
from collections import Counter

_TOKEN_RE = re.compile(r"[a-z0-9]+")

# Page-title patterns for low-value index/listing pages we want to exclude
# from the RAG corpus. These pages are just lists of references (FR numbers,
# bulletin titles, etc.) and don't contain substantive guidance text.
_NOISE_TITLE_PATTERNS = (
    "Federal Register Notices",
    "Customs Bulletin and Decisions",
    "Notices of Action",
    "Quota Bulletins",
    "CBP Trade-Related",
)


def _is_noise_chunk(rec: dict) -> bool:
    """Heuristic: skip chunks from low-information index/listing pages.

    A chunk is considered noise if:
      - its page title matches one of the index-page patterns, OR
      - the text is dominated by Federal Register references (e.g. "85 FR 15714")
        — more than 4 such references signals a list page.
    """
    title = rec.get("title", "") or ""
    if any(pat in title for pat in _NOISE_TITLE_PATTERNS):
        return True
    text = rec.get("text", "") or ""
    fr_refs = re.findall(r"\b\d{2,3}\s+FR\s+\d{4,6}\b", text)
    if len(fr_refs) >= 4:
        return True
    return False


def _tokenize(text: str) -> list[str]:
    """Lowercase, extract alphanumeric tokens, drop stopwords & short tokens."""
    return [t for t in _TOKEN_RE.findall(text.lower()) if len(t) > 2 and t not in _STOPWORDS]


@st.cache_resource(show_spinner=False)
def load_cbp_chunks_index() -> dict:
    """Load cbp_chunks.jsonl + precompute TF (per chunk) and IDF (global).

    Returned dict has:
      - chunks: list[dict]       (raw chunk records)
      - tfs: list[Counter]       (per-chunk term frequencies)
      - idf: dict[str, float]    (term -> inverse doc frequency)
      - N: int                   (total chunks)
      - avg_len: float           (avg token count per chunk, for BM25 length norm)
      - lens: list[int]          (per-chunk token count)
    """
    path = DATA_DIR / "cbp_chunks.jsonl"
    if not path.exists():
        return {"chunks": [], "tfs": [], "idf": {}, "N": 0, "avg_len": 0.0, "lens": []}

    chunks: list[dict] = []
    tfs: list[Counter] = []
    title_tokens_list: list[set[str]] = []
    lens: list[int] = []
    doc_freq: Counter = Counter()

    skipped_noise = 0
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            if _is_noise_chunk(rec):
                skipped_noise += 1
                continue
            text = rec.get("text", "")
            tokens = _tokenize(text)
            title_tokens = set(_tokenize(rec.get("title", "")))
            tf = Counter(tokens)
            chunks.append(rec)
            tfs.append(tf)
            title_tokens_list.append(title_tokens)
            lens.append(len(tokens))
            for term in tf.keys():
                doc_freq[term] += 1

    N = len(chunks)
    avg_len = (sum(lens) / N) if N else 0.0
    idf = {
        term: math.log(1 + (N - df + 0.5) / (df + 0.5))
        for term, df in doc_freq.items()
    }
    return {
        "chunks": chunks,
        "tfs": tfs,
        "title_tokens": title_tokens_list,
        "idf": idf,
        "N": N,
        "avg_len": avg_len,
        "lens": lens,
        "skipped_noise": skipped_noise,
    }


def search_cbp_chunks(
    query: str,
    index: dict,
    products: list[str],
    states: list[str],
    top_k: int = 3,
    k1: float = 1.5,
    b: float = 0.75,
) -> list[dict]:
    """BM25 retrieval over CBP chunks, deduplicated to one chunk per parent page."""
    chunks = index["chunks"]
    if not chunks:
        return []

    # Build query terms: tokenize question + add product category words + state names
    q_tokens = _tokenize(query)
    for p in products:
        q_tokens.extend(_tokenize(p.replace("_", " ")))
    for s in states:
        q_tokens.extend(_tokenize(s))
    if not q_tokens:
        return []

    q_set = set(q_tokens)
    tfs = index["tfs"]
    idf = index["idf"]
    lens = index["lens"]
    avg_len = index["avg_len"] or 1.0

    title_tokens_list = index.get("title_tokens", [])
    title_boost = 2.5  # multiply IDF when query term appears in the page title

    scored: list[tuple[float, int]] = []
    for i, tf in enumerate(tfs):
        score = 0.0
        dl = lens[i]
        norm = 1 - b + b * (dl / avg_len)
        title_set = title_tokens_list[i] if i < len(title_tokens_list) else set()
        for term in q_set:
            f = tf.get(term, 0)
            if not f:
                continue
            term_idf = idf.get(term, 0.0)
            if term in title_set:
                term_idf *= title_boost
            score += term_idf * (f * (k1 + 1)) / (f + k1 * norm)
        if score > 0:
            scored.append((score, i))

    scored.sort(key=lambda t: -t[0])

    seen_parents: set[str] = set()
    results: list[dict] = []
    for score, i in scored:
        chunk = chunks[i]
        parent = chunk.get("parent_id") or chunk.get("chunk_id")
        if parent in seen_parents:
            continue
        seen_parents.add(parent)
        # Keep the full chunk text for inline display, plus a shorter excerpt
        # for prompt injection (LLM context budget).
        full_text = (chunk.get("text") or "").strip().replace("\n", " ")
        # Collapse repeated whitespace
        full_text = re.sub(r"\s+", " ", full_text)
        if len(full_text) > 700:
            excerpt = full_text[:700].rsplit(" ", 1)[0] + "…"
        else:
            excerpt = full_text
        results.append({
            "title": chunk.get("title", ""),
            "url": chunk.get("url", ""),
            "section": chunk.get("section", ""),
            "page_type": chunk.get("page_type", ""),
            "published_date": chunk.get("published_date", ""),
            "excerpt": excerpt,       # short — used in prompt
            "full_text": full_text,   # full — shown in UI
            "score": round(score, 2),
            "chunk_id": chunk.get("chunk_id", ""),
        })
        if len(results) >= top_k:
            break

    return results


def find_relevant_alerts(
    alerts_df: pd.DataFrame,
    products: list[str],
    question: str,
) -> pd.DataFrame:
    """Return alerts matching detected products or 'all'-scoped Critical alerts."""
    if alerts_df.empty:
        return alerts_df
    question_lower = question.lower()
    mentions_china = "china" in question_lower or "chinese" in question_lower or "中国" in question

    def row_matches(row: pd.Series) -> bool:
        rel = str(row.get("relevant_products", "")).lower()
        # Always show Critical alerts that scope to "all"
        if row.get("severity") == "Critical" and "all" in rel:
            return True
        # Match by product overlap
        if products:
            rel_set = {p.strip() for p in rel.split(",") if p.strip()}
            if rel_set & set(products):
                return True
            if "all" in rel_set:
                return True
        # If user mentions China and alert is China-focused, include high-severity ones
        if mentions_china and str(row.get("country_focus", "")).lower() == "china":
            if row.get("severity") in ("Critical", "High"):
                return True
        return False

    mask = alerts_df.apply(row_matches, axis=1)
    matched = alerts_df[mask]
    # Sort: Critical → High → Medium → Info
    severity_order = {"Critical": 0, "High": 1, "Medium": 2, "Info": 3}
    matched = matched.assign(
        _sev_rank=matched["severity"].map(lambda s: severity_order.get(s, 99))
    ).sort_values("_sev_rank").drop(columns=["_sev_rank"])
    return matched.head(5)


# =============================================================================
# Section C — CBP News Fetcher
# =============================================================================

class _CBPNewsParser(HTMLParser):
    """Extracts anchor links and text from cbp.gov HTML.

    The CBP newsroom uses a Drupal site; news article links generally live under
    paths containing 'newsroom' or 'news-release'. We collect every <a> and filter
    afterwards rather than relying on a brittle CSS selector.
    """

    def __init__(self) -> None:
        super().__init__()
        self.links: list[dict] = []
        self._current_href: Optional[str] = None
        self._current_text: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, Optional[str]]]) -> None:
        if tag == "a":
            for k, v in attrs:
                if k == "href" and v:
                    self._current_href = v
                    self._current_text = []
                    return

    def handle_endtag(self, tag: str) -> None:
        if tag == "a" and self._current_href is not None:
            text = " ".join(self._current_text).strip()
            if text and len(text) > 15:  # filter noise (e.g., "Home", "Login")
                self.links.append({"href": self._current_href, "text": text})
            self._current_href = None
            self._current_text = []

    def handle_data(self, data: str) -> None:
        if self._current_href is not None:
            self._current_text.append(data.strip())


@st.cache_data(ttl=NEWS_CACHE_TTL, show_spinner=False)
def fetch_cbp_news() -> list[dict]:
    """Fetch the latest CBP newsroom headlines. Returns [] on any failure."""
    try:
        resp = requests.get(
            CBP_NEWSROOM_URL,
            headers={"User-Agent": USER_AGENT},
            timeout=NEWS_FETCH_TIMEOUT,
        )
        if resp.status_code != 200:
            return []
        parser = _CBPNewsParser()
        parser.feed(resp.text)
    except (requests.RequestException, ValueError):
        return []

    items: list[dict] = []
    seen_titles: set[str] = set()
    for link in parser.links:
        href = link["href"]
        text = link["text"]
        # Filter to CBP article-like URLs
        if not any(token in href for token in ("/newsroom/", "/news-release", "/spotlights/")):
            continue
        if text in seen_titles:
            continue
        seen_titles.add(text)
        # Normalize URL
        if href.startswith("/"):
            href = "https://www.cbp.gov" + href
        items.append({"title": text, "url": href})
        if len(items) >= 20:
            break
    return items


def filter_news_by_query(
    news_items: list[dict],
    question: str,
    products: list[str],
    states: list[str],
) -> list[dict]:
    """Return news items relevant to the user query (max 5)."""
    if not news_items:
        return []
    haystack_terms = set()
    haystack_terms.update(s.lower() for s in states)
    for p in products:
        haystack_terms.update(p.replace("_", " ").lower().split())
    for word in re.findall(r"[a-zA-Z]{4,}", question.lower()):
        haystack_terms.add(word)

    scored: list[tuple[int, dict]] = []
    for item in news_items:
        title_lower = item["title"].lower()
        score = sum(1 for term in haystack_terms if term in title_lower)
        if score > 0:
            scored.append((score, item))
    scored.sort(key=lambda t: -t[0])
    return [item for _, item in scored[:5]]


# =============================================================================
# Section D — Retrieval Helpers
# =============================================================================

def extract_states(text: str) -> list[str]:
    found: set[str] = set()
    text_lower = text.lower()
    for full in STATE_ABBREVIATIONS.values():
        if full.lower() in text_lower:
            found.add(full)
    for abbr, full in STATE_ABBREVIATIONS.items():
        if re.search(rf"\b{abbr}\b", text):
            found.add(full)
    return sorted(found)


def extract_product_categories(text: str, hts_df: pd.DataFrame) -> list[str]:
    text_lower = text.lower()
    found: list[str] = []
    available = set(hts_df["product_category"].tolist()) if not hts_df.empty else set()
    for category, keywords in PRODUCT_KEYWORDS.items():
        if category not in available:
            continue
        if any(kw in text_lower for kw in keywords):
            if category not in found:
                found.append(category)
    return found


def extract_sales_amount(text: str) -> Optional[float]:
    """Parse '$200k', '200,000', '$1.5M', '500000' from text. Returns USD amount."""
    cleaned = text.replace(",", "")
    pattern = r"\$?\s*(\d+(?:\.\d+)?)\s*([kKmM])?"
    best: Optional[float] = None
    for match in re.finditer(pattern, cleaned):
        num_str, suffix = match.group(1), match.group(2)
        try:
            val = float(num_str)
        except ValueError:
            continue
        if suffix in ("k", "K"):
            val *= 1_000
        elif suffix in ("m", "M"):
            val *= 1_000_000
        # Heuristic: only accept values that look like sales figures
        if val < 1_000:
            continue
        if best is None or val > best:
            best = val
    return best


def build_context(
    question: str,
    nexus_df: pd.DataFrame,
    hts_df: pd.DataFrame,
    tax_rates: dict,
    news_items: list[dict],
    alerts_df: Optional[pd.DataFrame] = None,
    chunk_index: Optional[dict] = None,
) -> tuple[str, dict]:
    """Assemble the retrieval context block plus a structured payload."""
    states = extract_states(question)
    products = extract_product_categories(question, hts_df)
    sales = extract_sales_amount(question)

    relevant_news = filter_news_by_query(news_items, question, products, states)
    relevant_alerts = (
        find_relevant_alerts(alerts_df, products, question)
        if alerts_df is not None and not alerts_df.empty
        else pd.DataFrame(columns=CBP_ALERTS_COLS)
    )
    relevant_chunks = (
        search_cbp_chunks(question, chunk_index, products, states, top_k=3)
        if chunk_index and chunk_index.get("N")
        else []
    )

    lines: list[str] = ["## Retrieved Knowledge Base Data", ""]

    if states:
        lines.append("### Nexus thresholds (source: nexus_thresholds.csv)")
        nexus_subset = nexus_df[nexus_df["state"].isin(states)]
        for _, row in nexus_subset.iterrows():
            lines.append(
                f"- **{row['state']}**: threshold ${int(row['threshold_usd']):,} | "
                f"transactions: {row['transaction_rule']} | notes: {row['notes']}"
            )
        lines.append("")
        lines.append("### State sales tax rates (source: tax_rates_by_state.json)")
        for s in states:
            rate = tax_rates.get(s)
            if rate is not None:
                lines.append(f"- **{s}**: {rate}%")
        lines.append("")

    if products:
        lines.append("### HTS / compliance (source: hts_duty_codes.csv)")
        hts_subset = hts_df[hts_df["product_category"].isin(products)]
        for _, row in hts_subset.iterrows():
            lines.append(
                f"- **{row['product_category']}**: HTS {row['hts_code']} | "
                f"duty {row['duty_rate']} | FCC: {row['fcc_needed']} | "
                f"UL: {row['ul_needed']} | FDA: {row['fda_needed']} | {row['notes']}"
            )
        lines.append("")

    if sales is not None:
        lines.append(f"### Detected sales amount: ${sales:,.0f} USD")
        lines.append("")

    if not relevant_alerts.empty:
        lines.append("### CBP enforcement & tariff alerts (source: cbp_alerts.csv)")
        for _, row in relevant_alerts.iterrows():
            lines.append(
                f"- **[{row['severity']}] {row['category']} — {row['title']}**: "
                f"{row['summary']} _Action:_ {row['action_required']} "
                f"({row['source_url']})"
            )
        lines.append("")

    if relevant_chunks:
        lines.append("### CBP knowledge base excerpts (source: cbp_chunks.jsonl)")
        for c in relevant_chunks:
            lines.append(f"- **{c['title']}** ({c['url']})")
            lines.append(f"  > {c['excerpt']}")
        lines.append("")

    if relevant_news:
        lines.append("### Recent CBP news (source: cbp.gov/newsroom)")
        for item in relevant_news:
            lines.append(f"- {item['title']} — {item['url']}")
        lines.append("")

    if not states and not products and not relevant_news and relevant_alerts.empty and not relevant_chunks:
        lines.append("_No matching state, product, alert, news, or CBP excerpt found._")

    payload = {
        "states": states,
        "products": products,
        "sales_usd": sales,
        "news": relevant_news,
        "alerts": relevant_alerts.to_dict(orient="records"),
        "chunks": relevant_chunks,
    }
    return "\n".join(lines), payload


# =============================================================================
# Section E — Risk Assessment
# =============================================================================

def assess_nexus_risk(sales_usd: Optional[float], threshold_usd: float) -> str:
    if sales_usd is None or threshold_usd <= 0:
        return "Unknown"
    ratio = sales_usd / threshold_usd
    if ratio >= 1.0:
        return "High"
    if ratio >= 0.7:
        return "Medium"
    return "Low"


def assess_compliance_risk(product_row: pd.Series) -> list[str]:
    flags: list[str] = []
    if str(product_row.get("fcc_needed", "")).strip().lower() == "yes":
        flags.append("FCC certification required")
    if str(product_row.get("ul_needed", "")).strip().lower() == "yes":
        flags.append("UL certification required")
    if str(product_row.get("fda_needed", "")).strip().lower() == "yes":
        flags.append("FDA clearance required")
    return flags


# =============================================================================
# Section F — LLM Integration
# =============================================================================

@st.cache_data(ttl=60, show_spinner=False)
def is_ollama_available() -> bool:
    try:
        resp = requests.get(f"{OLLAMA_URL}/api/tags", timeout=OLLAMA_PROBE_TIMEOUT)
        return resp.status_code == 200
    except requests.RequestException:
        return False


def call_ollama(system_prompt: str, user_prompt: str) -> str:
    payload = {
        "model": OLLAMA_MODEL,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "stream": False,
        "keep_alive": "10m",  # keep model loaded in RAM between calls
        "options": {
            "temperature": 0.2,
            "num_predict": 600,   # cap response length so it returns quickly
            "num_ctx": 4096,      # context window
        },
    }
    resp = requests.post(
        f"{OLLAMA_URL}/api/chat",
        json=payload,
        timeout=OLLAMA_TIMEOUT,
    )
    resp.raise_for_status()
    data = resp.json()
    return data.get("message", {}).get("content", "").strip()


# --- OrbitAI client (OpenAI-compatible) ---------------------------------------
# Used for premium agent features that need a stronger model than llama3.2:3b
# (market intelligence, localization, GTM roadmap generation, document analysis).

def is_orbitai_configured() -> bool:
    return bool(ORBITAI_API_KEY) and ORBITAI_API_KEY.startswith("sk-")


def call_orbitai(
    system_prompt: str,
    user_prompt: str,
    model: Optional[str] = None,
    temperature: float = 0.4,
) -> str:
    """Call OrbitAI's OpenAI-compatible chat completions endpoint.

    Raises requests.RequestException on network errors. Callers should catch
    and fall back gracefully (typically to Ollama or a deterministic template).
    """
    if not is_orbitai_configured():
        raise RuntimeError("ORBITAI_API_KEY is not set; cannot call OrbitAI.")
    payload = {
        "model": model or ORBITAI_MODEL,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "temperature": temperature,
        "stream": False,
    }
    headers = {
        "Authorization": f"Bearer {ORBITAI_API_KEY}",
        "Content-Type": "application/json",
        "User-Agent": USER_AGENT,
    }
    resp = requests.post(
        f"{ORBITAI_BASE_URL.rstrip('/')}/chat/completions",
        json=payload,
        headers=headers,
        timeout=ORBITAI_TIMEOUT,
    )
    resp.raise_for_status()
    data = resp.json()
    choices = data.get("choices") or []
    if not choices:
        return ""
    return (choices[0].get("message") or {}).get("content", "").strip()


# --- Fallback engine -----------------------------------------------------------

RISK_BADGE = {
    "Low": "🟢 Low",
    "Medium": "🟡 Medium",
    "High": "🔴 High",
    "Unknown": "⚪ Unknown",
}

RISK_BADGE_CN = {
    "Low": "🟢 低",
    "Medium": "🟡 中",
    "High": "🔴 高",
    "Unknown": "⚪ 未知",
}


def call_fallback(
    question: str,
    payload: dict,
    nexus_df: pd.DataFrame,
    hts_df: pd.DataFrame,
    tax_rates: dict,
) -> str:
    """Build a deterministic bilingual response using only CSV lookups (no LLM)."""
    states = payload["states"]
    products = payload["products"]
    sales = payload["sales_usd"]

    en_lines: list[str] = ["## [English]"]
    cn_lines: list[str] = ["## [中文]"]
    sources: set[str] = set()
    overall_risk = "Low"

    if not states and not products:
        en_lines.append("I don't have enough information. Please mention a US state and/or a product category (e.g., 'lithium batteries to Texas').")
        cn_lines.append("信息不足。请提供一个美国州名和/或产品类别（例如：销往德克萨斯州的锂电池）。")
        return "\n".join(en_lines + [""] + cn_lines)

    # ---- Nexus / sales tax ----
    if states:
        en_lines.append("### Sales tax & nexus checklist")
        cn_lines.append("### 销售税与经济关联检查清单")
        for state in states:
            row = nexus_df[nexus_df["state"] == state]
            if row.empty:
                en_lines.append(f"- **{state}**: I don't have enough information.")
                cn_lines.append(f"- **{state}**：信息不足。")
                continue
            sources.add("nexus_thresholds.csv")
            threshold = float(row.iloc[0]["threshold_usd"])
            tx_rule = row.iloc[0]["transaction_rule"]
            rate = tax_rates.get(state)
            risk = assess_nexus_risk(sales, threshold)

            if state in NO_TAX_STATES or threshold == 0:
                en_lines.append(
                    f"- **{state}**: No state sales tax obligation. Risk: {RISK_BADGE['Low']}."
                )
                cn_lines.append(
                    f"- **{state}**：无州销售税义务。风险：{RISK_BADGE_CN['Low']}。"
                )
            else:
                threshold_str = f"${int(threshold):,}"
                if sales is not None:
                    if risk == "High":
                        en_lines.append(
                            f"- **{state}**: Sales (${sales:,.0f}) exceed the {threshold_str} threshold — "
                            f"you MUST register and collect sales tax (rate: {rate}%). Risk: {RISK_BADGE['High']}."
                        )
                        cn_lines.append(
                            f"- **{state}**：销售额（${sales:,.0f}）超过 {threshold_str} 门槛 — "
                            f"必须注册并征收销售税（税率：{rate}%）。风险：{RISK_BADGE_CN['High']}。"
                        )
                        overall_risk = "High"
                    elif risk == "Medium":
                        en_lines.append(
                            f"- **{state}**: Sales (${sales:,.0f}) approaching the {threshold_str} threshold. "
                            f"Monitor closely. Rate when triggered: {rate}%. Risk: {RISK_BADGE['Medium']}."
                        )
                        cn_lines.append(
                            f"- **{state}**：销售额（${sales:,.0f}）接近 {threshold_str} 门槛。"
                            f"请密切监控。触发后税率：{rate}%。风险：{RISK_BADGE_CN['Medium']}。"
                        )
                        if overall_risk != "High":
                            overall_risk = "Medium"
                    else:
                        en_lines.append(
                            f"- **{state}**: Sales (${sales:,.0f}) below the {threshold_str} threshold — "
                            f"no collection required yet. Risk: {RISK_BADGE['Low']}."
                        )
                        cn_lines.append(
                            f"- **{state}**：销售额（${sales:,.0f}）低于 {threshold_str} 门槛 — "
                            f"暂无需征收。风险：{RISK_BADGE_CN['Low']}。"
                        )
                else:
                    en_lines.append(
                        f"- **{state}**: Threshold {threshold_str}, transaction rule: {tx_rule or 'none'}, "
                        f"rate: {rate}%. Provide your annual sales to assess risk."
                    )
                    cn_lines.append(
                        f"- **{state}**：门槛 {threshold_str}，交易规则：{tx_rule or '无'}，"
                        f"税率：{rate}%。请提供年销售额以评估风险。"
                    )
                    if overall_risk == "Low":
                        overall_risk = "Unknown"

    # ---- Product compliance ----
    if products:
        en_lines.append("")
        cn_lines.append("")
        en_lines.append("### Customs duty & federal compliance checklist")
        cn_lines.append("### 关税与联邦合规检查清单")
        for category in products:
            row = hts_df[hts_df["product_category"] == category]
            if row.empty:
                continue
            sources.add("hts_duty_codes.csv")
            r = row.iloc[0]
            flags = assess_compliance_risk(r)
            flags_str = ", ".join(flags) if flags else "no special certification flagged"
            en_lines.append(
                f"- **{category}**: HTS code `{r['hts_code']}`, duty rate **{r['duty_rate']}**. "
                f"Required: {flags_str}. {r['notes']}"
            )
            cn_lines.append(
                f"- **{category}**：HTS 代码 `{r['hts_code']}`，关税税率 **{r['duty_rate']}**。"
                f"所需认证：{flags_str}。{r['notes']}"
            )
            if flags and overall_risk == "Low":
                overall_risk = "Medium"

    # ---- CBP Alerts (static enforcement / tariff data) ----
    if payload.get("alerts"):
        sources.add("cbp_alerts.csv")
        en_lines.append("")
        cn_lines.append("")
        en_lines.append("### ⚠️ CBP enforcement & tariff alerts")
        cn_lines.append("### ⚠️ CBP 执法与关税警报")
        sev_badge = {"Critical": "🔴 Critical", "High": "🟠 High", "Medium": "🟡 Medium", "Info": "🟢 Info"}
        sev_badge_cn = {"Critical": "🔴 紧急", "High": "🟠 高", "Medium": "🟡 中", "Info": "🟢 信息"}
        for alert in payload["alerts"]:
            sev = alert.get("severity", "Info")
            en_lines.append(
                f"- {sev_badge.get(sev, sev)} **{alert['category']} — {alert['title']}**: "
                f"{alert['summary']} _Action:_ {alert['action_required']} "
                f"[source]({alert['source_url']})"
            )
            cn_lines.append(
                f"- {sev_badge_cn.get(sev, sev)} **{alert['category']} — {alert['title']}**："
                f"{alert['summary']} _建议行动：_ {alert['action_required']} "
                f"[来源]({alert['source_url']})"
            )
            if sev == "Critical":
                overall_risk = "High"
            elif sev == "High" and overall_risk != "High":
                overall_risk = "Medium" if overall_risk == "Low" else overall_risk

    # ---- CBP knowledge base excerpts (RAG from JSONL) ----
    if payload.get("chunks"):
        sources.add("cbp_chunks.jsonl")
        en_lines.append("")
        cn_lines.append("")
        en_lines.append("### 📚 CBP knowledge base excerpts")
        cn_lines.append("### 📚 CBP 知识库摘录")
        for c in payload["chunks"]:
            en_lines.append(f"- **[{c['title']}]({c['url']})** (relevance {c['score']})")
            en_lines.append(f"  > {c['excerpt']}")
            cn_lines.append(f"- **[{c['title']}]({c['url']})** （相关度 {c['score']}）")
            cn_lines.append(f"  > {c['excerpt']}")

    # ---- News ----
    if payload.get("news"):
        sources.add("cbp.gov/newsroom")
        en_lines.append("")
        cn_lines.append("")
        en_lines.append("### Recent CBP news (informational)")
        cn_lines.append("### 近期 CBP 新闻（仅供参考）")
        for item in payload["news"]:
            en_lines.append(f"- [{item['title']}]({item['url']})")
            cn_lines.append(f"- [{item['title']}]({item['url']})")

    # ---- Risk summary + sources ----
    en_lines.append("")
    en_lines.append(f"**Overall risk:** {RISK_BADGE[overall_risk]}")
    en_lines.append(f"**Sources:** {', '.join(sorted(sources)) if sources else 'none'}")

    cn_lines.append("")
    cn_lines.append(f"**总体风险：** {RISK_BADGE_CN[overall_risk]}")
    cn_lines.append(f"**来源：** {', '.join(sorted(sources)) if sources else '无'}")

    return "\n".join(en_lines + [""] + cn_lines)


def get_answer(
    question: str,
    nexus_df: pd.DataFrame,
    hts_df: pd.DataFrame,
    tax_rates: dict,
    news_items: list[dict],
    alerts_df: Optional[pd.DataFrame] = None,
    chunk_index: Optional[dict] = None,
    force_fallback: bool = False,
) -> tuple[str, str, dict]:
    """Returns (answer_markdown, mode_used, payload)."""
    context, payload = build_context(
        question, nexus_df, hts_df, tax_rates, news_items, alerts_df, chunk_index
    )

    use_llm = (not force_fallback) and is_ollama_available()

    if use_llm:
        user_prompt = (
            f"User question: {question}\n\n"
            f"{context}\n\n"
            "Using ONLY the data above, produce a structured bilingual answer with:\n"
            "1. [English] section: a checklist, risk flag (Low/Medium/High), and sources.\n"
            "2. [中文] section: the same content translated to Chinese.\n"
            "If the data doesn't cover the question, say 'I don't have enough information.'"
        )
        try:
            answer = call_ollama(SYSTEM_PROMPT, user_prompt)
            if answer:
                return answer, "🤖 Ollama (llama3.2:3b)", payload
        except requests.RequestException:
            pass  # fall through to template fallback

    answer = call_fallback(question, payload, nexus_df, hts_df, tax_rates)
    return answer, "📋 Rule-based fallback", payload


# =============================================================================
# Section G — UI Rendering
# =============================================================================

EXAMPLE_QUERIES = [
    "We sell power banks to Texas, $200k annual sales. Do we need to collect sales tax?",
    "Our startup ships lithium batteries to California and New York. What certifications do we need?",
    "Medical thermometer exports to Florida with $150k revenue — what are our obligations?",
]


def render_sidebar(ollama_ok: bool, news_count: int, news_enabled: bool, alerts_count: int = 0, chunks_count: int = 0) -> dict:
    with st.sidebar:
        st.markdown("# 🧭 Customs Compass")
        st.caption("AI compliance assistant for US exports")

        st.markdown("### System status")
        if ollama_ok:
            st.success(f"🟢 Ollama online ({OLLAMA_MODEL})")
        else:
            st.warning("🟡 Ollama offline — fallback mode")
            st.caption("Start Ollama and run: `ollama pull llama3.2:3b`")

        if news_enabled:
            if news_count > 0:
                st.success(f"🟢 CBP news ({news_count} items)")
            else:
                st.warning("🟡 CBP news unavailable")
        else:
            st.info("📴 CBP news disabled")

        news_toggle = st.checkbox("Enable live CBP news", value=news_enabled)
        refresh_news = st.button("🔄 Refresh news now", use_container_width=True)
        force_fallback = st.checkbox("Force fallback mode (skip LLM)", value=False)

        st.markdown("---")
        st.markdown("### Knowledge base")
        st.markdown("- 📄 `nexus_thresholds.csv` (51 states)")
        st.markdown("- 📄 `hts_duty_codes.csv` (14 categories)")
        st.markdown("- 📄 `tax_rates_by_state.json`")
        st.markdown(f"- ⚠️ `cbp_alerts.csv` ({alerts_count} curated alerts)")
        st.markdown(f"- 📚 `cbp_chunks.jsonl` ({chunks_count} RAG chunks)")
        st.markdown(f"- 🌐 [CBP Newsroom]({CBP_NEWSROOM_URL})")
        st.markdown(f"- 🌐 [CBP Trade]({CBP_TRADE_URL})")

        st.markdown("---")
        st.markdown("### Example queries")
        example_clicked: Optional[str] = None
        for i, ex in enumerate(EXAMPLE_QUERIES):
            if st.button(f"💬 Example {i + 1}", key=f"ex_{i}", use_container_width=True):
                example_clicked = ex
            with st.expander(f"Preview {i + 1}"):
                st.caption(ex)

        st.markdown("---")
        st.caption("⚠️ Educational tool. Not legal or tax advice.")

        return {
            "news_enabled": news_toggle,
            "refresh_news": refresh_news,
            "force_fallback": force_fallback,
            "example_clicked": example_clicked,
        }


def render_main_form(prefilled_question: str = "") -> dict:
    st.markdown("## Ask Customs Compass")
    st.caption("Describe your product and/or ask a compliance question. The assistant answers in English and 中文.")

    col1, col2 = st.columns([1, 1])
    with col1:
        product_desc = st.text_area(
            "Product description (optional)",
            placeholder="e.g., Lithium-ion power bank, 20000 mAh, with USB-C charging",
            height=120,
            key="product_desc",
        )
    with col2:
        question = st.text_area(
            "Your question",
            value=prefilled_question,
            placeholder="e.g., We sell to Texas and California, $300k sales. Do we need to collect sales tax?",
            height=120,
            key="question",
        )

    submitted = st.button("🔍 Analyze (分析)", type="primary", use_container_width=True)
    return {"product_desc": product_desc, "question": question, "submitted": submitted}


def render_response(answer: str, mode: str, payload: dict, news_items: list[dict]) -> None:
    st.markdown("---")
    st.markdown(f"#### Response — {mode}")
    st.markdown(answer)

    with st.expander("📊 Retrieved data used"):
        st.json({
            "states_detected": payload["states"],
            "products_detected": payload["products"],
            "sales_usd_detected": payload["sales_usd"],
        })

    if payload.get("alerts"):
        st.markdown(f"#### ⚠️ CBP alerts triggered ({len(payload['alerts'])})")
        import html as _html
        for alert in payload["alerts"]:
            sev = alert["severity"].lower()
            st.markdown(
                f"""
<div class="cc-alert cc-alert-{sev}">
  <div class="cc-alert-head">
    <span class="cc-badge cc-badge-{sev}">{alert['severity']}</span>
    <span>{alert['category']} — {alert['title']}</span>
  </div>
  <div class="cc-alert-body">{_html.escape(alert['summary'])}</div>
  <div class="cc-alert-action">✅ <b>Action:</b> {_html.escape(alert['action_required'])}</div>
  <div style="margin-top:6px"><a href="{alert['source_url']}" target="_blank">🔗 Source on cbp.gov</a></div>
</div>
                """,
                unsafe_allow_html=True,
            )

    if payload.get("chunks"):
        st.markdown(f"#### 📚 CBP knowledge base excerpts ({len(payload['chunks'])})")
        st.caption(
            "Full text from the most relevant CBP pages — no need to click out. The link goes to the source page."
        )
        import html as _html
        for c in payload["chunks"]:
            published = c.get("published_date", "")
            meta = f"{c['section']} · {c['page_type']} · relevance {c['score']}"
            if published:
                meta += f" · {published}"
            full = c.get("full_text") or c.get("excerpt", "")
            safe = _html.escape(full)
            st.markdown(
                f"""
<div class="cc-card">
  <div class="cc-card-title">{_html.escape(c['title'])}</div>
  <div class="cc-card-meta">{meta}</div>
  <div class="cc-chunk">{safe}</div>
  <div style="margin-top:10px"><a href="{c['url']}" target="_blank">🔗 View original page on cbp.gov</a></div>
</div>
                """,
                unsafe_allow_html=True,
            )

    if payload.get("news"):
        with st.expander(f"📰 Relevant CBP news ({len(payload['news'])})"):
            for item in payload["news"]:
                st.markdown(f"- [{item['title']}]({item['url']})")

    st.caption("⚠️ This is an educational tool. Always consult a licensed tax or trade professional before making compliance decisions.")


# =============================================================================
# Section H — Premium Visual Polish (custom CSS + hero)
# =============================================================================

_CSS = """
<style>
/* ----------- Fonts & base ----------- */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&family=JetBrains+Mono:wght@400;500&display=swap');

html, body, [class*="css"], .stApp {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
}

code, pre, .stCode {
    font-family: 'JetBrains Mono', monospace !important;
}

/* ----------- Hide default Streamlit chrome ----------- */
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
header[data-testid="stHeader"] {background: transparent;}

/* ----------- Color tokens ----------- */
:root {
    --cc-primary: #6366F1;
    --cc-primary-dark: #4F46E5;
    --cc-accent: #EC4899;
    --cc-surface: #F8FAFC;
    --cc-border: #E2E8F0;
    --cc-text: #0F172A;
    --cc-muted: #64748B;
    --cc-success: #10B981;
    --cc-warning: #F59E0B;
    --cc-danger: #EF4444;
    --cc-critical: #DC2626;
}

/* ----------- Hero header ----------- */
.cc-hero {
    background: linear-gradient(135deg, #6366F1 0%, #8B5CF6 50%, #EC4899 100%);
    border-radius: 16px;
    padding: 28px 32px;
    margin-bottom: 20px;
    color: white;
    box-shadow: 0 10px 30px -10px rgba(99,102,241,0.45);
}
.cc-hero h1 {
    font-size: 2.0rem;
    font-weight: 800;
    margin: 0 0 4px 0;
    color: white;
    letter-spacing: -0.02em;
}
.cc-hero p {
    margin: 0;
    font-size: 1.0rem;
    opacity: 0.95;
    font-weight: 400;
}
.cc-hero-stats {
    display: flex;
    gap: 12px;
    margin-top: 16px;
    flex-wrap: wrap;
}
.cc-stat {
    background: rgba(255,255,255,0.18);
    backdrop-filter: blur(8px);
    -webkit-backdrop-filter: blur(8px);
    padding: 8px 14px;
    border-radius: 100px;
    font-size: 0.85rem;
    font-weight: 500;
    border: 1px solid rgba(255,255,255,0.25);
}
.cc-stat b { font-weight: 700; }

/* ----------- Tabs ----------- */
.stTabs [data-baseweb="tab-list"] {
    gap: 4px;
    border-bottom: 2px solid var(--cc-border);
}
.stTabs [data-baseweb="tab"] {
    padding: 10px 18px;
    border-radius: 8px 8px 0 0;
    font-weight: 600;
    color: var(--cc-muted);
    background: transparent;
}
.stTabs [data-baseweb="tab"][aria-selected="true"] {
    color: var(--cc-primary);
    background: linear-gradient(180deg, rgba(99,102,241,0.06), transparent);
    border-bottom: 2px solid var(--cc-primary);
}

/* ----------- Buttons ----------- */
.stButton > button[kind="primary"] {
    background: linear-gradient(135deg, #6366F1, #8B5CF6);
    border: 0;
    font-weight: 600;
    padding: 0.6rem 1.3rem;
    box-shadow: 0 4px 14px rgba(99,102,241,0.35);
    transition: transform 0.15s ease, box-shadow 0.15s ease;
}
.stButton > button[kind="primary"]:hover {
    transform: translateY(-1px);
    box-shadow: 0 8px 24px rgba(99,102,241,0.45);
}

/* ----------- Card components (use via st.markdown + html) ----------- */
.cc-card {
    background: white;
    border: 1px solid var(--cc-border);
    border-radius: 12px;
    padding: 18px 20px;
    box-shadow: 0 1px 3px rgba(0,0,0,0.04);
    margin-bottom: 12px;
}
.cc-card-title {
    font-weight: 700;
    font-size: 1rem;
    margin-bottom: 4px;
    color: var(--cc-text);
}
.cc-card-meta {
    font-size: 0.78rem;
    color: var(--cc-muted);
    margin-bottom: 10px;
    font-weight: 500;
}
.cc-card-body {
    color: #334155;
    font-size: 0.92rem;
    line-height: 1.55;
}

/* ----------- Risk badges ----------- */
.cc-badge {
    display: inline-block;
    padding: 4px 10px;
    border-radius: 100px;
    font-size: 0.78rem;
    font-weight: 600;
    letter-spacing: 0.02em;
    text-transform: uppercase;
}
.cc-badge-low      { background: #D1FAE5; color: #047857; }
.cc-badge-medium   { background: #FEF3C7; color: #92400E; }
.cc-badge-high     { background: #FEE2E2; color: #B91C1C; }
.cc-badge-critical { background: #DC2626; color: white; }
.cc-badge-info     { background: #DBEAFE; color: #1E40AF; }
.cc-badge-unknown  { background: #E2E8F0; color: #475569; }

/* ----------- Alert cards (color-coded by severity) ----------- */
.cc-alert {
    border-left: 4px solid;
    padding: 14px 18px;
    border-radius: 8px;
    margin-bottom: 12px;
    background: white;
    box-shadow: 0 1px 2px rgba(0,0,0,0.04);
}
.cc-alert-critical { border-color: var(--cc-critical); background: #FEF2F2; }
.cc-alert-high     { border-color: var(--cc-danger);   background: #FFF7ED; }
.cc-alert-medium   { border-color: var(--cc-warning);  background: #FFFBEB; }
.cc-alert-info     { border-color: var(--cc-primary);  background: #EEF2FF; }

.cc-alert-head {
    display: flex;
    align-items: center;
    gap: 8px;
    margin-bottom: 6px;
    font-weight: 700;
    color: var(--cc-text);
}
.cc-alert-body {
    font-size: 0.9rem;
    color: #334155;
    line-height: 1.5;
}
.cc-alert-action {
    margin-top: 8px;
    padding: 8px 12px;
    background: rgba(99,102,241,0.07);
    border-radius: 6px;
    font-size: 0.85rem;
    color: #1E293B;
}

/* ----------- Big-number stat card for cost calculator ----------- */
.cc-total-card {
    background: linear-gradient(135deg, #6366F1, #8B5CF6);
    color: white;
    padding: 24px;
    border-radius: 14px;
    text-align: center;
    box-shadow: 0 10px 25px -8px rgba(99,102,241,0.5);
}
.cc-total-card .label {
    font-size: 0.85rem;
    opacity: 0.9;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    font-weight: 500;
}
.cc-total-card .amount {
    font-size: 2.4rem;
    font-weight: 800;
    margin: 4px 0;
    letter-spacing: -0.02em;
}
.cc-total-card .markup {
    font-size: 0.85rem;
    opacity: 0.92;
}

/* ----------- Chunk excerpts ----------- */
.cc-chunk {
    background: #F8FAFC;
    border-left: 4px solid var(--cc-primary);
    border-radius: 6px;
    padding: 12px 16px;
    font-size: 0.9em;
    line-height: 1.55;
    white-space: pre-wrap;
    color: #1E293B;
    margin: 8px 0;
}

/* ----------- Sidebar improvements ----------- */
[data-testid="stSidebar"] {
    background: linear-gradient(180deg, #0F172A 0%, #1E293B 100%);
}
[data-testid="stSidebar"] * { color: #E2E8F0 !important; }
[data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2, [data-testid="stSidebar"] h3 {
    color: white !important;
}
[data-testid="stSidebar"] a { color: #A5B4FC !important; }
[data-testid="stSidebar"] .stButton > button {
    background: rgba(255,255,255,0.08);
    color: white !important;
    border: 1px solid rgba(255,255,255,0.12);
    font-weight: 500;
}
[data-testid="stSidebar"] .stButton > button:hover {
    background: rgba(99,102,241,0.3);
    border-color: var(--cc-primary);
}

/* ----------- Inputs ----------- */
.stTextInput input, .stTextArea textarea, .stNumberInput input, .stSelectbox > div > div {
    border-radius: 8px !important;
    border: 1px solid var(--cc-border) !important;
}

/* ----------- Misc polish ----------- */
.cc-divider {
    height: 1px;
    background: var(--cc-border);
    margin: 18px 0;
}
.cc-pill {
    display: inline-block;
    padding: 3px 10px;
    background: var(--cc-surface);
    border: 1px solid var(--cc-border);
    border-radius: 100px;
    font-size: 0.75rem;
    color: var(--cc-muted);
    font-weight: 500;
}
</style>
"""


def inject_custom_css() -> None:
    st.markdown(_CSS, unsafe_allow_html=True)


def render_hero(nexus_count: int, alerts_count: int, chunks_count: int, ollama_ok: bool) -> None:
    ai_chip = "🟢 AI online" if ollama_ok else "🟡 Fallback mode"
    st.markdown(
        f"""
<div class="cc-hero">
  <h1>🧭 Customs Compass</h1>
  <p>AI compliance copilot for Chinese exporters entering the US market —
     sales tax, customs duties, federal certifications, and CBP enforcement.</p>
  <div class="cc-hero-stats">
    <span class="cc-stat">🇺🇸 <b>{nexus_count}</b> states + DC</span>
    <span class="cc-stat">⚠️ <b>{alerts_count}</b> CBP alerts</span>
    <span class="cc-stat">📚 <b>{chunks_count}</b> RAG chunks</span>
    <span class="cc-stat">{ai_chip}</span>
  </div>
</div>
        """,
        unsafe_allow_html=True,
    )


# =============================================================================
# Section I — Landed Cost Calculator
# =============================================================================

def _parse_duty_rate(rate_str: str) -> float:
    """Parse '3.4%', 'Free', '2.6%' → 0.034, 0.0, 0.026."""
    if not rate_str:
        return 0.0
    rate_str = str(rate_str).strip().lower()
    if rate_str in ("free", "0", "0%", "n/a", "none"):
        return 0.0
    m = re.search(r"(\d+(?:\.\d+)?)", rate_str)
    if not m:
        return 0.0
    val = float(m.group(1))
    return val / 100.0 if "%" in rate_str or val > 1 else val


def compute_landed_cost(
    customs_value_usd: float,
    duty_rate_str: str,
    apply_section_301: bool,
    shipping_usd: float,
    insurance_usd: float,
    destination_state: str,
    tax_rates: dict,
    use_sea_freight: bool = True,
) -> dict:
    """Compute the full landed cost breakdown for a Chinese export to the US.

    Returns a dict with every line item plus the final total.
    """
    duty_pct = _parse_duty_rate(duty_rate_str)
    base_duty = customs_value_usd * duty_pct
    s301_duty = customs_value_usd * SECTION_301_RATE if apply_section_301 else 0.0

    mpf = max(MPF_MIN, min(MPF_MAX, customs_value_usd * MPF_RATE))
    hmf = customs_value_usd * HMF_RATE if use_sea_freight else 0.0

    customs_total = base_duty + s301_duty + mpf + hmf

    cif = customs_value_usd + shipping_usd + insurance_usd
    state_tax_pct = float(tax_rates.get(destination_state, 0) or 0) / 100.0
    sales_tax = (cif + customs_total) * state_tax_pct

    total = cif + customs_total + sales_tax
    markup = (total / customs_value_usd - 1) * 100 if customs_value_usd > 0 else 0.0

    return {
        "customs_value": customs_value_usd,
        "base_duty": base_duty,
        "base_duty_pct": duty_pct,
        "section_301_duty": s301_duty,
        "section_301_applied": apply_section_301,
        "mpf": mpf,
        "hmf": hmf,
        "shipping": shipping_usd,
        "insurance": insurance_usd,
        "cif": cif,
        "customs_total": customs_total,
        "sales_tax": sales_tax,
        "sales_tax_pct": state_tax_pct,
        "destination_state": destination_state,
        "total_landed": total,
        "markup_pct": markup,
    }


def render_cost_calculator_tab(hts_df: pd.DataFrame, tax_rates: dict, alerts_df: pd.DataFrame) -> None:
    st.markdown("### 💰 Landed Cost Calculator")
    st.caption(
        "Compute the complete US import cost: customs duties + Section 301 + MPF + HMF + shipping + state sales tax."
    )

    col1, col2 = st.columns([1, 1])

    with col1:
        st.markdown("#### Product")
        categories = hts_df["product_category"].tolist() if not hts_df.empty else []
        category = st.selectbox(
            "Product category",
            options=categories,
            index=0 if categories else None,
            help="Picks the HTS code & duty rate from hts_duty_codes.csv",
        )
        customs_value = st.number_input(
            "Customs value (FOB, USD)",
            min_value=0.0,
            value=10000.0,
            step=500.0,
            help="The declared price of goods at the port of export (before shipping).",
        )
        units = st.number_input(
            "Number of units (optional)", min_value=1, value=100, step=10,
            help="Used to display per-unit landed cost.",
        )

    with col2:
        st.markdown("#### Shipping & destination")
        origin_china = st.toggle(
            "Origin: China 🇨🇳",
            value=True,
            help="If ON, applies +25% Section 301 surcharge to electronics-class HTS codes.",
        )
        use_sea = st.toggle(
            "Sea freight (adds Harbor Maintenance Fee)",
            value=True,
            help="0.125% HMF applies to sea/water imports; air shipments are exempt.",
        )
        shipping = st.number_input("Shipping cost (USD)", min_value=0.0, value=800.0, step=50.0)
        insurance = st.number_input("Insurance (USD)", min_value=0.0, value=100.0, step=25.0)
        state_options = sorted(tax_rates.keys()) if tax_rates else ["Texas"]
        destination = st.selectbox(
            "Destination state",
            options=state_options,
            index=state_options.index("Texas") if "Texas" in state_options else 0,
        )

    # Look up duty rate for selected category
    duty_rate_str = "0%"
    notes = ""
    fcc = ul_ = fda = ""
    if category and not hts_df.empty:
        row = hts_df[hts_df["product_category"] == category]
        if not row.empty:
            r = row.iloc[0]
            duty_rate_str = r["duty_rate"]
            notes = r["notes"]
            fcc, ul_, fda = r["fcc_needed"], r["ul_needed"], r["fda_needed"]

    breakdown = compute_landed_cost(
        customs_value_usd=customs_value,
        duty_rate_str=duty_rate_str,
        apply_section_301=origin_china,
        shipping_usd=shipping,
        insurance_usd=insurance,
        destination_state=destination,
        tax_rates=tax_rates,
        use_sea_freight=use_sea,
    )

    st.markdown('<div class="cc-divider"></div>', unsafe_allow_html=True)

    # --- Big total card ---
    per_unit = breakdown["total_landed"] / units if units > 0 else 0
    st.markdown(
        f"""
<div class="cc-total-card">
  <div class="label">Estimated total landed cost</div>
  <div class="amount">${breakdown['total_landed']:,.2f}</div>
  <div class="markup">{breakdown['markup_pct']:+.1f}% over FOB · ~${per_unit:,.2f} per unit ({units:,} units)</div>
</div>
        """,
        unsafe_allow_html=True,
    )

    # --- Breakdown table ---
    st.markdown("#### Breakdown")
    rows = [
        ("Customs value (FOB)", breakdown["customs_value"], ""),
        ("Shipping", breakdown["shipping"], ""),
        ("Insurance", breakdown["insurance"], ""),
        ("→ CIF subtotal", breakdown["cif"], ""),
        (f"Base customs duty ({breakdown['base_duty_pct']*100:.2f}%)", breakdown["base_duty"], f"HTS {category}"),
    ]
    if breakdown["section_301_applied"]:
        rows.append(("Section 301 surcharge (+25%)", breakdown["section_301_duty"], "China-origin electronics"))
    rows.extend([
        ("Merchandise Processing Fee (MPF)", breakdown["mpf"], f"0.3464%, min $32.71 max $634.62"),
    ])
    if use_sea:
        rows.append(("Harbor Maintenance Fee (HMF)", breakdown["hmf"], "0.125% (sea freight only)"))
    rows.extend([
        ("→ Customs total", breakdown["customs_total"], ""),
        (f"State sales tax — {destination} ({breakdown['sales_tax_pct']*100:.2f}%)", breakdown["sales_tax"], ""),
        ("→ TOTAL LANDED COST", breakdown["total_landed"], ""),
    ])

    df_show = pd.DataFrame(
        [{"Line item": r[0], "Amount (USD)": f"${r[1]:,.2f}", "Notes": r[2]} for r in rows]
    )
    st.dataframe(df_show, use_container_width=True, hide_index=True)

    # --- Compliance flags ---
    compliance_flags = []
    if str(fcc).strip().lower() == "yes":
        compliance_flags.append("📡 FCC certification required")
    if str(ul_).strip().lower() == "yes":
        compliance_flags.append("⚡ UL listing required")
    if str(fda).strip().lower() == "yes":
        compliance_flags.append("🩺 FDA clearance required")
    if origin_china:
        compliance_flags.append("⚠️ UFLPA documentation required (supply chain affidavits)")

    if compliance_flags:
        st.markdown("#### Compliance flags")
        for f in compliance_flags:
            st.markdown(f"- {f}")

    if notes:
        st.caption(f"📝 HTS notes: {notes}")

    st.caption(
        "💡 Estimates are educational. Final duties depend on the exact 10-digit HTS code, "
        "current Federal Register tariff actions, and CBP classification rulings."
    )


# =============================================================================
# Section J — Knowledge Base browser tab
# =============================================================================

def render_knowledge_tab(alerts_df: pd.DataFrame, chunk_index: dict) -> None:
    st.markdown("### 📚 Knowledge Base Browser")
    st.caption(
        "Browse the full curated alerts list and the underlying CBP RAG corpus that powers the assistant."
    )

    sub1, sub2 = st.tabs(["⚠️ CBP Alerts", "📄 CBP Pages (RAG corpus)"])

    with sub1:
        if alerts_df.empty:
            st.info("No alerts loaded.")
        else:
            severity_filter = st.multiselect(
                "Filter by severity",
                options=["Critical", "High", "Medium", "Info"],
                default=["Critical", "High", "Medium"],
            )
            filtered = alerts_df[alerts_df["severity"].isin(severity_filter)]
            st.caption(f"Showing {len(filtered)} / {len(alerts_df)} alerts.")
            for _, row in filtered.iterrows():
                sev = row["severity"].lower()
                st.markdown(
                    f"""
<div class="cc-alert cc-alert-{sev}">
  <div class="cc-alert-head">
    <span class="cc-badge cc-badge-{sev}">{row['severity']}</span>
    <span>{row['category']} — {row['title']}</span>
  </div>
  <div class="cc-alert-body">{row['summary']}</div>
  <div class="cc-alert-action">✅ <b>Action:</b> {row['action_required']}</div>
  <div style="margin-top:6px"><a href="{row['source_url']}" target="_blank">🔗 Source on cbp.gov</a></div>
</div>
                    """,
                    unsafe_allow_html=True,
                )

    with sub2:
        chunks = chunk_index.get("chunks") or []
        st.caption(f"{len(chunks)} substantive CBP chunks indexed ({chunk_index.get('skipped_noise', 0)} noise chunks filtered out).")
        # Group by parent page
        by_parent: dict[str, list[dict]] = {}
        for c in chunks:
            pid = c.get("parent_id", c.get("chunk_id"))
            by_parent.setdefault(pid, []).append(c)
        # Sort by title for browsability
        sorted_parents = sorted(by_parent.items(), key=lambda kv: (kv[1][0].get("title") or "").lower())
        query = st.text_input("Filter pages by title", placeholder="e.g. UFLPA, Section 301, IPR")
        for pid, parent_chunks in sorted_parents:
            title = parent_chunks[0].get("title", pid)
            url = parent_chunks[0].get("url", "")
            if query and query.lower() not in title.lower():
                continue
            with st.expander(f"📄 {title} ({len(parent_chunks)} chunks)"):
                st.markdown(f"🔗 [{url}]({url})")
                for c in parent_chunks[:3]:
                    import html as _html
                    safe = _html.escape((c.get("text") or "")[:1200])
                    st.markdown(f'<div class="cc-chunk">{safe}…</div>', unsafe_allow_html=True)
                if len(parent_chunks) > 3:
                    st.caption(f"… and {len(parent_chunks)-3} more chunks (collapsed)")


# =============================================================================
# Section K — PDF Report Generator (bilingual EN + 中文)
# =============================================================================

# ReportLab imports kept local to avoid slowing down app boot when feature unused
def _ensure_pdf_fonts() -> tuple[str, str]:
    """Register Microsoft YaHei (Chinese) and Helvetica (Latin) for PDF output.

    Returns (latin_font_name, cjk_font_name). Falls back to Helvetica-only
    if no CJK font is available — Chinese text will then render as boxes.
    """
    from reportlab.pdfbase import pdfmetrics
    from reportlab.pdfbase.ttfonts import TTFont

    latin = "Helvetica"
    cjk = latin  # fallback

    candidates = [
        ("YaHei", "C:/Windows/Fonts/msyh.ttc"),
        ("YaHeiBold", "C:/Windows/Fonts/msyhbd.ttc"),
        ("SimSun", "C:/Windows/Fonts/simsun.ttc"),
    ]
    for name, path in candidates:
        try:
            if name not in pdfmetrics.getRegisteredFontNames():
                pdfmetrics.registerFont(TTFont(name, path, subfontIndex=0))
            cjk = name
            break
        except Exception:
            continue
    return latin, cjk


def _orbitai_executive_summary(report_data: dict) -> dict:
    """Generate a bilingual executive summary via OrbitAI (premium model).

    Falls back to Ollama or a deterministic template if OrbitAI fails.
    Returns {"en": str, "cn": str, "source": "orbitai" | "ollama" | "template"}.
    """
    company = report_data.get("company_name") or "Your company"
    product = report_data.get("product_name") or "the product"
    states = ", ".join(report_data.get("target_states", []))
    revenue = report_data.get("annual_revenue_usd", 0)
    landed = report_data.get("landed_cost_total", 0)
    overall_risk = report_data.get("overall_risk", "Medium")

    system = (
        "You are a senior US trade compliance consultant writing the executive summary "
        "of a customs & sales tax report for a Chinese company entering the US market. "
        "Be concise, factual, and action-oriented. Write 4-6 sentences in English, then "
        "the same content translated to Mandarin Chinese (Simplified). No bullet points, "
        "no headers, just flowing professional prose for each language."
    )
    user = f"""Write the executive summary for this client engagement:

- Company: {company}
- Product: {product}
- Target US states: {states or 'unspecified'}
- Projected US sales: ${revenue:,.0f} per year
- Estimated landed cost (per shipment): ${landed:,.2f}
- Overall compliance risk: {overall_risk}

Output format (exactly two sections, separated by '---'):
[English]
<4-6 sentence executive summary in English>
---
[中文]
<same content translated to Simplified Chinese>
"""

    # 1) Try OrbitAI
    if is_orbitai_configured():
        try:
            content = call_orbitai(system, user, temperature=0.3)
            en, cn = _split_bilingual(content)
            if en or cn:
                return {"en": en, "cn": cn, "source": "orbitai"}
        except Exception:
            pass

    # 2) Try local Ollama
    if is_ollama_available():
        try:
            content = call_ollama(system, user)
            en, cn = _split_bilingual(content)
            if en or cn:
                return {"en": en, "cn": cn, "source": "ollama"}
        except Exception:
            pass

    # 3) Deterministic template fallback
    en = (
        f"{company} is preparing to export {product} to {states or 'the United States'} "
        f"with projected annual revenue of approximately ${revenue:,.0f}. Estimated landed "
        f"cost per shipment is ${landed:,.2f}, with an overall compliance risk rated {overall_risk}. "
        "This report outlines sales-tax nexus obligations, applicable customs duties (including "
        "Section 301 China surcharges where relevant), federal certification requirements "
        "(FCC, UL, FDA where applicable), and CBP enforcement priorities. We recommend "
        "engaging a licensed US customs broker before initiating commercial shipments."
    )
    cn = (
        f"{company}计划向{states or '美国'}出口{product}，预计年销售额约 ${revenue:,.0f}。"
        f"每批货物的预计到岸成本为 ${landed:,.2f}，整体合规风险评级为{overall_risk}。"
        "本报告概述了销售税经济关联义务、适用的关税（包括相关的301条款对华附加税）、"
        "联邦认证要求（FCC、UL、FDA 视情况而定）以及 CBP 的执法重点。"
        "建议在开始商业出货前聘请持牌的美国报关行。"
    )
    return {"en": en, "cn": cn, "source": "template"}


def _split_bilingual(text: str) -> tuple[str, str]:
    """Parse '[English]...---[中文]...' or '[English]...[中文]...' into (en, cn)."""
    if not text:
        return "", ""
    text = text.strip()
    # Try '---' separator first
    if "---" in text:
        parts = text.split("---", 1)
        en = re.sub(r"^\s*\[English\]\s*", "", parts[0], flags=re.IGNORECASE).strip()
        cn = re.sub(r"^\s*\[中文\]\s*", "", parts[1]).strip()
        return en, cn
    # Try [中文] marker
    m = re.search(r"\[中文\]", text)
    if m:
        en = re.sub(r"^\s*\[English\]\s*", "", text[:m.start()], flags=re.IGNORECASE).strip()
        cn = text[m.end():].strip()
        return en, cn
    return text.strip(), ""


def generate_pdf_report(report_data: dict) -> bytes:
    """Render the report as a polished bilingual PDF and return the bytes."""
    from io import BytesIO
    from reportlab.lib.pagesizes import LETTER
    from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
    from reportlab.lib.units import inch
    from reportlab.lib import colors
    from reportlab.lib.enums import TA_LEFT, TA_CENTER, TA_JUSTIFY
    from reportlab.platypus import (
        SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle,
        PageBreak, KeepTogether,
    )
    from datetime import datetime, timezone

    latin_font, cjk_font = _ensure_pdf_fonts()

    buf = BytesIO()
    doc = SimpleDocTemplate(
        buf, pagesize=LETTER,
        leftMargin=0.7 * inch, rightMargin=0.7 * inch,
        topMargin=0.7 * inch, bottomMargin=0.7 * inch,
        title=f"Customs Compass Report — {report_data.get('company_name', 'Untitled')}",
        author="Customs Compass",
    )

    # Brand colors
    PRIMARY = colors.HexColor("#6366F1")
    PRIMARY_DARK = colors.HexColor("#4F46E5")
    MUTED = colors.HexColor("#64748B")
    SURFACE = colors.HexColor("#F8FAFC")
    SUCCESS = colors.HexColor("#10B981")
    WARNING = colors.HexColor("#F59E0B")
    DANGER = colors.HexColor("#EF4444")
    CRITICAL = colors.HexColor("#DC2626")

    sev_color = {
        "Critical": CRITICAL, "High": DANGER,
        "Medium": WARNING, "Info": PRIMARY, "Low": SUCCESS, "Unknown": MUTED,
    }

    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        "Title", parent=styles["Title"], fontName=latin_font, fontSize=26,
        leading=30, textColor=PRIMARY_DARK, spaceAfter=4, alignment=TA_LEFT,
    )
    subtitle_style = ParagraphStyle(
        "Subtitle", parent=styles["Normal"], fontName=latin_font, fontSize=11,
        leading=15, textColor=MUTED, spaceAfter=22,
    )
    h2_style = ParagraphStyle(
        "H2", parent=styles["Heading2"], fontName=latin_font, fontSize=14,
        leading=18, textColor=PRIMARY_DARK, spaceBefore=14, spaceAfter=8,
    )
    h3_style = ParagraphStyle(
        "H3", parent=styles["Heading3"], fontName=latin_font, fontSize=11,
        leading=14, textColor=colors.HexColor("#0F172A"),
        spaceBefore=8, spaceAfter=4,
    )
    body_style = ParagraphStyle(
        "Body", parent=styles["Normal"], fontName=latin_font, fontSize=10,
        leading=14, textColor=colors.HexColor("#1E293B"),
        spaceAfter=8, alignment=TA_JUSTIFY,
    )
    cn_style = ParagraphStyle(
        "CN", parent=body_style, fontName=cjk_font, fontSize=10, leading=15,
    )
    cn_bold_style = ParagraphStyle(
        "CNBold", parent=cn_style, textColor=PRIMARY_DARK,
    )
    callout_style = ParagraphStyle(
        "Callout", parent=body_style, fontSize=9.5, textColor=colors.HexColor("#334155"),
        backColor=SURFACE, borderColor=PRIMARY, borderWidth=0, leftIndent=10,
        rightIndent=10, spaceBefore=6, spaceAfter=10,
    )
    small_style = ParagraphStyle(
        "Small", parent=styles["Normal"], fontName=latin_font, fontSize=8.5,
        leading=11, textColor=MUTED,
    )

    story = []

    # ---- Cover header ----
    company = report_data.get("company_name", "Client")
    product = report_data.get("product_name", "Product")
    now = datetime.now(timezone.utc).strftime("%B %d, %Y")
    story.append(Paragraph("Customs Compass <font color='#EC4899'>Report</font>", title_style))
    story.append(Paragraph(
        f"Prepared for <b>{company}</b> · {product} · {now}", subtitle_style
    ))

    # Risk + key stats cards row
    risk = report_data.get("overall_risk", "Unknown")
    risk_col = sev_color.get(risk, MUTED)
    stats = [
        ["Overall Risk", risk, f"{report_data.get('target_states_count', 0)}",
         f"${report_data.get('landed_cost_total', 0):,.0f}",
         f"${report_data.get('annual_revenue_usd', 0):,.0f}"],
        ["Risk Level", "—", "Target states", "Landed cost / shipment", "Annual revenue"],
    ]
    headerless = [
        [
            Paragraph(f"<b>OVERALL RISK</b>", small_style),
            Paragraph(f"<b>STATES</b>", small_style),
            Paragraph(f"<b>LANDED / SHIPMENT</b>", small_style),
            Paragraph(f"<b>ANNUAL REVENUE</b>", small_style),
        ],
        [
            Paragraph(f"<font color='{risk_col.hexval()}'><b>{risk}</b></font>",
                      ParagraphStyle("v", parent=body_style, fontSize=14, leading=16)),
            Paragraph(f"<b>{report_data.get('target_states_count', 0)}</b>",
                      ParagraphStyle("v", parent=body_style, fontSize=14, leading=16)),
            Paragraph(f"<b>${report_data.get('landed_cost_total', 0):,.0f}</b>",
                      ParagraphStyle("v", parent=body_style, fontSize=14, leading=16)),
            Paragraph(f"<b>${report_data.get('annual_revenue_usd', 0):,.0f}</b>",
                      ParagraphStyle("v", parent=body_style, fontSize=14, leading=16)),
        ],
    ]
    stat_table = Table(headerless, colWidths=[1.6 * inch] * 4, hAlign="LEFT")
    stat_table.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (-1, -1), SURFACE),
        ("BOX", (0, 0), (-1, -1), 1, colors.HexColor("#E2E8F0")),
        ("INNERGRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#E2E8F0")),
        ("LEFTPADDING", (0, 0), (-1, -1), 12),
        ("RIGHTPADDING", (0, 0), (-1, -1), 12),
        ("TOPPADDING", (0, 0), (-1, -1), 10),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 10),
    ]))
    story.append(stat_table)
    story.append(Spacer(1, 14))

    # ---- Executive Summary ----
    story.append(Paragraph("Executive Summary", h2_style))
    exec_summary = report_data.get("exec_summary") or {}
    if exec_summary.get("en"):
        story.append(Paragraph(exec_summary["en"], body_style))
    if exec_summary.get("cn"):
        story.append(Paragraph("中文摘要", h3_style))
        story.append(Paragraph(exec_summary["cn"], cn_style))
    src = exec_summary.get("source")
    if src:
        story.append(Paragraph(
            f"<i>Summary generated via {src.upper()}.</i>", small_style
        ))

    # ---- Company Profile ----
    story.append(Paragraph("1. Company Profile", h2_style))
    profile_rows = [
        ["Company", company],
        ["Industry / sector", report_data.get("industry", "—")],
        ["Country of origin", report_data.get("origin_country", "China")],
        ["Manufacturing location", report_data.get("manufacturing_location", "—")],
        ["Xinjiang (XUAR) exposure", "Yes ⚠️" if report_data.get("xinjiang_exposure") else "No"],
        ["Projected annual US revenue", f"${report_data.get('annual_revenue_usd', 0):,.0f}"],
        ["Sales channel", report_data.get("sales_channel", "—")],
    ]
    story.append(_kv_table(profile_rows, latin_font, SURFACE, MUTED))

    # ---- Product Analysis ----
    story.append(Paragraph("2. Product Analysis", h2_style))
    prod_rows = [
        ["Product name", product],
        ["HTS category", report_data.get("hts_category", "—")],
        ["HTS code", report_data.get("hts_code", "—")],
        ["Base duty rate", report_data.get("duty_rate", "—")],
        ["FCC required", report_data.get("fcc_needed", "—")],
        ["UL listing required", report_data.get("ul_needed", "—")],
        ["FDA clearance required", report_data.get("fda_needed", "—")],
        ["FOB declared value / shipment", f"${report_data.get('fob_value', 0):,.2f}"],
        ["Units per shipment", f"{report_data.get('units_per_shipment', 0):,}"],
    ]
    story.append(_kv_table(prod_rows, latin_font, SURFACE, MUTED))
    notes = report_data.get("hts_notes")
    if notes:
        story.append(Paragraph(f"<i>HTS notes:</i> {notes}", small_style))

    # ---- Landed Cost Breakdown ----
    breakdown = report_data.get("cost_breakdown") or {}
    if breakdown:
        story.append(Paragraph("3. Landed Cost Breakdown", h2_style))
        cost_rows = [
            ["Line item", "Amount (USD)", "Notes"],
            ["Customs value (FOB)", f"${breakdown.get('customs_value', 0):,.2f}", ""],
            ["Shipping", f"${breakdown.get('shipping', 0):,.2f}", ""],
            ["Insurance", f"${breakdown.get('insurance', 0):,.2f}", ""],
            ["CIF subtotal", f"${breakdown.get('cif', 0):,.2f}", ""],
            [f"Base duty ({breakdown.get('base_duty_pct', 0)*100:.2f}%)",
             f"${breakdown.get('base_duty', 0):,.2f}", "HTS classification"],
        ]
        if breakdown.get("section_301_applied"):
            cost_rows.append([
                "Section 301 surcharge (+25%)",
                f"${breakdown.get('section_301_duty', 0):,.2f}",
                "China-origin electronics",
            ])
        cost_rows.extend([
            ["MPF", f"${breakdown.get('mpf', 0):,.2f}", "0.3464% capped"],
            ["HMF", f"${breakdown.get('hmf', 0):,.2f}", "Sea freight only"],
            [f"State sales tax ({breakdown.get('sales_tax_pct', 0)*100:.2f}%)",
             f"${breakdown.get('sales_tax', 0):,.2f}", breakdown.get("destination_state", "")],
            ["TOTAL LANDED COST", f"${breakdown.get('total_landed', 0):,.2f}", ""],
        ])
        cost_tbl = Table(cost_rows, colWidths=[2.3 * inch, 1.5 * inch, 2.7 * inch], hAlign="LEFT")
        cost_tbl.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), PRIMARY),
            ("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
            ("FONTNAME", (0, 0), (-1, -1), latin_font),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("BOTTOMPADDING", (0, 0), (-1, 0), 8),
            ("TOPPADDING", (0, 0), (-1, 0), 8),
            ("ROWBACKGROUNDS", (0, 1), (-1, -2),
             [colors.white, SURFACE]),
            ("BACKGROUND", (0, -1), (-1, -1), colors.HexColor("#FEF3C7")),
            ("FONTNAME", (0, -1), (-1, -1), latin_font),
            ("GRID", (0, 0), (-1, -1), 0.4, colors.HexColor("#E2E8F0")),
            ("LEFTPADDING", (0, 0), (-1, -1), 8),
            ("RIGHTPADDING", (0, 0), (-1, -1), 8),
            ("ALIGN", (1, 0), (1, -1), "RIGHT"),
        ]))
        story.append(cost_tbl)

    # ---- State-by-State Nexus Analysis ----
    nexus_rows = report_data.get("nexus_analysis") or []
    if nexus_rows:
        story.append(Paragraph("4. State-by-State Nexus Analysis", h2_style))
        tbl_data = [["State", "Threshold (USD)", "Your sales", "Status", "Tax rate"]]
        for r in nexus_rows:
            tbl_data.append([
                r["state"],
                f"${r['threshold']:,.0f}" if r["threshold"] else "No tax",
                f"${r['projected_sales']:,.0f}",
                r["status"],
                f"{r['tax_rate']:.2f}%",
            ])
        nexus_tbl = Table(tbl_data, colWidths=[1.4 * inch, 1.3 * inch, 1.3 * inch, 1.3 * inch, 1.0 * inch], hAlign="LEFT")
        nexus_tbl.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), PRIMARY),
            ("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
            ("FONTNAME", (0, 0), (-1, -1), latin_font),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, SURFACE]),
            ("GRID", (0, 0), (-1, -1), 0.4, colors.HexColor("#E2E8F0")),
            ("BOTTOMPADDING", (0, 0), (-1, 0), 8),
            ("TOPPADDING", (0, 0), (-1, 0), 8),
            ("LEFTPADDING", (0, 0), (-1, -1), 8),
            ("RIGHTPADDING", (0, 0), (-1, -1), 8),
        ]))
        story.append(nexus_tbl)

    # ---- Compliance Checklist ----
    checklist = report_data.get("checklist") or []
    if checklist:
        story.append(Paragraph("5. Compliance Checklist", h2_style))
        for item in checklist:
            done = "✅" if item.get("done") else "☐"
            label = item.get("label", "")
            story.append(Paragraph(f"{done}  {label}", body_style))

    # ---- CBP Enforcement Alerts ----
    alerts = report_data.get("alerts") or []
    if alerts:
        story.append(PageBreak())
        story.append(Paragraph("6. CBP Enforcement Priorities for Your Product", h2_style))
        for alert in alerts:
            sev = alert.get("severity", "Info")
            col = sev_color.get(sev, MUTED)
            head = (
                f"<font color='{col.hexval()}'><b>[{sev}]</b></font> "
                f"<b>{alert.get('category','')}</b> — {alert.get('title','')}"
            )
            story.append(Paragraph(head, h3_style))
            story.append(Paragraph(alert.get("summary", ""), body_style))
            action = alert.get("action_required", "")
            if action:
                story.append(Paragraph(f"<b>Required action:</b> {action}", callout_style))
            src = alert.get("source_url", "")
            if src:
                story.append(Paragraph(
                    f"Source: <link href='{src}'><font color='{PRIMARY_DARK.hexval()}'>{src}</font></link>",
                    small_style,
                ))
            story.append(Spacer(1, 6))

    # ---- Sources ----
    story.append(PageBreak())
    story.append(Paragraph("7. Sources & Methodology", h2_style))
    story.append(Paragraph(
        "This report draws on the following data sources, all sourced from public "
        "US Customs and Border Protection (CBP) and state Department of Revenue publications:",
        body_style,
    ))
    sources_list = [
        "<b>nexus_thresholds.csv</b> — Economic nexus thresholds for all 50 US states + DC",
        "<b>hts_duty_codes.csv</b> — Harmonized Tariff Schedule codes & duty rates",
        "<b>tax_rates_by_state.json</b> — State sales tax base rates",
        "<b>cbp_alerts.csv</b> — Curated CBP enforcement priorities (Section 301, UFLPA, AD/CVD, IPR)",
        "<b>cbp_chunks.jsonl</b> — 398 substantive CBP page excerpts (BM25-light retrieval)",
        "<b>cbp.gov/newsroom</b> — Live news feed (cached hourly)",
    ]
    for s in sources_list:
        story.append(Paragraph(f"• {s}", body_style))

    story.append(Paragraph(
        "<i><b>Disclaimer:</b> This report is educational. It does not constitute "
        "legal, tax, or customs advice. Final determinations depend on the exact "
        "10-digit HTS code, current Federal Register tariff actions, and CBP rulings. "
        "Always consult a licensed customs broker and qualified tax counsel before "
        "initiating commercial shipments.</i>",
        callout_style,
    ))
    story.append(Paragraph(
        f"Generated by Customs Compass · {now} · "
        f"AI engine: {exec_summary.get('source', 'template').upper()}",
        small_style,
    ))

    doc.build(story)
    return buf.getvalue()


def _kv_table(rows: list, font_name: str, bg_color, muted_color):
    """Render a 2-column key/value table for profile/product sections."""
    from reportlab.lib.units import inch
    from reportlab.platypus import Table, TableStyle
    from reportlab.lib import colors

    tbl = Table(rows, colWidths=[2.3 * inch, 4.2 * inch], hAlign="LEFT")
    tbl.setStyle(TableStyle([
        ("FONTNAME", (0, 0), (-1, -1), font_name),
        ("FONTSIZE", (0, 0), (-1, -1), 9.5),
        ("TEXTCOLOR", (0, 0), (0, -1), muted_color),
        ("BACKGROUND", (0, 0), (0, -1), bg_color),
        ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 6),
        ("TOPPADDING", (0, 0), (-1, -1), 6),
        ("LEFTPADDING", (0, 0), (-1, -1), 10),
        ("RIGHTPADDING", (0, 0), (-1, -1), 10),
        ("GRID", (0, 0), (-1, -1), 0.4, colors.HexColor("#E2E8F0")),
    ]))
    return tbl


# =============================================================================
# Section L — Report Wizard UI (multi-step form)
# =============================================================================

WIZARD_STEPS = [
    "Company",
    "Product",
    "Markets & Sales",
    "Compliance status",
    "Review & Generate",
]


def _wiz_progress(current: int):
    """Render the wizard progress indicator."""
    pills = []
    for i, label in enumerate(WIZARD_STEPS):
        if i < current:
            cls = "cc-badge-low"
            icon = "✓"
        elif i == current:
            cls = "cc-badge-info"
            icon = f"{i+1}"
        else:
            cls = "cc-badge-unknown"
            icon = f"{i+1}"
        pills.append(
            f'<span class="cc-badge {cls}" style="margin-right:8px;">'
            f'{icon} {label}</span>'
        )
    st.markdown(
        f'<div style="margin-bottom:18px">{"".join(pills)}</div>',
        unsafe_allow_html=True,
    )


def render_report_wizard_tab(
    nexus_df: pd.DataFrame,
    hts_df: pd.DataFrame,
    tax_rates: dict,
    alerts_df: pd.DataFrame,
) -> None:
    st.markdown("### 📄 Generate Customs Compass Report")
    st.caption(
        "Answer a few questions about your company and product. We'll generate a polished "
        "bilingual PDF report you can take to your customs broker, tax advisor, or board."
    )

    ss = st.session_state
    if "wiz_step" not in ss:
        ss["wiz_step"] = 0
    if "wiz_data" not in ss:
        ss["wiz_data"] = {}

    _wiz_progress(ss["wiz_step"])

    data = ss["wiz_data"]

    # ---- Step 0: Company ----
    if ss["wiz_step"] == 0:
        with st.form("wiz_step_0"):
            st.markdown("#### Step 1 / 5 — Company profile")
            data["company_name"] = st.text_input(
                "Company name *",
                value=data.get("company_name", ""),
                placeholder="e.g. Shenzhen PowerCell Technology Co., Ltd."
            )
            data["industry"] = st.selectbox(
                "Industry / sector",
                options=[
                    "Batteries & energy storage",
                    "Robotics & embodied AI",
                    "Clean energy / solar / wind",
                    "Consumer electronics",
                    "Medical devices",
                    "Industrial machinery",
                    "EV & EV charging",
                    "Smart home / IoT",
                    "Other hardware",
                ],
                index=0,
            )
            col_a, col_b = st.columns(2)
            with col_a:
                data["origin_country"] = st.text_input(
                    "Country of origin", value=data.get("origin_country", "China")
                )
            with col_b:
                data["manufacturing_location"] = st.text_input(
                    "Manufacturing city / province",
                    value=data.get("manufacturing_location", ""),
                    placeholder="e.g. Shenzhen, Guangdong",
                )
            data["xinjiang_exposure"] = st.checkbox(
                "⚠️ Any sourcing or operations in Xinjiang (XUAR)?",
                value=data.get("xinjiang_exposure", False),
                help="Critical for UFLPA compliance assessment.",
            )
            submit = st.form_submit_button("Next →", type="primary", use_container_width=True)
            if submit:
                if not data["company_name"]:
                    st.warning("Company name is required.")
                else:
                    ss["wiz_step"] = 1
                    st.rerun()

    # ---- Step 1: Product ----
    elif ss["wiz_step"] == 1:
        with st.form("wiz_step_1"):
            st.markdown("#### Step 2 / 5 — Product details")
            data["product_name"] = st.text_input(
                "Product name *",
                value=data.get("product_name", ""),
                placeholder="e.g. 20000mAh USB-C Power Bank",
            )
            data["product_description"] = st.text_area(
                "Short description",
                value=data.get("product_description", ""),
                placeholder="What the product is, what's inside, who buys it.",
                height=80,
            )
            categories = hts_df["product_category"].tolist() if not hts_df.empty else []
            current_cat = data.get("hts_category", categories[0] if categories else "")
            idx = categories.index(current_cat) if current_cat in categories else 0
            data["hts_category"] = st.selectbox("HTS category *", options=categories, index=idx)

            col_a, col_b = st.columns(2)
            with col_a:
                data["fob_value"] = st.number_input(
                    "FOB value per shipment (USD) *",
                    min_value=0.0, value=float(data.get("fob_value", 10000.0)),
                    step=500.0,
                )
            with col_b:
                data["units_per_shipment"] = st.number_input(
                    "Units per shipment",
                    min_value=1, value=int(data.get("units_per_shipment", 100)),
                    step=10,
                )
            col_c, col_d = st.columns(2)
            with col_c:
                data["shipping_cost"] = st.number_input(
                    "Shipping cost / shipment (USD)",
                    min_value=0.0, value=float(data.get("shipping_cost", 800.0)),
                    step=50.0,
                )
                data["use_sea_freight"] = st.checkbox(
                    "Sea freight (vs. air)", value=data.get("use_sea_freight", True)
                )
            with col_d:
                data["insurance_cost"] = st.number_input(
                    "Insurance / shipment (USD)",
                    min_value=0.0, value=float(data.get("insurance_cost", 100.0)),
                    step=25.0,
                )

            cback, cnext = st.columns([1, 1])
            with cback:
                back = st.form_submit_button("← Back", use_container_width=True)
            with cnext:
                nxt = st.form_submit_button("Next →", type="primary", use_container_width=True)
            if back:
                ss["wiz_step"] = 0
                st.rerun()
            if nxt:
                if not data["product_name"]:
                    st.warning("Product name is required.")
                else:
                    ss["wiz_step"] = 2
                    st.rerun()

    # ---- Step 2: Markets ----
    elif ss["wiz_step"] == 2:
        with st.form("wiz_step_2"):
            st.markdown("#### Step 3 / 5 — Target US markets & sales")
            state_options = sorted(tax_rates.keys()) if tax_rates else []
            data["target_states"] = st.multiselect(
                "Which US states do you sell into? *",
                options=state_options,
                default=data.get("target_states", []),
                help="Pick all states where you have customers or projected sales.",
            )
            data["annual_revenue_usd"] = st.number_input(
                "Projected annual US revenue (USD) *",
                min_value=0.0,
                value=float(data.get("annual_revenue_usd", 500_000.0)),
                step=50_000.0,
            )
            data["sales_channel"] = st.selectbox(
                "Primary sales channel",
                options=[
                    "B2B direct (distributor)",
                    "B2B direct (end customer / OEM)",
                    "Online marketplace (Amazon, eBay, etc.)",
                    "Own DTC e-commerce site",
                    "Retail (brick-and-mortar partner)",
                    "Mixed",
                ],
                index=0,
            )
            cback, cnext = st.columns([1, 1])
            with cback:
                back = st.form_submit_button("← Back", use_container_width=True)
            with cnext:
                nxt = st.form_submit_button("Next →", type="primary", use_container_width=True)
            if back:
                ss["wiz_step"] = 1
                st.rerun()
            if nxt:
                if not data["target_states"]:
                    st.warning("Pick at least one target state.")
                else:
                    ss["wiz_step"] = 3
                    st.rerun()

    # ---- Step 3: Compliance status ----
    elif ss["wiz_step"] == 3:
        with st.form("wiz_step_3"):
            st.markdown("#### Step 4 / 5 — Current compliance status")
            st.caption(
                "Tell us what you've already done — the report will flag what's still missing."
            )
            col_a, col_b = st.columns(2)
            with col_a:
                data["has_customs_broker"] = st.checkbox(
                    "Already engaged a US customs broker",
                    value=data.get("has_customs_broker", False),
                )
                data["has_registered_agent"] = st.checkbox(
                    "Has a US registered agent",
                    value=data.get("has_registered_agent", False),
                )
                data["has_fcc"] = st.checkbox(
                    "FCC certification (or SDoC) obtained",
                    value=data.get("has_fcc", False),
                )
            with col_b:
                data["has_ul"] = st.checkbox(
                    "UL listing or Nationally Recognized Test Lab cert",
                    value=data.get("has_ul", False),
                )
                data["has_fda"] = st.checkbox(
                    "FDA clearance / 510(k) (if medical)",
                    value=data.get("has_fda", False),
                )
                data["has_uflpa_docs"] = st.checkbox(
                    "UFLPA supply-chain documentation prepared",
                    value=data.get("has_uflpa_docs", False),
                )
            data["has_sales_tax_reg"] = st.text_input(
                "States where already registered for sales tax (comma-separated)",
                value=data.get("has_sales_tax_reg", ""),
                placeholder="e.g. Texas, California",
            )
            cback, cnext = st.columns([1, 1])
            with cback:
                back = st.form_submit_button("← Back", use_container_width=True)
            with cnext:
                nxt = st.form_submit_button("Next →", type="primary", use_container_width=True)
            if back:
                ss["wiz_step"] = 2
                st.rerun()
            if nxt:
                ss["wiz_step"] = 4
                st.rerun()

    # ---- Step 4: Review + Generate ----
    elif ss["wiz_step"] == 4:
        st.markdown("#### Step 5 / 5 — Review & generate")

        # Compact summary card
        st.markdown(
            f"""
<div class="cc-card">
  <div class="cc-card-title">{data.get('company_name','—')}</div>
  <div class="cc-card-meta">{data.get('industry','—')} · from {data.get('manufacturing_location') or data.get('origin_country','—')}</div>
  <div class="cc-card-body">
    <b>Product:</b> {data.get('product_name','—')} ({data.get('hts_category','—')})<br/>
    <b>FOB / shipment:</b> ${float(data.get('fob_value', 0)):,.0f} · {data.get('units_per_shipment',0):,} units<br/>
    <b>Target states:</b> {', '.join(data.get('target_states', [])) or '—'}<br/>
    <b>Annual US revenue:</b> ${float(data.get('annual_revenue_usd', 0)):,.0f}
  </div>
</div>
            """,
            unsafe_allow_html=True,
        )

        st.info(
            "Click **Generate PDF** below to compile your personalized report. "
            "The executive summary will be drafted by AI (OrbitAI → Ollama → template fallback) "
            "and the full PDF is rendered locally — nothing is sent to third-party servers "
            "except the prompt for the exec summary."
        )

        cback, cgen = st.columns([1, 2])
        with cback:
            if st.button("← Back to edit", use_container_width=True):
                ss["wiz_step"] = 3
                st.rerun()
        with cgen:
            if st.button("🚀 Generate PDF Report", type="primary", use_container_width=True):
                with st.spinner("Compiling analysis and rendering PDF…"):
                    pdf_bytes = _build_full_report(data, nexus_df, hts_df, tax_rates, alerts_df)
                    ss["wiz_pdf_bytes"] = pdf_bytes
                    ss["wiz_step"] = 5
                    st.rerun()

    # ---- Step 5: Done — download ----
    elif ss["wiz_step"] == 5:
        st.success("✅ Your Customs Compass Report is ready.")
        pdf_bytes = ss.get("wiz_pdf_bytes")
        company = (data.get("company_name") or "report").replace(" ", "_")[:40]
        if pdf_bytes:
            st.download_button(
                "📥 Download PDF Report",
                data=pdf_bytes,
                file_name=f"CustomsCompass_{company}.pdf",
                mime="application/pdf",
                type="primary",
                use_container_width=True,
            )
            st.caption(f"PDF size: {len(pdf_bytes)/1024:,.1f} KB")
        col_a, col_b = st.columns(2)
        with col_a:
            if st.button("🔁 Generate another", use_container_width=True):
                ss["wiz_step"] = 0
                ss["wiz_data"] = {}
                ss.pop("wiz_pdf_bytes", None)
                st.rerun()
        with col_b:
            if st.button("✏️ Edit answers", use_container_width=True):
                ss["wiz_step"] = 0
                st.rerun()


def _build_full_report(
    data: dict,
    nexus_df: pd.DataFrame,
    hts_df: pd.DataFrame,
    tax_rates: dict,
    alerts_df: pd.DataFrame,
) -> bytes:
    """Run the full analysis pipeline and render the PDF."""
    # Lookup HTS row
    hts_row = hts_df[hts_df["product_category"] == data.get("hts_category", "")]
    hts_info = {}
    if not hts_row.empty:
        r = hts_row.iloc[0]
        hts_info = {
            "hts_code": r["hts_code"],
            "duty_rate": r["duty_rate"],
            "fcc_needed": r["fcc_needed"],
            "ul_needed": r["ul_needed"],
            "fda_needed": r["fda_needed"],
            "hts_notes": r["notes"],
        }

    # Landed cost
    breakdown = compute_landed_cost(
        customs_value_usd=float(data.get("fob_value", 0)),
        duty_rate_str=hts_info.get("duty_rate", "0%"),
        apply_section_301=(data.get("origin_country", "China").lower() == "china"),
        shipping_usd=float(data.get("shipping_cost", 0)),
        insurance_usd=float(data.get("insurance_cost", 0)),
        destination_state=(data.get("target_states") or ["Texas"])[0],
        tax_rates=tax_rates,
        use_sea_freight=bool(data.get("use_sea_freight", True)),
    )

    # State-by-state nexus (assume equal distribution if multiple states)
    target_states = data.get("target_states") or []
    annual = float(data.get("annual_revenue_usd", 0))
    per_state = annual / len(target_states) if target_states else 0
    nexus_rows = []
    overall_risk = "Low"
    for s in target_states:
        nrow = nexus_df[nexus_df["state"] == s]
        if nrow.empty:
            continue
        threshold = float(nrow.iloc[0]["threshold_usd"])
        risk = assess_nexus_risk(per_state, threshold)
        if risk == "High":
            status = "🔴 Obligation triggered — must register"
            overall_risk = "High"
        elif risk == "Medium":
            status = "🟡 Approaching threshold"
            if overall_risk == "Low":
                overall_risk = "Medium"
        elif threshold == 0:
            status = "🟢 No state sales tax"
        else:
            status = "🟢 Below threshold"
        nexus_rows.append({
            "state": s,
            "threshold": threshold,
            "projected_sales": per_state,
            "status": status,
            "tax_rate": float(tax_rates.get(s, 0) or 0),
        })

    # Relevant alerts (search by product + question proxy)
    fake_question = (
        f"{data.get('product_description','')} {data.get('product_name','')} "
        f"from {data.get('origin_country','China')} to {' '.join(target_states)}"
    )
    rel_alerts = find_relevant_alerts(alerts_df, [data.get("hts_category", "")], fake_question)
    alerts_payload = rel_alerts.to_dict(orient="records") if not rel_alerts.empty else []

    # Compliance checklist
    checklist = [
        {"label": "Engage US customs broker before first shipment",
         "done": data.get("has_customs_broker", False)},
        {"label": "Appoint US registered agent (for legal service)",
         "done": data.get("has_registered_agent", False)},
    ]
    if str(hts_info.get("fcc_needed", "")).lower() == "yes":
        checklist.append({"label": "Obtain FCC certification (or SDoC)",
                          "done": data.get("has_fcc", False)})
    if str(hts_info.get("ul_needed", "")).lower() == "yes":
        checklist.append({"label": "Obtain UL listing or NRTL certification",
                          "done": data.get("has_ul", False)})
    if str(hts_info.get("fda_needed", "")).lower() == "yes":
        checklist.append({"label": "Obtain FDA clearance / 510(k) where applicable",
                          "done": data.get("has_fda", False)})
    if data.get("origin_country", "China").lower() == "china":
        checklist.append({"label": "Prepare UFLPA supply-chain affidavits",
                          "done": data.get("has_uflpa_docs", False)})
        checklist.append({"label": "Verify HTS code is NOT on Section 301 exclusion list",
                          "done": False})
    registered_states = {
        s.strip() for s in (data.get("has_sales_tax_reg", "") or "").split(",") if s.strip()
    }
    for row in nexus_rows:
        if row["status"].startswith("🔴"):
            checklist.append({
                "label": f"Register for sales tax in {row['state']}",
                "done": row["state"] in registered_states,
            })

    # Executive summary (AI-generated)
    report_data = {
        **data,
        **hts_info,
        "cost_breakdown": breakdown,
        "landed_cost_total": breakdown["total_landed"],
        "nexus_analysis": nexus_rows,
        "target_states_count": len(target_states),
        "alerts": alerts_payload,
        "checklist": checklist,
        "overall_risk": overall_risk,
    }
    report_data["exec_summary"] = _orbitai_executive_summary(report_data)

    return generate_pdf_report(report_data)


# =============================================================================
# Section M — Live Intel Agent (Federal Register + Web Search + OrbitAI)
# =============================================================================
#
# Unlike the static CSV / RAG layers, this section pulls FRESH data from public
# APIs on demand, then asks an LLM (OrbitAI preferred, Ollama fallback) to
# synthesize a citable analysis. Results are cached 15 minutes to stay fast
# and avoid hammering upstream services.

FEDERAL_REGISTER_API = "https://www.federalregister.gov/api/v1/documents.json"
DDG_HTML_ENDPOINT = "https://html.duckduckgo.com/html/"
LIVE_INTEL_CACHE_TTL = 900  # 15 minutes


@st.cache_data(ttl=LIVE_INTEL_CACHE_TTL, show_spinner=False)
def fetch_federal_register(query: str, days_back: int = 365, max_results: int = 4) -> list[dict]:
    """Fetch recent Federal Register notices matching a query.

    The Federal Register API is free, public, and well-documented:
      https://www.federalregister.gov/developers/documentation/api/v1

    Returns: list of {title, date, url, abstract, agencies}. Empty list on failure.
    """
    from datetime import date, timedelta
    cutoff = (date.today() - timedelta(days=days_back)).isoformat()
    params = {
        "conditions[term]": query,
        "conditions[publication_date][gte]": cutoff,
        "per_page": max_results,
        "order": "relevance",
        "fields[]": ["title", "publication_date", "html_url", "abstract", "agency_names"],
    }
    try:
        resp = requests.get(
            FEDERAL_REGISTER_API,
            params=params,
            headers={"User-Agent": USER_AGENT, "Accept": "application/json"},
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
    except (requests.RequestException, ValueError):
        return []

    out = []
    for d in data.get("results", []) or []:
        out.append({
            "title": d.get("title", "") or "",
            "date": d.get("publication_date", "") or "",
            "url": d.get("html_url", "") or "",
            "abstract": (d.get("abstract") or "")[:400],
            "agencies": d.get("agency_names") or [],
        })
    return out


@st.cache_data(ttl=LIVE_INTEL_CACHE_TTL, show_spinner=False)
def web_search_ddg(query: str, max_results: int = 4) -> list[dict]:
    """Scrape DuckDuckGo's HTML search endpoint.

    Returns: list of {title, url, snippet}. Empty list on failure / block.
    """
    try:
        resp = requests.post(
            DDG_HTML_ENDPOINT,
            data={"q": query, "kl": "us-en"},
            headers={
                "User-Agent": USER_AGENT,
                "Accept": "text/html,application/xhtml+xml",
                "Accept-Language": "en-US,en;q=0.9",
            },
            timeout=10,
        )
        resp.raise_for_status()
    except requests.RequestException:
        return []

    try:
        from bs4 import BeautifulSoup
    except ImportError:
        return []

    soup = BeautifulSoup(resp.text, "html.parser")
    from urllib.parse import unquote, urlparse, parse_qs

    results: list[dict] = []
    for result_div in soup.select(".result")[: max_results * 2]:
        title_el = result_div.select_one(".result__title a, h2 a")
        snippet_el = result_div.select_one(".result__snippet, .result__body")
        if not title_el:
            continue
        title = title_el.get_text(strip=True)
        url = title_el.get("href", "") or ""
        snippet = snippet_el.get_text(" ", strip=True) if snippet_el else ""

        # DDG wraps real URLs in a redirect: //duckduckgo.com/l/?uddg=<encoded>
        if "duckduckgo.com/l/" in url:
            try:
                full = "https:" + url if url.startswith("//") else url
                parsed = parse_qs(urlparse(full).query)
                if "uddg" in parsed:
                    url = unquote(parsed["uddg"][0])
            except Exception:
                pass

        if not url or not title:
            continue
        results.append({"title": title, "url": url, "snippet": snippet[:400]})
        if len(results) >= max_results:
            break
    return results


def _live_intel_query_terms(question: str, hts_df: pd.DataFrame) -> tuple[str, str]:
    """Derive Federal Register & web search queries from a user question."""
    products = extract_product_categories(question, hts_df)
    states = extract_states(question)

    base_terms: list[str] = []
    for p in products:
        base_terms.append(p.replace("_", " "))
    if not base_terms:
        # Fall back to content words from the question
        words = re.findall(r"\b[a-zA-Z]{4,}\b", question.lower())
        skipwords = _STOPWORDS | {"china", "chinese", "import", "export"}
        for w in words[:4]:
            if w not in skipwords:
                base_terms.append(w)
    if not base_terms:
        base_terms = ["import tariff"]

    core = " ".join(base_terms[:3])
    fed_query = f"{core} tariff"
    web_query = f"{core} US import tariff Section 301 2026"
    if states:
        web_query += f" {states[0]}"
    return fed_query, web_query


def run_live_intel_agent(question: str, hts_df: pd.DataFrame) -> dict:
    """Live intelligence agent: fetch fresh sources, synthesize via OrbitAI/Ollama.

    Returns a dict with:
      - summary (str): AI-synthesized answer with [n] citations
      - sources (list[dict]): normalized source list for the UI
      - fed_reg_items (list[dict])
      - web_items (list[dict])
      - engine (str): "orbitai" | "ollama" | "raw"
      - fed_query, web_query (str): the actual queries used (for transparency)
    """
    fed_query, web_query = _live_intel_query_terms(question, hts_df)
    fed_reg_items = fetch_federal_register(fed_query)
    web_items = web_search_ddg(web_query)

    # Build a numbered source list (capped to keep LLM prompt small enough for fast inference)
    sources: list[dict] = []
    for i, item in enumerate(fed_reg_items[:4], start=1):
        sources.append({
            "n": i,
            "type": "Federal Register",
            "title": item["title"],
            "date": item["date"],
            "url": item["url"],
            "snippet": (item["abstract"] or item["title"])[:220],
        })
    offset = len(sources)
    for i, item in enumerate(web_items[:3], start=offset + 1):
        sources.append({
            "n": i,
            "type": "Web",
            "title": item["title"],
            "date": "",
            "url": item["url"],
            "snippet": item["snippet"][:220],
        })

    result_skeleton = {
        "fed_reg_items": fed_reg_items,
        "web_items": web_items,
        "sources": sources,
        "fed_query": fed_query,
        "web_query": web_query,
    }

    if not sources:
        return {
            **result_skeleton,
            "summary": (
                "⚠️ Could not fetch any live sources (Federal Register or web search both failed). "
                "Possible causes: network offline, upstream rate-limiting, or anti-bot block. "
                "The static knowledge base in the **Analyze** tab still works fully offline."
            ),
            "engine": "raw",
        }

    # Compact synthesis prompt — kept short so llama3.2:3b on CPU returns in <60s
    context_block = "\n".join(
        f"[{s['n']}] {s['title']} ({s.get('date') or 'undated'}). {s['snippet']}"
        for s in sources
    )
    system = (
        "You are a US trade-intelligence analyst. Synthesize the sources into 3-5 "
        "bullet points answering the user's question. Cite sources as [1], [2], etc. "
        "Only state facts present in the sources. End with one line: "
        "'Risk: Low/Medium/High — most recent source: <date>'."
    )
    user = f"Question: {question}\n\nSources:\n{context_block}\n\nAnswer:"

    # Try OrbitAI first (premium model)
    if is_orbitai_configured():
        try:
            answer = call_orbitai(system, user, temperature=0.3)
            if answer:
                return {**result_skeleton, "summary": answer, "engine": "orbitai"}
        except Exception:
            pass

    # Fallback: local Ollama
    if is_ollama_available():
        try:
            answer = call_ollama(system, user)
            if answer:
                return {**result_skeleton, "summary": answer, "engine": "ollama"}
        except Exception:
            pass

    # Final fallback: raw list (no LLM)
    bullets = "\n".join(
        f"- **[{s['n']}] {s['title']}** ({s.get('date', '')})  \n  {s['snippet']}"
        for s in sources
    )
    raw = (
        "_LLM unavailable — showing raw fetched sources below. Static analysis "
        "in the Analyze tab remains fully functional._\n\n" + bullets
    )
    return {**result_skeleton, "summary": raw, "engine": "raw"}


def render_live_intel_tab(hts_df: pd.DataFrame) -> None:
    st.markdown("### 🌐 Live Intel Agent")
    st.caption(
        "Pulls FRESH data from Federal Register (US government) + web search, then "
        "synthesizes via OrbitAI (premium) or Ollama (local fallback). Results cached 15 min."
    )

    ss = st.session_state
    if "live_intel_result" not in ss:
        ss["live_intel_result"] = None

    quick_examples = [
        "What are the latest Section 301 tariff changes on lithium batteries from China?",
        "Recent UFLPA enforcement actions on solar panel imports",
        "New CBP rulings on consumer electronics origin verification",
        "Federal Register notices about HTS chapter 85 in 2026",
    ]

    with st.container():
        col_a, col_b = st.columns([3, 1])
        with col_a:
            question = st.text_area(
                "Your question (live data, fresh from upstream)",
                value=ss.get("live_intel_question", ""),
                placeholder="e.g. Are there new Section 301 tariffs on Chinese lithium batteries this year?",
                height=80,
                key="live_intel_question_input",
            )
        with col_b:
            st.markdown("&nbsp;", unsafe_allow_html=True)
            run_clicked = st.button("🔍 Fetch Live", type="primary", use_container_width=True)
            if st.button("🧹 Clear cache", use_container_width=True,
                         help="Force a fresh fetch (clears the 15-min cache)"):
                fetch_federal_register.clear()  # type: ignore[attr-defined]
                web_search_ddg.clear()          # type: ignore[attr-defined]
                st.success("Live cache cleared.")

    st.markdown("##### Quick examples")
    cols = st.columns(len(quick_examples))
    for i, ex in enumerate(quick_examples):
        with cols[i]:
            if st.button(f"💡 {ex.split('?')[0][:42]}…", key=f"live_ex_{i}", use_container_width=True):
                ss["live_intel_question"] = ex
                st.rerun()

    if run_clicked and (question or "").strip():
        ss["live_intel_question"] = question
        with st.spinner("Fetching Federal Register + web search + synthesizing with AI…"):
            result = run_live_intel_agent(question.strip(), hts_df)
        ss["live_intel_result"] = result

    result = ss.get("live_intel_result")
    if not result:
        st.info(
            "👆 Type a question and click **Fetch Live**, or pick a quick example. "
            "Unlike the static **Analyze** tab, this fetches fresh data from "
            "[federalregister.gov](https://www.federalregister.gov) and the open web."
        )
        return

    # ---- Render the result ----
    from datetime import datetime, timezone
    now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    engine = result.get("engine", "raw").upper()
    engine_badge = {
        "ORBITAI": ("cc-badge-info", "🛰️ OrbitAI"),
        "OLLAMA": ("cc-badge-low", "🦙 Ollama local"),
        "RAW": ("cc-badge-unknown", "📋 Raw fetch (no LLM)"),
    }.get(engine, ("cc-badge-unknown", engine))

    st.markdown(
        f'<div style="margin:10px 0">'
        f'<span class="cc-badge {engine_badge[0]}">{engine_badge[1]}</span> '
        f'<span class="cc-pill">As of {now}</span> '
        f'<span class="cc-pill">{len(result.get("sources", []))} sources</span>'
        f'</div>',
        unsafe_allow_html=True,
    )

    st.markdown("#### 🧠 Synthesized analysis")
    st.markdown(result["summary"])

    sources = result.get("sources") or []
    if sources:
        with st.expander(f"📎 Sources ({len(sources)})", expanded=True):
            import html as _html
            for s in sources:
                badge_cls = "cc-badge-info" if s["type"] == "Federal Register" else "cc-badge-unknown"
                st.markdown(
                    f"""
<div class="cc-card">
  <div class="cc-card-title">
    <span class="cc-badge {badge_cls}">[{s['n']}] {s['type']}</span>
    &nbsp;{_html.escape(s['title'])}
  </div>
  <div class="cc-card-meta">{s.get('date', '') or 'undated'} · <a href="{s['url']}" target="_blank">{s['url']}</a></div>
  <div class="cc-card-body">{_html.escape(s['snippet'])}</div>
</div>
                    """,
                    unsafe_allow_html=True,
                )

    with st.expander("🔎 Debug: queries used"):
        st.json({
            "fed_register_query": result.get("fed_query"),
            "web_search_query": result.get("web_query"),
            "fed_register_count": len(result.get("fed_reg_items", [])),
            "web_count": len(result.get("web_items", [])),
        })


# =============================================================================
# Section N — Main Entry Point
# =============================================================================

def main() -> None:
    # Inject premium CSS first so everything renders polished
    inject_custom_css()

    # Load data
    try:
        nexus_df = load_nexus_thresholds()
        hts_df = load_hts_duty_codes()
        tax_rates = load_tax_rates()
        alerts_df = load_cbp_alerts()
        chunk_index = load_cbp_chunks_index()
    except (ValueError, json.JSONDecodeError) as e:
        st.error(f"Data file error: {e}")
        return

    if nexus_df.empty or hts_df.empty:
        st.error(
            "Required data files are missing or empty. Ensure "
            "`nexus_thresholds.csv`, `hts_duty_codes.csv`, and "
            "`tax_rates_by_state.json` exist in the same folder as `app.py`."
        )
        return

    # Session state defaults
    if "news_enabled" not in st.session_state:
        st.session_state["news_enabled"] = True
    if "prefilled" not in st.session_state:
        st.session_state["prefilled"] = ""

    # Initial sidebar render (we need to know if news is enabled before fetching)
    news_items: list[dict] = []
    if st.session_state["news_enabled"]:
        news_items = fetch_cbp_news()

    ollama_ok = is_ollama_available()

    sidebar_state = render_sidebar(
        ollama_ok=ollama_ok,
        news_count=len(news_items),
        news_enabled=st.session_state["news_enabled"],
        alerts_count=len(alerts_df),
        chunks_count=chunk_index.get("N", 0),
    )

    if sidebar_state["refresh_news"]:
        fetch_cbp_news.clear()  # type: ignore[attr-defined]
        st.rerun()

    if sidebar_state["news_enabled"] != st.session_state["news_enabled"]:
        st.session_state["news_enabled"] = sidebar_state["news_enabled"]
        st.rerun()

    if sidebar_state["example_clicked"]:
        st.session_state["prefilled"] = sidebar_state["example_clicked"]
        # Switch to Analyze tab when clicking an example
        st.session_state["active_tab"] = "analyze"
        st.rerun()

    # Hero header
    render_hero(
        nexus_count=len(nexus_df),
        alerts_count=len(alerts_df),
        chunks_count=chunk_index.get("N", 0),
        ollama_ok=ollama_ok,
    )

    tab_analyze, tab_cost, tab_live, tab_report, tab_kb = st.tabs([
        "🔍  Analyze",
        "💰  Landed Cost Calculator",
        "🌐  Live Intel",
        "📄  Generate Report",
        "📚  Knowledge Base",
    ])

    # ---- Tab 1: Analyze (Q&A flow) ----
    with tab_analyze:
        form = render_main_form(prefilled_question=st.session_state.get("prefilled", ""))

        if form["submitted"]:
            combined = (form["product_desc"] + "\n" + form["question"]).strip()
            if not combined:
                st.warning("Please enter a product description or a question.")
            else:
                with st.spinner("Analyzing… (consulting knowledge base and LLM)"):
                    answer, mode, payload = get_answer(
                        question=combined,
                        nexus_df=nexus_df,
                        hts_df=hts_df,
                        tax_rates=tax_rates,
                        news_items=news_items,
                        alerts_df=alerts_df,
                        chunk_index=chunk_index,
                        force_fallback=sidebar_state["force_fallback"],
                    )

                render_response(answer, mode, payload, news_items)

    # ---- Tab 2: Cost Calculator ----
    with tab_cost:
        render_cost_calculator_tab(hts_df, tax_rates, alerts_df)

    # ---- Tab 3: Live Intel (real-time web + Federal Register) ----
    with tab_live:
        render_live_intel_tab(hts_df)

    # ---- Tab 4: Generate Report (wizard + PDF) ----
    with tab_report:
        render_report_wizard_tab(nexus_df, hts_df, tax_rates, alerts_df)

    # ---- Tab 5: Knowledge Base ----
    with tab_kb:
        render_knowledge_tab(alerts_df, chunk_index)


if __name__ == "__main__":
    main()