src/streamlit_app.py

"""
app.py — Macro Data Tool
──────────────────────────────────────────────────────────────
Streamlit UI that calls your 4 existing scripts in sequence:
  Step 1: python generate_excel.py
  Step 2: python generate_narratives_v1.py
  Step 3: python update_template.py  →  python generate_word.py

Run:
    streamlit run app.py
"""

import streamlit as st
import json
import time
import os
import subprocess
from pathlib import Path
import sys

st.set_page_config(page_title="Macro Data Tool", page_icon="📊",
                   layout="wide", initial_sidebar_state="expanded")

st.markdown("""
<style>
    .block-container { padding-top: 2rem; max-width: 1000px; }
    .log-line { font-family: 'Consolas', monospace; font-size: 13px; color: #8BC34A; margin: 2px 0; }
    .log-container { background: #1a1a2e; border-radius: 8px; padding: 16px 20px; margin: 8px 0; }
    .log-warning { color: #FFC107; }
    .log-error   { color: #FF5252; }
    .log-info    { color: #64B5F6; }
    .step-header {
        background: linear-gradient(90deg, #4472C4 0%, #5BA3C9 100%);
        color: white; padding: 10px 20px; border-radius: 8px;
        font-weight: 600; font-size: 16px; margin: 24px 0 12px 0;
    }
</style>
""", unsafe_allow_html=True)

# ── Session state ────────────────────────────────────────────
for k, d in [("step",0),("indicators",None),("ratings",None),
             ("narratives",None),("template_path",None)]:
    if k not in st.session_state:
        st.session_state[k] = d

OUTPUT_DIR = Path("output")
OUTPUT_DIR.mkdir(exist_ok=True)


# ══════════════════════════════════════════════════════════════
# HELPERS
# ══════════════════════════════════════════════════════════════

def load_json(path):
    p = Path(path)
    return json.loads(p.read_text(encoding="utf-8")) if p.exists() else None

def save_json(path, data):
    Path(path).write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")

# def build_env(or_key, fc_key):
#     env = os.environ.copy()
#     env["PYTHONIOENCODING"] = "utf-8"
#     env["FIRECRAWL_API_KEY"] = fc_key or ""
#     env["OPENROUTER_API_KEY"] = or_key or ""
#     # Ensure subprocess uses the same Python as Streamlit
#     env["PATH"] = os.path.dirname(sys.executable) + os.pathsep + env.get("PATH", "")
#     return env

def build_env(or_key, fc_key):
    env = os.environ.copy()
    env["PYTHONIOENCODING"] = "utf-8"
    env["FIRECRAWL_API_KEY"] = fc_key or ""
    env["OPENROUTER_API_KEY"] = or_key or ""
    env["COUNTRY_NAME"] = country
    env["PATH"] = os.path.dirname(sys.executable) + os.pathsep + env.get("PATH", "")
    return env

def run_script(name, env, timeout=180):
    r = subprocess.run([sys.executable, name], capture_output=True, text=True,
                       encoding="utf-8", errors="replace", timeout=timeout, env=env)
    return r.returncode == 0, r.stdout or "", r.stderr or ""

def show_output(stdout, stderr):
    lines = [l for l in stdout.strip().split("\n") if l.strip()]
    if not lines:
        return
    html = '<div class="log-container">'
    for line in lines:
        css = ""
        ll = line.lower()
        if any(w in ll for w in ["error","failed","❌"]):
            css = "log-error"
        elif any(w in ll for w in ["⚠","warn","skip"]):
            css = "log-warning"
        elif any(w in ll for w in ["fetch","call","load","scrape","generat","running","read"]):
            css = "log-info"
        html += f'<div class="log-line {css}">{line}</div>'
    html += '</div>'
    st.markdown(html, unsafe_allow_html=True)
    # Show errors if any
    if stderr and any(w in stderr.lower() for w in ["error","traceback","exception"]):
        with st.expander("Show error details"):
            st.code(stderr[:1500])


# ══════════════════════════════════════════════════════════════
# HEADER
# ══════════════════════════════════════════════════════════════

st.markdown("## 📊 Macro Data Tool")
st.markdown("*AI-powered country risk report generator*")
st.divider()


# ══════════════════════════════════════════════════════════════
# SIDEBAR
# ══════════════════════════════════════════════════════════════

with st.sidebar:
    st.markdown("### ⚙️ Configuration")
    or_key = st.text_input("OpenRouter API Key", type="password",
                           help="https://openrouter.ai/keys")
    fc_key = st.text_input("Firecrawl API Key", type="password",
                           help="https://www.firecrawl.dev/app/api-keys")
    st.caption("🔒 Keys are stored only in your browser session and never saved.")
    
    st.divider()
    st.markdown("### 🔗 Custom Data Sources")
    st.caption("*(Coming soon)*")
    st.text_input("GDP Source URL", disabled=True, placeholder="https://...")
    st.text_input("Monetary Policy URL", disabled=True, placeholder="https://...")


