import argparse
import asyncio
import sys

from fsm_crawl.browser.manager import BrowserManager
from fsm_crawl.browser.engine.playwright_engine import PlaywrightEngine
from fsm_crawl.experiment.normal_crawl import NormalCrawlExpirement
from fsm_crawl.experiment.normal_crawl_with_cookies import NormalCrawlWithCookiesExpirement
from fsm_crawl.experiment.explorative_crawl import ExplorativeCrawlExpirement
from fsm_crawl.utils.data_paths import get_url_file_path


EXPERIMENTS = {
    "normal": NormalCrawlExpirement,
    "normal_with_cookies": NormalCrawlWithCookiesExpirement,
    "explorative": ExplorativeCrawlExpirement,
}

# Default paths using bundled data if available
DEFAULT_PATH = get_url_file_path("first_1000.csv")
DEFAULT_PATH2 = get_url_file_path("first_2000.csv")


def parse_args(args=sys.argv[1:]):
    p = argparse.ArgumentParser(description="Run FSM web crawler experiments")
    
    p.add_argument("--shard-index", type=int, default=0,
                   help="Shard ID (for distributed runs)")
    
    p.add_argument("--shard-count", type=int, default=1,
                   help="Total number of shards (for distributed runs)")

    p.add_argument("--experiment", "-e", choices=EXPERIMENTS.keys(), default="normal",
                   help="Which experiment to run")
    p.add_argument("--path", "-p", default=DEFAULT_PATH,
                   help=f"Path to input CSV for URL manager (default: {DEFAULT_PATH})")
    p.add_argument("--prefix", default="run",
                   help="Filename prefix for output logs")
    p.add_argument("--two-runs", action="store_true", help="Run two experiments sequentially (first and second 1k).", default=True)
    p.add_argument("--path2", default=DEFAULT_PATH2,
                   help="Path to second CSV when --two-runs is used (default: blocking/first_2000.csv)")
    p.add_argument("--prefix2", default="run_second",
                   help="Filename prefix for second run; defaults to <prefix>_second")
    p.add_argument("--headless", action="store_true", help="Run browser in headless mode")
    p.add_argument("--engine", choices=["playwright"], default="playwright",
                   help="Browser engine to use (only 'playwright' implemented)")
    p.add_argument("--num-tabs", type=int, default=10,
                   help="Number of parallel tabs to use for crawling")
    return p.parse_args(args=args)


def _get_engine(args):
    if args.engine == "playwright":
        engine = PlaywrightEngine(headless=args.headless)
    else:
        raise RuntimeError(f"Unsupported engine: {args.engine}")
    return engine

async def _run_experiment(args):
    print(args)
    ExpClass = EXPERIMENTS[args.experiment]
    exp1 = ExpClass(args.path, args.prefix, headless=args.headless, shard_index=args.shard_index, shard_count=args.shard_count, num_tabs=args.num_tabs)

    if args.prefix2 is None:
        args.prefix2 = f"{args.prefix}_second"

    # exp2 = None
    # if args.two_runs:
    #     exp2 = ExpClass(args.path2, args.prefix2, headless=args.headless)
    engine = _get_engine(args)
    async with BrowserManager(engine) as mgr:
        await mgr.start_experiment(exp1)
    # if exp2 is not None:
    #     engine = _get_engine(args)
    #     async with BrowserManager(engine) as mgr:
    #         await mgr.start_experiment(exp2)


def main():
    args = parse_args()
    try:
        asyncio.run(_run_experiment(args))
    except KeyboardInterrupt:
        print("Interrupted by user", file=sys.stderr)
        return 130
    except Exception as e:
        print(f"Experiment failed: {e}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
