
import json
import time
import csv
import argparse
import logging
from functools import partial
import config
import sqlite_utils as db_tools
import iafd_scraper as scraper
import utils 

config.setup_logging()

debug = False
force = False

# 按星座获取演员列表，无翻页
def fetch_performers_by_astro():    
    for astro in scraper.astro_list:
        url = scraper.astr_base_url + astro
        logging.info(f"Fetching data for {astro}, url {url} ...")

        soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="div", identifier="astro", attr_type="id"))
        if soup:
            list_data, next_url = scraper.parse_page_astro(soup, astro)
            if list_data:
                for row in list_data :
                    # 写入演员数据表
                    perfomer_id = db_tools.insert_performer_index(name=row['person'], href=row.get('href', '').lower(), from_astro_list=1)
                    if perfomer_id:
                        logging.debug(f'insert performer index to db. performer_id:{perfomer_id}, name: {row['person']}, href:{row['href']}')
                    else:
                        logging.warning(f'insert performer index failed. name: {row['person']}, href:{row['href']}')

            else:
                logging.warning(f'fetch astro error. {url} ...')
        elif status_code  and status_code == 404:
            logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
        else:
            logging.warning(f'fetch astro error. {url} ...')

        # 调试添加break
        if debug:
            break


# 按生日获取演员列表，无翻页
def fetch_performers_by_birth():    
    for month in range(1, 13):  # 遍历1到12月
        for day in range(1, 32):  # 遍历1到31天
            url = scraper.birth_base_url.format(month=month, day=day)
            logging.info(f"Fetching data for birth, url {url}")
            soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="div", identifier="col-sm-12 col-lg-9", attr_type="class"))
            if soup:
                list_data, next_url = scraper.parse_page_birth(soup, month, day)
                if list_data:
                    for row in list_data :
                        # 写入演员数据表
                        perfomer_id = db_tools.insert_performer_index(name=row['person'], href=row.get('href', '').lower(), from_birth_list=1)
                        if perfomer_id:
                            logging.debug(f'insert performer index to db. performer_id:{perfomer_id}, name: {row['person']}, href:{row['href']}')
                        else:
                            logging.warning(f'insert performer index failed. name: {row['person']}, href:{row['href']}')
                else:
                    logging.warning(f'fetch astro error. {url} ...')
            elif status_code  and status_code == 404:
                logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
            else:
                logging.warning(f'fetch astro error. {url} ...')

            # 调试添加break
            if debug:
                return True

# 更新人种列表
def fetch_ethic_list():
    url = scraper.ethnic_list_url
    logging.info(f"Fetching data for performer's ethnic list, url {url} ...")
    soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="select", identifier="ethnicity1", attr_type="id"))
    if soup:
        list_data = scraper.parse_page_ethnic_list(soup, url)
        if list_data:
            for row in list_data :
                dist_id = db_tools.insert_or_update_ethnic({'name': row['name'], 'href': row.get('href', '')})
                if dist_id:
                    logging.debug(f'insert one record into ethnic table. id:{dist_id}, name: {row['name']}, href:{row.get('href', '')}')
        else:
            logging.warning(f'fetch ethnic error. {url} ...')
    elif status_code  and status_code == 404:
        logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
    else:
        logging.warning(f'fetch page error. {url} ...')


