import streamlit as st
import faiss
import pickle
import numpy as np
import requests
import os
from sentence_transformers import SentenceTransformer

st.set_page_config(page_title="SEC 10-K Analyser", page_icon="📊", layout="wide")

st.title("📊 SEC 10-K Filing Analyser")
st.markdown("Ask questions about **Apple, Microsoft, Google, JPMorgan and Tesla** annual reports (2023-2025)")

HF_TOKEN = os.environ.get("HF_TOKEN")
API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"

@st.cache_resource
def load_components():
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    index = faiss.read_index("sec_index.faiss")
    with open("chunks.pkl", "rb") as f:
        chunks = pickle.load(f)
    return embedder, index, chunks

def query_mistral(prompt, hf_token):
    headers = {
        "Authorization": f"Bearer {hf_token}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "mistralai/Mistral-7B-Instruct-v0.2",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 400,
        "temperature": 0.1
    }
    API_URL = "https://router.huggingface.co/hf-inference/v1/chat/completions"
    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"API error: {response.status_code} - {response.text}"

def answer_question(query, embedder, index, chunks, hf_token, top_k=5):
    query_embedding = embedder.encode([query], convert_to_numpy=True)
    faiss.normalize_L2(query_embedding)
    distances, indices = index.search(query_embedding, top_k)
    
    results = []
    for i, idx in enumerate(indices[0]):
        results.append({
            "ticker": chunks[idx]["ticker"],
            "section": chunks[idx]["section"],
            "text": chunks[idx]["text"],
            "score": float(distances[0][i])
        })
    
    context = ""
    for i, r in enumerate(results):
        context += f"[Source {i+1} - {r['ticker']} {r['section']}]: {r['text']}\n\n"
    
    prompt = f"""<s>[INST] You are a financial analyst assistant specialising in SEC 10-K filings.
Answer the question using ONLY the context provided below.
Always state which company and section your answer comes from.
If the context does not contain enough information, say "I cannot confidently answer this from the available filings."
Do not make up information.

Context:
{context}

Question: {query} [/INST]"""
    
    answer = query_mistral(prompt, hf_token)
    sources = list(set([f"{r['ticker']} ({r['section']})" for r in results]))
    return answer, sources, results

with st.sidebar:
    st.header("About")
    st.markdown("""
    This tool analyses SEC 10-K filings using RAG (Retrieval Augmented Generation).
    
    **Companies covered:**
    - Apple (AAPL)
    - Microsoft (MSFT)
    - Google (GOOGL)
    - JPMorgan Chase (JPM)
    - Tesla (TSLA)
    
    **Filing years:** 2023-2025
    
    **Sections indexed:**
    - Risk Factors
    - MD&A
    """)
    
    st.header("Example Questions")
    examples = [
        "What are Apple's main supply chain risks?",
        "How does Tesla describe competition in EVs?",
        "What does Google say about AI regulation risks?",
        "How has JPMorgan managed credit risk?",
        "What are Microsoft's cybersecurity concerns?",
    ]
    for ex in examples:
        if st.button(ex, key=ex):
            st.session_state.query = ex

query = st.text_input(
    "Ask a question about the filings",
    value=st.session_state.get("query", ""),
    placeholder="e.g. What are Apple's main risk factors?"
)

if st.button("Analyse", type="primary") and query:
    if not HF_TOKEN:
        st.error("HF_TOKEN not set. Please add it in Space secrets.")
    else:
        with st.spinner("Searching filings and generating answer..."):
            embedder, index, chunks = load_components()
            answer, sources, results = answer_question(query, embedder, index, chunks, HF_TOKEN)
        
        st.markdown("### Answer")
        st.write(answer)
        
        st.markdown("### Sources Used")
        for s in sources:
            st.markdown(f"- {s}")
        
        with st.expander("View retrieved chunks"):
            for i, r in enumerate(results):
                st.markdown(f"**Chunk {i+1} [{r['ticker']} - {r['section']}] (score: {r['score']:.3f})**")
                st.text(r["text"][:300])
                st.divider()