segment.py

import os
import json
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.messages import SystemMessage, HumanMessage
from tavily import TavilyClient

load_dotenv()

llm = ChatGroq(
    model="openai/gpt-oss-120b",
    temperature=0,
    api_key=os.getenv("api_key"),
)

tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))


# ── Step 1: LLM generates targeted search keywords ──────────────────────────
def get_search_keywords(company_name: str) -> list[str]:
    messages = [
        SystemMessage(content="""
You are an equity research analyst.

Given a company name, generate 3–5 targeted search queries to find information about:
- Business segments and divisions
- Revenue drivers and sources
- Key products/services breakdown

Return ONLY a JSON array of search query strings. No explanation.

Example output:
["NVIDIA business segments revenue 2024", "NVIDIA data center gaming revenue breakdown", "NVIDIA annual report segment analysis"]
"""),
        HumanMessage(content=company_name),
    ]
    response = llm.invoke(messages)

    # Parse the JSON array from LLM response
    content = response.content.strip()
    # Strip markdown code fences if present
    if content.startswith("```"):
        content = content.split("```")[1]
        if content.startswith("json"):
            content = content[4:]
    keywords = json.loads(content.strip())
    return keywords


# ── Step 2: Search Tavily ────────────────────────────────────────────────────

BLOCKED_DOMAINS = [
    "cliffsnotes.com", "studocu.com", "coursehero.com",
    "chegg.com", "quizlet.com", "wikipedia.org",
    "reddit.com", "quora.com",
    # Add these
    "stocklight.com", "riskintelligenceservice.com",
    "last10k.com", "wisesheets.io", "macrotrends.net",
]

def search_tavily(queries: list[str], max_words: int = 3000) -> str:
    all_results = []

    for query in queries:
        results = tavily.search(
            query=query,
            search_depth="advanced",
            max_results=3,
            include_raw_content=False,
             exclude_domains=BLOCKED_DOMAINS,
        )
        for r in results.get("results", []):
            all_results.append(f"SOURCE: {r['url']}\n{r['content']}")

    combined = "\n\n---\n\n".join(all_results)
    
    # ── Truncate to stay within token limits ──
    words = combined.split()
    if len(words) > max_words:
        combined = " ".join(words[:max_words])
        print(f"  ⚠️  Context truncated to {max_words} words to fit token limit")
    
    return combined

# ── Step 3: LLM summarizes the search results ────────────────────────────────
def summarize_business_segments(company_name: str, raw_context: str) -> str:
    messages = [
        SystemMessage(content="""
You are a senior equity research analyst writing a company profile.

Using the provided search results, produce a structured analysis covering:

1. **Business Segments** — List each segment, what it does, and approximate % of revenue if available
2. **Revenue Drivers** — Key factors/products/geographies driving growth
3. **Revenue Mix Trend** — Any notable shifts in segment contribution over time

Be factual, cite approximate figures where available, and flag uncertainty.
Keep the output concise but information-dense (bullet points preferred).
"""),
        HumanMessage(content=f"""
Company: {company_name}

Search Results:
{raw_context}
"""),
    ]
    response = llm.invoke(messages)
    return response.content


# ── Main pipeline ────────────────────────────────────────────────────────────
def analyze_company(company_name: str) -> str:
    print(f"\n[1/3] Generating search keywords for: {company_name}")
    keywords = get_search_keywords(company_name)
    print(f"  Keywords: {keywords}")

    print(f"\n[2/3] Searching Tavily ({len(keywords)} queries)...")
    raw_context = search_tavily(keywords)
    print(f"  Retrieved ~{len(raw_context.split())} words of context")

    print(f"\n[3/3] Summarizing with LLM...")
    summary = summarize_business_segments(company_name, raw_context)

    return summary


# ── Run ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    company = input("Enter company name or ticker: ").strip()
    if not company:
        print("No company entered. Exiting.")
        exit()
    
    result = analyze_company(company)
    print("\n" + "="*60)
    print(f"BUSINESS SEGMENTS & REVENUE DRIVERS — {company.upper()}")
    print("="*60)
    print(result)