# 按人种获取演员列表，有翻页
def fetch_performers_by_ethnic():    
    # 先刷新列表
    fetch_ethic_list()

    ethnic_list = db_tools.query_ethnic_hrefs()
    for row in ethnic_list:
        url = row['href']
        ethnic = row['name']
        next_url = url

        while next_url:
            logging.info(f"Fetching data for {ethnic}, url {url} ...")
            soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="div", identifier="row headshotrow", attr_type="class"), 
                        parser="lxml", preprocessor=scraper.preprocess_html)
            if soup:
                list_data, next_url = scraper.parse_page_ethnic(soup, ethnic)
                if list_data:
                    for row in list_data :
                        # 写入演员数据表
                        perfomer_id = db_tools.insert_performer_index(name=row['person'], href=row.get('href', '').lower(), from_ethnic_list=1)                      
                        if perfomer_id:
                            logging.debug(f'insert performer index to db. performer_id:{perfomer_id}, name: {row['person']}, href:{row['href']}')
                        else:
                            logging.warning(f'insert performer index failed. name: {row['person']}, href:{row['href']}')
                else:
                    logging.warning(f'fetch astro error. {url} ...')
            elif status_code  and status_code == 404:
                logging.warning(f'fetch page error. httpcode: {status_code}, url: {next_url}, Skiping...')
                break
            else:
                logging.warning(f'fetch astro error. {url} ...')

            # 调试添加break
            if debug:
                return True

# 获取distributors列表
def fetch_distributors_list():
    url = scraper.distributors_list_url
    logging.info(f"Fetching data for distributors list, url {url} ...")
    soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="select", identifier="Distrib", attr_type="name"))
    if soup:
        list_data, next_url = scraper.parse_page_dist_stu_list(soup, "Distrib")
        if list_data:
            for row in list_data :
                dis_url = scraper.distributors_base_url + row['href']
                dist_id = db_tools.insert_or_update_distributor({'name': row['name'], 'href': dis_url})
                if dist_id:
                    logging.debug(f'insert one record into distributors table. id:{dist_id}, name: {row['name']}, href:{dis_url}')
        else:
            logging.warning(f'fetch astro error. {url} ...')
    elif status_code  and status_code == 404:
        logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
    else:
        logging.warning(f'fetch astro error. {url} ...')

# 获取studios列表
def fetch_studios_list():
    url = scraper.studios_list_url
    logging.info(f"Fetching data for studios list, url {url} ...")
    soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="select", identifier="Studio", attr_type="name"))
    if soup:
        list_data, next_url = scraper.parse_page_dist_stu_list(soup, "Studio")
        if list_data:
            for row in list_data :
                stu_url = scraper.studios_base_url + row['href']
                stu_id = db_tools.insert_or_update_studio({'name': row['name'], 'href': stu_url})
                if stu_id:
                    logging.debug(f'insert one record into studios table. id:{stu_id}, name: {row['name']}, href:{stu_url}')
        else:
            logging.warning(f'fetch astro error. {url} ...')
    elif status_code  and status_code == 404:
        logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
    else:
        logging.warning(f'fetch astro error. {url} ...')


# 更新distributors列表中的影片信息
def fetch_movies_by_dist():
    # 先刷新一下列表
    fetch_distributors_list()

    url_list = db_tools.query_studio_hrefs()
    if debug:
        url_list = db_tools.query_distributor_hrefs(name='vixen.com')
    for url in url_list:
        logging.info(f"Fetching data for distributor url {url} ...")
        soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="table", identifier="distable", attr_type="id"))
        if soup:
            list_data, next_url = scraper.parse_page_dist_stu(soup, 'distable')
            if list_data:                
                for movie in list_data:
                    tmp_id = db_tools.insert_movie_index(title=movie['title'], href=movie['href'], release_year=utils.to_number(movie['year']), from_dist_list=1)
                    if tmp_id:
                        logging.debug(f'insert one movie index to db. movie_id: {tmp_id}, title: {movie['title']}, href: {movie['href']}')
                    else:
                        logging.warning(f'insert movie index failed. title: {movie['title']}, href: {movie['href']}')
            else :
                logging.warning(f'parse_page_movie error. url: {url}')
        elif status_code  and status_code == 404:
            logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
        else:
            logging.warning(f'fetching page error. {url}')
        # 调试增加brak
        if debug:
            break