# ══════════════════════════════════════════════════════════════
# TOP — Upload + Country + Period
# ══════════════════════════════════════════════════════════════

c1, c2, c3 = st.columns([2, 1, 1])
with c1:
    template_file = Path("Country_Risk_profile_template.xlsx")
    if not template_file.exists():
        # Try parent directory (Docker copies files to /app)
        template_file = Path("/app/Country_Risk_profile_template.xlsx")

    if template_file.exists():
        st.success("📁 Template loaded: Country_Risk_profile_template.xlsx")
        st.session_state.template_path = str(template_file)
    else:
        st.error("❌ Template file not found. Please ensure it is included in the deployment.")
with c2:
    country = st.selectbox("🌍 Country", ["India", "Indonesia"])
    # Reset state when country changes
    if "prev_country" not in st.session_state:
        st.session_state.prev_country = country
    if country != st.session_state.prev_country:
        st.session_state.prev_country = country
        st.session_state.step = 0
        st.session_state.indicators = None
        st.session_state.ratings = None
        st.session_state.narratives = None
with c3:
    quarter = st.selectbox("📅 Period", ["Q1 2026","Q2 2026","Q3 2026","Q4 2026"])
    q_part, y_part = quarter.split(" ")


# ══════════════════════════════════════════════════════════════
# STEP 1 — generate_excel.py
# ══════════════════════════════════════════════════════════════

st.markdown('<div class="step-header">Step 1 — Collect Macroeconomic Data</div>', unsafe_allow_html=True)

_, btn_col1 = st.columns([4, 1])
with btn_col1:
    fetch_clicked = st.button("🔍  Fetch Data", type="primary", use_container_width=True)
if fetch_clicked:
    if not st.session_state.template_path:
        st.error("⚠️ Please upload the Excel template first.")
    elif not fc_key:
        st.error("⚠️ Please enter your Firecrawl API key in the sidebar.")
    else:
        env = build_env(or_key, fc_key)
        prog = st.progress(0, "Fetching macro data from World Bank, IMF, ECB ...")

        prog.progress(10, "Connecting to World Bank, IMF, ECB data sources ...")
        ok, out, err = run_script("generate_excel.py", env, timeout=120)

        prog.progress(80, "Processing fetched data ...")
        show_output(out, err)

        if ok:
            data = load_json(OUTPUT_DIR / f"{country}_data.json")
            if data:
                st.session_state.indicators = data["indicators"]
                st.session_state.ratings = data.get("ratings", {})
                st.session_state.step = max(st.session_state.step, 1)
                prog.progress(100, "Data fetch complete!")
                time.sleep(0.5)
                prog.empty()
            else:
                st.error("❌ Data file not generated. Check the log above for errors.")
                prog.empty()
        else:
            st.error("❌ Data fetch failed. Please check your internet connection and Firecrawl API key, then click 'Fetch Data' again.")
            prog.empty()

# ── Editable indicators ─────────────────────────────────────
if st.session_state.indicators:
    st.markdown("#### 📋 Fetched Indicators")
    st.caption("Edit values if needed. Changes save automatically.")
    
    hdr_a, hdr_b, hdr_c = st.columns([3, 2, 2])
    hdr_a.markdown("**Indicator**")
    hdr_b.markdown("**Value**")
    hdr_c.markdown("**Source**")

    
    edited = []
    for i, ind in enumerate(st.session_state.indicators):
        a, b, c = st.columns([3, 2, 2])
        a.markdown(f"`{ind['name']}`")
        cur = str(ind["value"]) if ind["value"] is not None else ""
        nv = b.text_input(f"v{i}", value=cur, key=f"iv_{i}", label_visibility="collapsed")

        # Source column
        # Source column — fix broken World Bank links
        src = ind.get("source", "")
        broken_wb = ["SL.UEM.TOTL.ZS", "FI.RES.TOTL.CD"]
        
        if src and src.startswith("http"):
            # Check if this is a known broken World Bank link
            is_broken = any(code in src for code in broken_wb)
            is_ecb_api = "data-api.ecb" in src
            
            if is_broken:
                c.caption("World Bank API")
            elif is_ecb_api:
                c.caption("European Central Bank (ECB)")
            elif "worldbank" in src:
                c.caption("[World Bank](%s)" % src)
            elif "imf.org" in src:
                c.caption("IMF World Economic Outlook")
            elif "tradingeconomics" in src:
                c.caption("[Trading Economics](%s)" % src)
            else:
                c.caption("[Link](%s)" % src)
        elif src:
            c.caption(src)
        else:
            c.caption("—")
        try:
            pv = float(nv) if nv else None
        except ValueError:
            pv = ind["value"]
        edited.append({**ind, "value": pv})

    # if edited != st.session_state.indicators:
    st.session_state.indicators = edited
    dp = OUTPUT_DIR / f"{country}_data.json"
    ex = load_json(dp) or {}
    ex["indicators"] = edited
    save_json(dp, ex)

    if st.session_state.ratings:
        st.markdown("#### 🏦 Credit Ratings")
        rc = st.columns(3)
        for i, ag in enumerate(["Fitch", "Moody's", "S&P"]):
            d = st.session_state.ratings.get(ag, {})
            rc[i].metric(ag, d.get("rating","N/A"), d.get("outlook",""))
            rc[i].caption(d.get("date", ""))


