import requests
import yaml
from typing import List, Dict, Optional


class PlanItem:
    def __init__(self, id: str, description: str, status: str, subtasks: List['PlanItem'] = None):
        self.id = id
        self.status = status
        self.description = description
        self.subtasks = subtasks or []

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'status': self.status,
            'description': self.description,
            'subtasks': [subtask.to_dict() for subtask in self.subtasks]
        }


class CodeBlock:
    def __init__(self, path: str, tag: str, code: str):
        self.path = path
        self.tag = tag
        self.code = code

    def to_dict(self) -> dict:
        return {
            'path': self.path,
            'tag': self.tag,
            'code': self.code
        }

class CodeConstructionContext:
    def __init__(self, specification: str, plan: List[PlanItem], code_blocks: List[Optional[CodeBlock]],
                 code_index: Dict[str, List[str]],
                 compiler_error_messages: List[str], test_results: List[str]):
        self.specification = specification
        self.plan = plan
        self.code_index = code_index
        self.code_blocks = code_blocks
        self.compiler_error_messages = compiler_error_messages
        self.test_results = test_results


class CodeBlockUpdate:
    def __init__(self, path: str, tag: str, code: str):
        self.path, self.tag, self.code = path, tag, code


def parse_plan_update(data: dict) -> PlanItem:
    id = data['id']
    status = data['status']
    description = data.get('description')
    subtasks = [parse_plan_update(subtask) for subtask in data.get('subtasks', [])]

    return PlanItem(id, description, status, subtasks)


class CodeBlockRequest:
    def __init__(self, path: str, tag: str):
        self.path, self.tag = path, tag


class CodeConstructionOutput:
    def __init__(self, code_block_updates: List[CodeBlockUpdate], plan_updates: List[PlanItem],
                 code_block_requests: List[CodeBlockRequest]):
        self.code_block_updates, self.plan_updates, self.code_block_requests = code_block_updates, plan_updates, code_block_requests


class GenerationEngineState:
    code_blocks: Dict[str, List[CodeBlock]]
    plan: List[PlanItem]

    def __init__(self, specification: str):
        self.specification = specification
        self.plan = [
            PlanItem(id='start', status='todo', description='Implement a first sketch of the solution'),
            PlanItem(id='plan', status='todo', description='Formulate a more detailed plan'),
            PlanItem(id='test', status='todo', description='Create tests'),
        ]
        self.code_blocks = {}

    def get_initial_context(self) -> CodeConstructionContext:
        return CodeConstructionContext(self.specification, self.plan, [], {}, [], [])

    def _find_code_block(self, path: str, tag: str) -> Optional[CodeBlock]:
        if path not in self.code_blocks:
            return None

        for block in self.code_blocks[path]:
            if block.tag == tag:
                return block

        return None

    def _apply_code_blocks(self, code_block_updates: List[CodeBlockUpdate]):
        for update in code_block_updates:
            if update.path not in self.code_blocks:
                self.code_blocks[update.path] = []
            path_blocks = self.code_blocks[update.path]
            for i, block in enumerate(path_blocks):
                if block.tag == update.tag:
                    path_blocks[i] = CodeBlock(update.path, update.tag, update.code)
                    break
            else:
                path_blocks.append(CodeBlock(update.path, update.tag, update.code))

    def _apply_plan_updates(self, plan_updates: List[PlanItem]):
        def update_plan_item(item: PlanItem, update: PlanItem) -> PlanItem:
            if item.id == update.id:
                item.status = update.status
                if update.description:
                    item.description = update.description

                for update_subtask in update.subtasks:
                    for item_subtask in item.subtasks:
                        if item_subtask.id == update_subtask.id:
                            update_plan_item(item_subtask, update_subtask)
                            break
                    else:
                        item.subtasks.append(update_subtask)
            return item

        for update in plan_updates:
            item_found = False
            for i, item in enumerate(self.plan):
                if item.id == update.id:
                    self.plan[i] = update_plan_item(item, update)
                    item_found = True
                    break
            if not item_found:
                self.plan.append(PlanItem(update.id, update.description, update.status, update.subtasks))

    def _resolve_code_block_requests(self, code_block_requests: List[CodeBlockRequest]) -> List[CodeBlock]:
        blocks = [self._find_code_block(request.path, request.tag) for request in code_block_requests]
        # remove None values
        return [block for block in blocks if block]

    def _generate_code_index(self) -> Dict[str, List[str]]:
        return {path: [block.tag for block in blocks] for path, blocks in self.code_blocks.items()}

    def apply_step(self, output: CodeConstructionOutput) -> CodeConstructionContext:
        self._apply_code_blocks(output.code_block_updates)
        self._apply_plan_updates(output.plan_updates)
        requested_code_blocks = self._resolve_code_block_requests(output.code_block_requests)

        return CodeConstructionContext(self.specification, self.plan, requested_code_blocks,
                                       self._generate_code_index(), [], [])