# 更新distributors列表中的影片信息
def fetch_movies_by_stu():
    # 先刷新一下列表
    fetch_studios_list()

    url_list = db_tools.query_studio_hrefs()
    if debug:
        url_list = db_tools.query_studio_hrefs(name='vixen.com')
    for url in url_list:
        logging.info(f"Fetching data for studio url {url} ...")
        soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="table", identifier="studio", attr_type="id"))
        if soup:
            list_data, next_url = scraper.parse_page_dist_stu(soup, 'studio')
            if list_data:                
                for movie in list_data:
                    tmp_id = db_tools.insert_movie_index(title=movie['title'], href=movie['href'], release_year=utils.to_number(movie['year']), from_stu_list=1)
                    if tmp_id:
                        logging.debug(f'insert one movie index to db. movie_id: {tmp_id}, title: {movie['title']}, href: {movie['href']}')
                    else:
                        logging.warning(f'insert movie index failed. title: {movie['title']}, href: {movie['href']}')
            else :
                logging.warning(f'parse_page_movie error. url: {url}')
        elif status_code  and status_code == 404:
            logging.warning(f'fetch page error. httpcode: {status_code}, url: {url}, Skiping...')
        else:
            logging.warning(f'fetching page error. {url}')
        # 调试增加brak
        if debug:
            break

# 更新演员信息，单次循环
def fetch_performers_detail_once(perfomers_list):
    last_performer_id = 0
    for performer in perfomers_list:
        url = performer['href']
        person = performer['name']
        logging.info(f"Fetching data for performer ({person}), url {url} ...")
        soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="div", identifier="headshot", attr_type="id"))
        if soup:
            data = scraper.parse_page_performer(soup)
            if data:
                performer_id = db_tools.insert_or_update_performer({
                    'href': url,
                    'person': person,
                    **data
                })
                if performer_id:
                    logging.debug(f'insert one person, id: {performer_id}, person: ({person}), url: {url}')
                    last_performer_id = performer_id
                else:
                    logging.warning(f'insert person: ({person}) {url} failed.')

                # 写入到本地json文件
                utils.write_person_json(person, url, {
                    'href': url,
                    'person': person,
                    **data
                })
            else:
                logging.warning(f'parse_page_performer error. person: ({person}), url: {url}')
        elif status_code  and status_code == 404:
            performer_id = db_tools.insert_or_update_performer_404(name=person, href=url)
            logging.warning(f'404 page. id: {performer_id}, name: {person}, url: {url}, Skiping...')
        else:
            logging.warning(f'fetch_page error. person: ({person}), url: {url}')
        time.sleep(1)
    return last_performer_id

# 更新演员信息
def fetch_performers_detail():
    limit_count = 5 if debug else 100
    perfomers_list = []

    # 获取新演员的列表
    while True:
        perfomers_list = db_tools.query_performer_hrefs(is_full_data=0, limit=limit_count)
        if len(perfomers_list) < 1:
            logging.info(f'all new performers fetched. ')
            break
        last_perfomer_id = fetch_performers_detail_once(perfomers_list)
        logging.info(f'insert {len(perfomers_list)} person. last performer id: {last_perfomer_id}')
        if debug:
            break

    # 获取待更新的演员的列表
    while True:
        perfomers_list = db_tools.get_performers_needed_update(limit=limit_count)
        if len(perfomers_list) < 1:
            logging.info(f'all existed performers updated. ')
            break
        last_perfomer_id = fetch_performers_detail_once(perfomers_list)
        logging.info(f'insert {len(perfomers_list)} person. last performer id: {last_perfomer_id}')
        if debug:
            break

