import os
import pprint
import asyncio
import json
import csv
import argparse
import uvicorn
from scanscan.util import debug_print, colorize, calculate_jaccard_similarity, get_short_drug_name, create_run_name, write_to_google_sheet, display_image_to_terminal
from scanscan.ocr import DEFAULT_OCR_ENGINES

lock = asyncio.Lock()

# Usage:
#
# ScanForDrugs = LLMDocumentParser(
#     f_ground_truth=load_parsed_drugs_ground_truth,
#     f_image_to_text=perform_ocr,
#     f_process=compare_drugs_to_ground_truth,
#     f_score_and_log=score_and_log,
#     config = {
#         'system_prompt': "You are a helpful assistant that can parse drugs from a text.",
#         'temperature': 0.0
#     }
# )
# ScanForDrugs.handle_args()


# TODO: this is quite ugly at the moment: we inject 'image_to_text' which we then pass to 'f_process',
# because it needs to use it several times (for multiple OCR engines)
class LLMDocumentParser:
    def __init__(self, *, f_ground_truth: callable, f_image_to_text: callable, f_process: callable, f_score_and_log: callable, config: dict):
        self.f_ground_truth = f_ground_truth
        self.f_image_to_text = f_image_to_text
        self.f_process = f_process
        self.f_score_and_log = f_score_and_log
        self.config = config


    # TODO: change image_path to directory and filename
    async def parse_single(self, image_path: str):
        directory = os.path.dirname(image_path)
        filename = os.path.basename(image_path)

        model = self.config["model"]
        ocr_engines = self.config["ocr_engines"]
        prompt_file = self.config["prompt_file"]
        country = self.config["country"]

        if filename not in self.f_ground_truth(country):
            return f"Filename {filename} not found in ground truth."

        ground_truth = self.f_ground_truth(country)[filename]

        results = await self.f_process(directory, filename, country, model, prompt_file, ocr_engines, self.config, self.f_image_to_text)
        results["_ground_truth"] = [d['registered_name'] for d in ground_truth]

        return results

    async def evaluate_all(self, directory: str, f_ground_truth: callable, f_process: callable, f_eval_log_entry: callable):
        run_name = create_run_name(self.config["country"])
        evaluation_log = []
        model = self.config["model"]
        prompt_file = self.config["prompt_file"]
        ocr_engines = self.config["ocr_engines"]
        include_files = self.config["include_files"]
        exclude_files = self.config["exclude_files"]
        only_ground_truth = self.config["only_ground_truth"]
        comment = self.config["comment"]
        country = self.config["country"]

        # process using injected functions
        async def process_file(model, prompt_file, file, country):
            # ground truth
            print(colorize(f'Processing {file}', 'yellow'))
            ground_truth = f_ground_truth(country)[file]

            # process concurrently
            results = await f_process(directory, file, country, model, prompt_file, ocr_engines, self.config, self.f_image_to_text)
            pprint.pprint(results, width=160)

            # build log entry
            evaluation_log_entry = f_eval_log_entry(results, file, run_name, comment, model, prompt_file, country, ground_truth)
            async with lock:
                if evaluation_log_entry:
                    evaluation_log.append(evaluation_log_entry)

        # filter files
        files = [file for file in os.listdir(directory) if file.endswith(('.png', '.jpg', '.jpeg'))]
        files = sorted(files, key=str.lower)
        if only_ground_truth:
            files = [file for file in files if file in f_ground_truth(country)]
        if include_files:
            files = [file for file in files if any(pat in file for pat in include_files.split(','))]
        if exclude_files:
            files = [file for file in files if not any(pat in file for pat in exclude_files.split(','))]
        print(f'files: {files}')

        # process all
        if self.config["use_concurrency"]:
            print(f'USE_CONCURRENCY: True, tasks: {len(files)}')
            tasks = []
            for file in files:
                tasks.append(asyncio.create_task(process_file(model, prompt_file, file, country)))
            await asyncio.gather(*tasks)
        else:
            print(f'USE_CONCURRENCY: False, files: {len(files)}')
            for file in files:
                await process_file(model, prompt_file, file, country)

        # TODO: writing the log should probably also be an injectable function (to be able to write to DB etc.)
        # write to file
        with open(f'logs/evaluation_log_{run_name}.json', 'w') as file:
            json.dump(evaluation_log, file, indent=4)

        # write to csv
        try:
            with open(f'logs/evaluation_log_{run_name}.csv', 'w', newline='') as file:
                writer = csv.DictWriter(file, fieldnames=evaluation_log[0].keys())
                writer.writeheader()
                writer.writerows(evaluation_log)
        except Exception as e:
            breakpoint()
            print(f'Error writing to csv: {e}')

        # write to google sheet
        write_to_google_sheet(evaluation_log, 'evaluation-log-sheet-credentials.json', os.getenv('EVALUATION_LOG_SHEET_ID'), 'log')

    def handle_args(self):
        parser = argparse.ArgumentParser(description="Parse images and extract data")
        parser.add_argument("--parse-single", type=str, help="Path to the image file to be scanned")
        # parser.add_argument("--parse-all", type=str, help="Path to the image file to be scanned")
        parser.add_argument("--evaluate-all", type=str, help="Path to the directory to be evaluated")
        parser.add_argument("--country", type=str, default="SI", help="Country to use for ground truth comparison")
        parser.add_argument("--ocr-engines", type=str, default=DEFAULT_OCR_ENGINES, help="OCR engine to use for comparison (azure, tesseract, paddleocr, ios, android)")
        parser.add_argument("--model", type=str, default="gpt-4o", help="Model to use for LLM parsing (gpt-4o, gpt-4o-mini, llama3-8b-8192)")
        parser.add_argument("--prompt-file", type=str, default="prompt3", help="Prompt file to use for LLM parsing")
        parser.add_argument("--server", action="store_true", help="Run the uvicorn server")
        parser.add_argument("--comment", type=str, help="Additional comments on the run, e.g. why you runnign it or something you changed that's not reflected in the params")
        parser.add_argument("--include-files", type=str, help="Include files by pattern (comma separated)")
        parser.add_argument("--exclude-files", type=str, help="Exclude files by pattern (comma separated)")
        parser.add_argument("--only-ground-truth", action="store_true", help="Only evaluate files that are in the ground truth")
        parser.add_argument("--use-concurrency", action="store_true", help="Use concurrency for evaluation")

        args = parser.parse_args()

        self.config.update({
            "use_concurrency": args.use_concurrency,
            "model": args.model,
            "comment": args.comment,
            "prompt_file": args.prompt_file,
            "country": args.country,
            "ocr_engines": args.ocr_engines,
            "include_files": args.include_files,
            "exclude_files": args.exclude_files,
            "only_ground_truth": args.only_ground_truth,
        })

        if args.server:
            uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)
        elif args.parse_single:
            display_image_to_terminal(args.parse_single)

            result = asyncio.run(
                self.parse_single(
                    args.parse_single,
                )
            )
            print(result)
        elif args.evaluate_all:
            asyncio.run(
                self.evaluate_all(
                    directory=args.evaluate_all,
                    f_ground_truth=self.f_ground_truth,
                    f_process=self.f_process,
                    f_eval_log_entry=self.f_score_and_log,
                )
            )

        else:
            parser.print_help()
