from duckduckgo_search import DDGS
import json
import time

def save_datas(file_path, datas, json_flag=True, all_flag=False, with_indent=False, mode='w'):
    """保存文本文件"""
    with open(file_path, mode, encoding='utf-8') as f:
        if all_flag:
            if json_flag:
                f.write(json.dumps(datas, ensure_ascii=False, indent= 4 if with_indent else None))
            else:
                f.write(''.join(datas))
        else:
            for data in datas:
                if json_flag:
                    f.write(json.dumps(data, ensure_ascii=False) + '\n') 
                else:
                    f.write(data + '\n')


def search_companies(companies):
    results = []
    for company in companies:
        if '公司' in company:
            results.append({
                'company': company,
                'search_results': 'company'
            })
            continue
        text = f'{company} 公司名全稱'
        
        search_results = None
        while search_results is None:
            try:
                search_results = DDGS().text(text, max_results=10)
                if search_results: break
            except Exception as e:
                print('sleep 2s')
                time.sleep(2)
                continue
        results.append({
            'company': company,
            'search_results': search_results
        })
        time.sleep(2)
    save_datas('data/公司簡稱檢索結果.json', results)


def get_datas(file_path, json_flag=True, all_flag=False, mode='r'):
    """讀取文本文件"""
    results = []
    
    with open(file_path, mode, encoding='utf-8') as f:
        for line in f.readlines():
            if json_flag:
                results.append(json.loads(line))
            else:
                results.append(line.strip())
        if all_flag:
            if json_flag:
                return json.loads(''.join(results))
            else:
                return '\n'.join(results)
        return results
    

if __name__ == '__main__':
    search_companies(get_datas('data/公司簡稱列表.txt', json_flag=False))

from duckduckgo_search import DDGS
import json
import time


PROMPT = """你是一個助手，你的任務是基于輸入的公司名簡稱以及搜索結果，分析并提取出公司名簡稱對應的公司名全稱。要求如下：
1. "簡稱"是公司名簡稱，"搜索結果"是基于互聯網的搜索后的資源，需要根據"簡稱"和"搜索結果"進行分析，并輸出公司全稱，如果無法確認，請返回"無"；
2. 如果檢索結果不包含公司全稱，請基于你所學習的知識可以進一步判斷; 
3. 輸出結果只包含公司名的全稱信息，且只能包含一個，不需要輸出解釋信息;
4. 輸入的公司名簡稱均是科技領域的知名公司，這點請注意；

示例：
簡稱: 京東
搜索結果：
1. 京東（中國1998年創立的自營式電商企業）_百度百科\n京東（股票代碼：jd），中國自營式電商企業，創始人劉強東初期擔任京東集團董事局主席兼首席執行官，2021年9月，徐雷獲任集團總裁。京東旗下設有京東商城、京東金融、拍拍網、京東智能、o2o及海外事業部等。1998年6月18日，劉強東在 中關村成立京東公司。
2. 京東集團 - 維基百科，自由的百科全書\n東集團. 京東集團 （NASDAQ： JD 、 港交所： 9618 、 港交所： 89618 （人民幣結算）），前稱 360buy 和 京東商城，由劉強東于1998年6月18日創立，是一家總部位于 北京 的 中國 電子商務公司，主要為 B2C 模式的購物網站 。. 2014年，京東集團在 美國 納斯達克證券交易 ...
3. 京東集團股份有限公司 - 愛企查\n簡介： 京東集團股份有限公司（JD.com, Inc.）于2006年11月6日在在英屬維爾京群島注冊成立的公司，通過中國境內的子公司和VIE開展經營活動，公司總部位于北京。. 京東是專業的綜合性網上購物商城，是中國B2C市場最大的3C網購專業平臺，是中國電子商務領域最受 ...
輸出: 京東集團股份有限公司

現在，請按照要求完成：
簡稱: {company_name}
搜索結果: {search_results}
輸出: 
"""

def save_datas(file_path, datas, json_flag=True, all_flag=False, with_indent=False, mode='w'):
    """保存文本文件"""
    with open(file_path, mode, encoding='utf-8') as f:
        if all_flag:
            if json_flag:
                f.write(json.dumps(datas, ensure_ascii=False, indent= 4 if with_indent else None))
            else:
                f.write(''.join(datas))
        else:
            for data in datas:
                if json_flag:
                    f.write(json.dumps(data, ensure_ascii=False) + '\n') 
                else:
                    f.write(data + '\n')

def get_datas(file_path, json_flag=True, all_flag=False, mode='r'):
    """讀取文本文件"""
    results = []
    
    with open(file_path, mode, encoding='utf-8') as f:
        for line in f.readlines():
            if json_flag:
                results.append(json.loads(line))
            else:
                results.append(line.strip())
        if all_flag:
            if json_flag:
                return json.loads(''.join(results))
            else:
                return '\n'.join(results)
        return results
    


def get_company_full_names():
    results = []
    for ele in get_datas('data/公司簡稱檢索結果.json'):
        company_name = ele['company']
        search_results = ele['search_results']
        if isinstance(search_results, str):
            results.append(f'{company_name}\t{company_name}')
            continue
        
        prompt = PROMPT.format(company_name=company_name, search_results=search_results)
        result = ''
        while result == '':
            try:
                result = DDGS().chat(prompt, model='gpt-4o-mini')
                if result.strip(): break
            except Exception as e:
                time.sleep(2)
                continue
        results.append(f'{company_name}\t{result}')
    save_datas('data/公司全稱提取結果.txt', results, json_flag=False)

if __name__ == '__main__':
    get_company_full_names()