import os
import argparse
import requests
import re
from bs4 import BeautifulSoup

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document

from core.sec_tools import get_latest_10k_url, HEADERS

def ingest_10k(ticker: str):
    """Downloads, cleans, and embeds a 10-K into a local Chroma Vector Database."""
    ticker = ticker.upper()
    persist_directory = f"./chroma_db/{ticker}_10k"
    
    if os.path.exists(persist_directory):
        print(f"[Ingest: Vector DB for {ticker} 10-K already exists at {persist_directory}. Skipping...]")
        return
    
    print(f"\n==============================================")
    print(f"Starting Ingestion Pipeline for {ticker}")
    print(f"==============================================")

    url = get_latest_10k_url(ticker)
    if url.startswith("Error") or url.startswith("No 10-K"):
        print(f"[Error: SEC URL Fetch failed: {url}]")
        return
        
    print(f"[1/4] Downloading raw 10-K from SEC: {url}")
    response = requests.get(url, headers=HEADERS)
    response.raise_for_status()
    raw_text = response.text
    
    print(f"[2/4] Parsing HTML and isolating text payload...")
    doc_match = re.search(r'<DOCUMENT>(.*?)</DOCUMENT>', raw_text, re.DOTALL | re.IGNORECASE)
    if doc_match:
        raw_text = doc_match.group(1)
        
    soup = BeautifulSoup(raw_text, "html.parser")
    clean_text = soup.get_text(separator=" ", strip=True)
    
    print(f"[3/4] Chunking document...")
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000, 
        chunk_overlap=200
    )
    docs = [Document(page_content=clean_text, metadata={"source": url, "ticker": ticker})]
    chunks = text_splitter.split_documents(docs)
    
    print(f"[4/4] Embedding {len(chunks)} chunks into Chroma DB. (This may take a minute) ...")
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    Chroma.from_documents(
        documents=chunks, 
        embedding=embeddings, 
        persist_directory=persist_directory
    )
    
    print(f"[Success] {ticker} 10-K successfully ingested.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Ingest SEC 10-K filings into Chroma DB.")
    parser.add_argument(
        "--tickers", 
        nargs="+", 
        required=True, 
        help="List of stock tickers to ingest (e.g., --tickers AAPL MSFT TSLA)"
    )
    args = parser.parse_args()
    
    os.makedirs("./chroma_db", exist_ok=True)
    
    for t in args.tickers:
        try:
            ingest_10k(t)
        except Exception as e:
            print(f"[Error] Failed to ingest {t}: {str(e)}")