# 更新影片信息
def fetch_movies_detail():
    limit_count = 10 if debug else 100
    movies_list = []
    while True:
        movies_list = db_tools.query_movie_hrefs(is_full_data=0, limit=limit_count)
        if len(movies_list) < 1:
            logging.info(f'all movies fetched.')
            break
        last_movie_id = 0
        succ_count = 0
        for movie in movies_list:
            url = movie['href']
            title = movie['title']
            logging.debug(f"Fetching data for movie ({title}), url {url} ...")
            soup, status_code = scraper.fetch_page(url, partial(scraper.generic_validator, tag="div", identifier="col-xs-12 col-sm-3", attr_type="class"))
            if soup:
                movie_data = scraper.parse_page_movie(soup, url, title)
                if movie_data :
                    # 修复url不规范的问题
                    if movie_data['DistributorHref']:
                        movie_data['DistributorHref'] = utils.dist_stu_href_rewrite(movie_data['DistributorHref'].lower())
                    if movie_data['StudioHref']:
                        movie_data['StudioHref'] = utils.dist_stu_href_rewrite(movie_data['StudioHref'].lower())
                    movie_id = db_tools.insert_or_update_movie(movie_data)
                    if movie_id:
                        logging.debug(f'insert one movie, id: {movie_id}, title: ({title}) url: {url}')
                        last_movie_id = movie_id
                        succ_count += 1
                    else:
                        logging.warning(f'insert movie {url} failed.')

                    # 写入到本地json文件
                    utils.write_movie_json(url, movie_data)
                else:
                    logging.warning(f'parse_page_movie error. url: {url}')
            elif status_code  and status_code == 404:
                # 标记为已处理
                movie_id = db_tools.insert_or_update_movie_404(title=title, href=url)
                logging.warning(f'404 page. id: {movie_id}, title: ({title}), url: {url}, Skiping...')
            else:
                logging.warning(f'fetch_page error. url: {url}')
            time.sleep(1)
        logging.info(f'total request: {len(movies_list)}, succ: {succ_count}. last movie id: {last_movie_id}')
        # 调试增加break
        if debug:
            return True


# 建立缩写到函数的映射
function_map = {
    "astro": fetch_performers_by_astro,
    "birth": fetch_performers_by_birth,
    "ethnic": fetch_performers_by_ethnic,
    "dist" : fetch_movies_by_dist,
    "stu" : fetch_movies_by_stu,
    "performers": fetch_performers_detail,
    "movies" : fetch_movies_detail,
}   

# 主函数
def main(cmd, args_debug, args_force):
    global debug
    debug = args_debug

    global force
    force = args_force

    # 开启任务
    task_id = db_tools.insert_task_log()
    if task_id is None:
        logging.warning(f'insert task log error.')
        return None
    
    logging.info(f'running task. id: {task_id}, debug: {debug}, force: {force}, cmd: {cmd}')

    # 执行指定的函数
    if cmd:
        function_names = args.cmd.split(",")  # 拆分输入
        for short_name in function_names:
            func = function_map.get(short_name.strip())  # 从映射中获取对应的函数
            if callable(func):
                db_tools.update_task_log(task_id, task_status=f'Running {func}')
                func()
            else:
                print(f"Warning: {short_name} is not a valid function shortcut.")
    else: # 全量执行
        for name, func in function_map.items():
            if callable(func):
                db_tools.update_task_log(task_id, task_status=f'Running {func}')
                func()
            else:
                print(f"Warning: {name} is not a valid function shortcut.")

    logging.info(f'all process completed!')
    db_tools.finalize_task_log(task_id)

    # TODO:
    # 1, movies 更新之后，要给相应的 performers 表打个 is_full_data = 0, 然后刷新获取
    # 2, distributors 和 studios 对movie列表的互相检验
    # 3, 数据不规范问题，可以先手动导入所有 performers 和 movies ，然后用本程序增量获取新的

if __name__ == "__main__":
    # 命令行参数处理
    keys_str = ",".join(function_map.keys())

    parser = argparse.ArgumentParser(description='fetch iafd data.')
    parser.add_argument("--cmd", type=str, help=f"Comma-separated list of function shortcuts: {keys_str}")
    parser.add_argument('--debug', action='store_true', help='Enable debug mode (limit records)')
    parser.add_argument('--force', action='store_true', help='force update (true for rewrite all)')
    args = parser.parse_args()
    
    main(args.cmd, args.debug, args.force)