# ══════════════════════════════════════════════════════════════
# STEP 2 — generate_narratives_v1.py
# ══════════════════════════════════════════════════════════════

st.markdown('<div class="step-header">Step 2 — Generate Analyst Narratives</div>', unsafe_allow_html=True)

s2 = st.session_state.step >= 1
_, btn_col2 = st.columns([4, 1])
with btn_col2:
    narr_clicked = st.button("🤖  Generate Narratives", type="primary" if s2 else "secondary",
                              disabled=not s2, use_container_width=True)
if narr_clicked:
    if not or_key:
        st.error("⚠️ Please enter your OpenRouter API key in the sidebar (click ⚙️ top-left).")
    else:
        # Force save current indicators before generating narratives
        dp = OUTPUT_DIR / f"{country}_data.json"
        current_data = load_json(dp) or {}
        if st.session_state.indicators:
            current_data["indicators"] = st.session_state.indicators
        if st.session_state.ratings:
            current_data["ratings"] = st.session_state.ratings
        save_json(dp, current_data)
        
        env = build_env(or_key, fc_key)
        prog = st.progress(0, "Initializing AI pipeline ...")

        # Show "thinking" indicator
        status = st.empty()
        status.markdown(
            '<div class="log-container">'
            '<div class="log-line log-info">🤖  Calling AI model ...</div>'
            '<div class="log-line">    Generating macroeconomic narratives ...</div>'
            '<div class="log-line">    This may take 30-60 seconds ...</div>'
            '</div>', unsafe_allow_html=True)

        prog.progress(15, "🤖 AI is analyzing macroeconomic data ...")
        ok, out, err = run_script("generate_narratives_v1.py", env, timeout=180)

        status.empty()
        prog.progress(85, "Loading results ...")
        show_output(out, err)

        if ok:
            narr = load_json(OUTPUT_DIR / f"{country}_narratives.json")
            if narr and "narratives" in narr:
                st.session_state.narratives = narr["narratives"]
                st.session_state.step = max(st.session_state.step, 2)
                prog.progress(100, "AI generation complete!")
                time.sleep(0.5)
                prog.empty()
            else:
                st.error("❌ Narratives not generated. Check the log above.")
                prog.empty()
        else:
            st.error("❌ AI generation failed. Please check your OpenRouter API key in the sidebar and ensure you have sufficient credits, then try again.")
            prog.empty()

# ── Editable narrative cards ─────────────────────────────────
if st.session_state.narratives:
    st.markdown("#### 📝 Generated Narratives")
    st.caption("Review and edit. Changes save automatically.")

    labels = {"gdp":"GDP", "ratings":"Rating Update", "debt_to_gdp":"Debt to GDP", "unemployment":"Unemployment",
              "fiscal_policy":"Fiscal Policy", "monetary_policy":"Monetary Policy"}

    te = {"India": "india", "Indonesia": "indonesia"}.get(country, "india")  # Trading Economics slug
    source_map = {
        "gdp": (
            "Sources: World Bank API (GDP, GDP Growth Rate, GDP per Capita) — "
            "derived from key indicators above"
        ),
        "ratings": (
            "Sources: Wikipedia — [List of countries by credit rating]"
            "(https://en.wikipedia.org/wiki/List_of_countries_by_credit_rating); "
            "IMF WEO (Government Debt to GDP) — derived from key indicators above"
        ),
        "unemployment": (
            "Sources: World Bank API (Unemployment Rate, Population, GDP Growth Rate) — "
            "derived from key indicators above"
        ),
        "debt_to_gdp": (
            "Sources: IMF WEO (Debt to GDP) — derived from key indicators above; "
            f"[Trading Economics](https://tradingeconomics.com/{te}/government-debt-to-gdp)"
        ),
        "fiscal_policy": (
            "Sources: IMF WEO, World Bank API — derived from key indicators above; "
            f"Trading Economics ([budget](https://tradingeconomics.com/{te}/government-budget), "
            f"[debt](https://tradingeconomics.com/{te}/government-debt))"
        ),
        "monetary_policy": (
            "Sources: World Bank API (CPI), Firecrawl (Key Rate) — derived from key indicators above; "
            f"[Trading Economics](https://tradingeconomics.com/{te}/currency)"
        ),
    }

    ed = {}
    for key, label in labels.items():
        st.markdown(f"**{label}**")
        ed[key] = st.text_area(f"n_{key}", value=st.session_state.narratives.get(key,""),
                               height=120, key=f"ne_{key}", label_visibility="collapsed")
        # Show source below the text area
        src = source_map.get(key, "")
        if src:
            st.caption(src)

    if ed != st.session_state.narratives:
        st.session_state.narratives = ed
        np = OUTPUT_DIR / f"{country}_narratives.json"
        ex = load_json(np) or {}
        ex["narratives"] = ed
        save_json(np, ex)


