import byzerllm
from typing import Optional
from autocoder.common import AutoCoderArgs
from autocoder.common.code_auto_merge import CodeAutoMerge
from autocoder.common.code_auto_merge_diff import CodeAutoMergeDiff
from autocoder.common.code_auto_merge_strict_diff import CodeAutoMergeStrictDiff
from autocoder.common.code_auto_merge_editblock import CodeAutoMergeEditBlock
from autocoder.common.code_auto_generate import CodeAutoGenerate
from autocoder.common.code_auto_generate_diff import CodeAutoGenerateDiff
from autocoder.common.code_auto_generate_strict_diff import CodeAutoGenerateStrictDiff
from autocoder.common.code_auto_generate_editblock import CodeAutoGenerateEditBlock
from autocoder.common.v2.code_editblock_manager import CodeEditBlockManager
from autocoder.index.entry import build_index_and_filter_files
from autocoder.regexproject import RegexProject
from autocoder.utils.conversation_store import store_code_model_conversation
from autocoder.common.printer import Printer
import time
from autocoder.utils.llms import get_llm_names
from autocoder.common import SourceCodeList
from autocoder.common.global_cancel import global_cancel
from autocoder.events.event_manager_singleton import get_event_manager
from autocoder.events import event_content as EventContentCreator
from autocoder.events.event_types import EventMetadata
from loguru import logger
class ActionRegexProject:
    def __init__(
        self, args: AutoCoderArgs, llm: Optional[byzerllm.ByzerLLM] = None
    ) -> None:
        self.args = args
        self.llm = llm
        self.pp = None
        self.printer = Printer()

    def run(self):
        args = self.args
        if not args.project_type.startswith(
            "human://"
        ) and not args.project_type.startswith("regex://"):
            return False

        args = self.args
        pp = RegexProject(args=args, llm=self.llm)
        self.pp = pp
        pp.run()
        source_code_list = SourceCodeList(pp.sources)
        if self.llm:
            if args.in_code_apply:
                old_query = args.query
                args.query = (args.context or "") + "\n\n" + args.query
            source_code_list = build_index_and_filter_files(
                llm=self.llm, args=args, sources=pp.sources
            )
            if args.in_code_apply:
                args.query = old_query
        self.process_content(source_code_list)

    def process_content(self, source_code_list: SourceCodeList):
        args = self.args
        content = source_code_list.to_str()

        if args.execute and self.llm and not args.human_as_model:
            if len(content) > self.args.model_max_input_length:
                logger.warning(
                    f"Content length is {len(content)}, which is larger than the maximum input length {self.args.model_max_input_length}. chunk it..."
                )
                content = content[: self.args.model_max_input_length]

        start_time = time.time()
        
        global_cancel.check_and_raise(token=self.args.cancel_token)

        if (args.enable_auto_fix_merge or args.enable_auto_fix_lint) and args.execute and args.auto_merge=="editblock":
            code_merge_manager = CodeEditBlockManager(llm=self.llm, args=self.args,action=self)
            code_merge_manager.run(query=args.query, source_code_list=source_code_list)
            return
        
        if args.execute:
            self.printer.print_in_terminal("code_generation_start")
            
            if args.auto_merge == "diff":
                generate = CodeAutoGenerateDiff(
                    llm=self.llm, args=self.args, action=self
                )
            elif args.auto_merge == "strict_diff":
                generate = CodeAutoGenerateStrictDiff(
                    llm=self.llm, args=self.args, action=self
                )
            elif args.auto_merge == "editblock":
                generate = CodeAutoGenerateEditBlock(
                    llm=self.llm, args=self.args, action=self
                )
            else:
                generate = CodeAutoGenerate(llm=self.llm, args=self.args, action=self)
            
            generate_result = generate.single_round_run(
                query=args.query, source_code_list=source_code_list
            )

            elapsed_time = time.time() - start_time
            speed = generate_result.metadata.get('generated_tokens_count', 0) / elapsed_time if elapsed_time > 0 else 0
            model_names = ",".join(get_llm_names(self.llm))
            input_tokens_cost = generate_result.metadata.get('input_tokens_cost', 0)
            generated_tokens_cost = generate_result.metadata.get('generated_tokens_cost', 0)
            self.printer.print_in_terminal(
                "code_generation_complete",
                duration=elapsed_time,
                input_tokens=generate_result.metadata.get('input_tokens_count', 0),
                output_tokens=generate_result.metadata.get('generated_tokens_count', 0),
                input_cost=input_tokens_cost,
                output_cost=generated_tokens_cost,
                speed=round(speed, 2),
                model_names=model_names,
                sampling_count=len(generate_result.contents)
            )

            get_event_manager(self.args.event_file).write_result(
                EventContentCreator.create_result(content=EventContentCreator.ResultTokenStatContent(
                    model_name=model_names,
                    elapsed_time=elapsed_time,
                    input_tokens=generate_result.metadata.get(
                        'input_tokens_count', 0),
                    output_tokens=generate_result.metadata.get(
                        'generated_tokens_count', 0),
                    input_cost=input_tokens_cost,
                    output_cost=generated_tokens_cost,
                    speed=round(speed, 2),                    
                )).to_dict(),metadata=EventMetadata(
                    action_file=self.args.file
                ).to_dict())

            global_cancel.check_and_raise(token=self.args.cancel_token)
            
            merge_result = None
            if args.execute and args.auto_merge:
                self.printer.print_in_terminal("code_merge_start")
                if args.auto_merge == "diff":
                    code_merge = CodeAutoMergeDiff(llm=self.llm, args=self.args)
                    merge_result = code_merge.merge_code(generate_result=generate_result)
                elif args.auto_merge == "strict_diff":
                    code_merge = CodeAutoMergeStrictDiff(llm=self.llm, args=self.args)
                    merge_result = code_merge.merge_code(generate_result=generate_result)
                elif args.auto_merge == "editblock":
                    code_merge = CodeAutoMergeEditBlock(llm=self.llm, args=self.args)
                    merge_result = code_merge.merge_code(generate_result=generate_result)
                else:
                    code_merge = CodeAutoMerge(llm=self.llm, args=self.args)
                    merge_result = code_merge.merge_code(generate_result=generate_result)
            
            if merge_result is not None:
                content = merge_result.contents[0]
                store_code_model_conversation(
                    args=self.args,
                    instruction=self.args.query,
                    conversations=merge_result.conversations[0],
                    model=self.llm.default_model_name,
                )
            else:
                content = generate_result.contents[0]
                store_code_model_conversation(
                    args=self.args,
                    instruction=self.args.query,
                    conversations=generate_result.conversations[0],
                    model=self.llm.default_model_name,
                )
            
