from prompts import Prompts
from config import Config
from common_functionality import send_openai_request
import openai
import os
import yfinance as yf


def get_basic_info(sector: str, n_companies: int, required_fields: list, filter_tickers=False) -> list:
    max_num_companies = Config.OPENAI_MAX_NUM_COMPANIES
    data = []
    if n_companies > max_num_companies:
        print(f"Getting {n_companies} companies")
        ignore_companies = []
        for _ in range(n_companies // max_num_companies):
            prompt = Prompts.BASIC_INFORMATION_PROMPT.format(num_companies=max_num_companies,
                                                             sector=sector,
                                                             ignore_companies=ignore_companies,
                                                             required_fields=required_fields)
            resp = send_openai_request(prompt, max_tokens=7100, temperature=0.1, model="gpt-4")
            print(f"Ignore companies: {ignore_companies}")
            ignore_companies.extend([company['name'] for company in resp])
            if type(resp) == dict:
                data.append(resp)
            else:
                data.extend(resp)
    else:
        prompt = Prompts.BASIC_INFORMATION_PROMPT_SINGLE_USE.format(num_companies=n_companies, sector=sector,
                                                                    required_fields=required_fields)
        data.extend(send_openai_request(prompt, max_tokens=7100, temperature=0.1, model="gpt-4"))

    if filter_tickers:
        filtered_tickers = []
        for company in data:
            try:
                ticker = yf.Ticker(company['ticker'])
                if ticker.info.get("currentPrice"):
                    filtered_tickers.append(company["ticker"])
            except:
                continue
        data = [company for company in data if company['ticker'] in filtered_tickers]
    return data