import sys
import json
from model_downloader import download_model
from ai_agent import AIAgent
import argparse

def main():
    parser = argparse.ArgumentParser(description='AI 웹 정보 수집 시스템')
    parser.add_argument('--topics', nargs='+', help='수집할 주제 목록', default=None)
    parser.add_argument('--config', default='./config.json', help='설정 파일 경로')
    parser.add_argument('--auto-topics', action='store_true', help='AI가 스스로 주제를 선정하여 조사')
    parser.add_argument('--save-path', help='데이터 저장 경로 (예: /content/drive/MyDrive/MyFolder)', default=None)

    args = parser.parse_args()

    print("AI 웹 정보 수집 시스템 시작")

    # 1. 모델 다운로드 (필요한 경우)
    print("모델 확인 중...")
    model, tokenizer = download_model(args.config)

    if model is None:
        print("모델 다운로드 실패. 프로그램을 종료합니다.")
        sys.exit(1)

    # 2. 설정 파일 로드 및 수정
    with open(args.config, 'r') as f:
        config = json.load(f)

    if args.save_path:
        config['data_storage']['drive_mount_path'] = args.save_path
        # 수정된 config 저장
        with open(args.config, 'w') as f:
            json.dump(config, f, indent=2)
        print(f"저장 경로 설정됨: {args.save_path}")

    # 3. AI 에이전트 초기화
    print("AI 에이전트 초기화 중...")
    agent = AIAgent(args.config)

    # 3. 주제 결정
    if args.auto_topics or args.topics is None:
        print("AI가 스스로 주제를 선정합니다...")
        topics = agent.generate_topics(num_topics=3)
        print(f"선정된 주제: {topics}")
    else:
        topics = args.topics

    # 4. 정보 수집 실행
    print(f"다음 주제들에 대해 정보를 수집합니다: {topics}")
    results = agent.collect_information(topics)

    # 5. 결과 출력
    print("\n=== 수집 결과 ===")
    for result in results:
        print(f"주제: {result['topic']}")
        print(f"응답: {result['response']}")
        print("-" * 50)

    # 6. 정리
    agent.close()
    print("프로그램 완료")

if __name__ == "__main__":
    main()