def parse_output(data: dict) -> CodeConstructionOutput:
    code_block_updates = [CodeBlockUpdate(d['path'], d['tag'], d['code']) for d in data.get('code_blocks', [])]
    plan_updates = [parse_plan_update(d) for d in data.get('plan_updates', [])]
    code_block_requests = [CodeBlockRequest(d['path'], d['tag']) for d in data.get('code_block_requests', [])]

    return CodeConstructionOutput(code_block_updates, plan_updates, code_block_requests)


class CodeConstructionDriver:
    def __init__(self, api_key: str):
        self.api_key = api_key
        #openai.api_key = api_key
        self.state = None

    def initialize(self, specification: str):
        self.state = GenerationEngineState(specification)

    def generate_code(self, context: CodeConstructionContext) -> CodeConstructionOutput:
        if not self.state:
            raise Exception("Call initialize() first")

        content = yaml.dump({
            'specification': self.state.specification,
            'plan': [item.to_dict() for item in context.plan],
            'code_blocks': [block.to_dict() for block in context.code_blocks],
            'code_index': context.code_index
        })

        messages = [
            {"role": "system", "content":
                """You're a code construction AI which creates code iteratively.
In each step a context is supplied. It includes
specification, plan, code blocks requested in the previous step.
Generated response should be in YAML format as in example below

```
code_blocks: # list of code blocks to be inserted or updated
 - path: "src/main_yaml.py"
   tag: "def_main"
   code: "new code here"
plan_updates:
 - id: "start"
   status: "done"
 - id: "new task"
   status: "in_progress
   description: "description of new task"
   subtasks:
    - id: "subtask1.1"
      status: "todo"
      description: "subtask1.1 description"
code_block_requests: # list of code blocks requested in the next step
 - path: "path_to_previously_created_file.py"
   tag: "tag2"
```

code_block_requests are used to get previously written code into the context for the next step.
code_block_updates are the newly generated code, possibly updating existing code blocks.
Do not include "```" in the response, include only raw YAML data.
"""},
            {"role": "user", "content": content}
        ]

        print(content)

        headers = {"Authorization":"Basic"}
        response_1 = requests.post(
            "URL to scale AI deployment",
            json={"input": {"content": content}},
            headers=headers
        )
        print(response_1)
        response = response_1.json()["output"]

        print("Response:")
        print(response)
        print("Response ENDS")

        output_data = yaml.load(response, Loader=yaml.SafeLoader)

        return parse_output(output_data)

    def run(self, specification: str, num_steps: int):
        self.initialize(specification)
        context = self.state.get_initial_context()
        for _ in range(num_steps):
            output = self.generate_code(context)
            context = self.state.apply_step(output)


if __name__ == "__main__":
    driver = CodeConstructionDriver("openAI key actually not used")
    driver.run("Write utility similar to grep in C++", 10)