# ══════════════════════════════════════════════════════════════
# STEP 3 — update_template.py + generate_word.py
# ══════════════════════════════════════════════════════════════

st.markdown('<div class="step-header">Step 3 — Generate Reports</div>', unsafe_allow_html=True)

s3 = st.session_state.step >= 2
_, btn_col3 = st.columns([4, 1])
with btn_col3:
    report_clicked = st.button("📄  Generate Reports", type="primary" if s3 else "secondary",
                                disabled=not s3, use_container_width=True)
if report_clicked:
    if not st.session_state.template_path:
        st.error("⚠️ Please upload the Excel template first.")
    else:
        
        # Force save current data before generating reports
        dp = OUTPUT_DIR / f"{country}_data.json"
        current_data = load_json(dp) or {}
        if st.session_state.indicators:
            current_data["indicators"] = st.session_state.indicators
        if st.session_state.ratings:
            current_data["ratings"] = st.session_state.ratings
        save_json(dp, current_data)

        np = OUTPUT_DIR / f"{country}_narratives.json"
        narr_data = load_json(np) or {}
        if st.session_state.narratives:
            narr_data["narratives"] = st.session_state.narratives
        save_json(np, narr_data)
        
        env = build_env(or_key, fc_key)
        prog = st.progress(0, "Generating reports ...")

        # Step 3a: Excel
        prog.progress(20, "Writing data into Excel template ...")
        ok1, out1, err1 = run_script("update_template.py", env, timeout=60)
        show_output(out1, err1)
        if not ok1:
            st.error("❌ Excel generation failed. Please check the template file.")
            prog.empty()
            st.stop()

        # Step 3b: Word
        prog.progress(60, "Formatting Word report ...")
        ok2, out2, err2 = run_script("generate_word.py", env, timeout=60)
        show_output(out2, err2)
        if not ok2:
            st.error("❌ Word report failed. Make sure Node.js and docx package are installed.")
            prog.empty()
            st.stop()

        prog.progress(100, "Reports generated!")
        time.sleep(0.5)
        prog.empty()
        st.session_state.step = 3
        st.success("✅ Reports generated successfully!")


# ── Download buttons ─────────────────────────────────────────
if st.session_state.step >= 3:
    st.markdown("#### 📥 Download Reports")
    d1, d2 = st.columns(2)

    xlsx = OUTPUT_DIR / f"{country}_{y_part}_{q_part}_updated.xlsx"
    docx = OUTPUT_DIR / f"{country}_{q_part}_{y_part}_report.docx"

    with d1:
        if xlsx.exists():
            st.download_button("📊  Download Excel Report",
                data=xlsx.read_bytes(), 
                file_name=f"{country}_Country_Risk_Profile_{q_part}_{y_part}.xlsx", 
                use_container_width=True,
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
        else:
            st.warning(f"File not found: {xlsx.name}")

    with d2:
        if docx.exists():
            st.download_button("📝  Download Word Report",
                data=docx.read_bytes(), 
                file_name=f"{country}_Country_Snapshot_{q_part}_{y_part}.docx",
                use_container_width=True,
                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
        else:
            st.warning(f"File not found: {docx.name}")


# ── Footer ───────────────────────────────────────────────────
st.divider()

with st.expander("ℹ️ How this tool works"):
    st.markdown("""
**Data Sources:** World Bank API, IMF World Economic Outlook, European Central Bank, Trading Economics

**AI Model:** Configurable — supports DeepSeek, GPT-4o, Mistral, Gemini

**Process:** Collect data → AI analysis → Human review → Export reports

**Privacy:** No data is stored on the server. API keys exist only in your browser session and are never saved.
    """)

st.caption(f"Macro Data Tool v0.1 — {country} — {quarter}")