# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Dict, List, Literal, Optional

from ..base import Template
from ..constant import LLMTemplateType, MLLMTemplateType
from ..register import TemplateMeta, register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Context, Prompt, findall

today = datetime.now().strftime('%Y-%m-%d')

mistral_2501_system = (
    'You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup '
    'headquartered in Paris.\n'
    f'Your knowledge base was last updated on 2023-10-01. The current date is {today}.\n\n'
    "When you're not sure about some information, you say that you don't have the information and don't "
    'make up anything.\n'
    "If the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer "
    'the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. '
    '"What are some good restaurants around me?" => "Where are you?" or "When is the next flight to Tokyo" => "'
    'Where do you travel from?")')


@dataclass
class Mistral3TemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: ['<s>'])
    prompt: Prompt = field(default_factory=lambda: ['[INST]{{QUERY}}[/INST]'])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['</s>'])
    suffix: Prompt = field(default_factory=lambda: ['</s>'])
    system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<s>[SYSTEM_PROMPT]{{SYSTEM}}[/SYSTEM_PROMPT]'])


register_template(Mistral3TemplateMeta(LLMTemplateType.mistral_2501, default_system=mistral_2501_system))


class Mistral2503Template(Template):
    placeholder_tokens = ['[IMG]']
    image_token = 10

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type == 'image'
        return ['[IMG]']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        processor = self.processor
        images = inputs.images
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        idx_list = findall(input_ids, self.image_token)
        patch_size = processor.patch_size * processor.spatial_merge_size
        if idx_list:
            image_inputs = processor.image_processor(images, patch_size=patch_size, return_tensors='pt')
            encoded['pixel_values'] = image_inputs['pixel_values'].to(self.model_info.torch_dtype)
            encoded['image_sizes'] = image_sizes = image_inputs['image_sizes']

            def _get_new_tokens(i):
                height, width = image_sizes[i]
                num_height_tokens = height // patch_size
                num_width_tokens = width // patch_size
                replace_tokens = [[processor.image_token] * num_width_tokens + [processor.image_break_token]
                                  ] * num_height_tokens
                # Flatten list
                replace_tokens = [item for sublist in replace_tokens for item in sublist]
                replace_tokens[-1] = processor.image_end_token
                replace_str = ''.join(replace_tokens)
                return processor.encode(replace_str, add_special_tokens=False)

            encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
                input_ids, labels, loss_scale, idx_list, _get_new_tokens)

        return encoded


register_template(
    Mistral3TemplateMeta(
        MLLMTemplateType.mistral_2503, default_system=mistral_2501_system, template_cls=Mistral2503Template))

devstral_small_2505_system = (  # from https://huggingface.co/mistralai/Devstral-Small-2505/blob/main/SYSTEM_PROMPT.txt
    'You are Devstral, a helpful agentic model trained by Mistral AI and using the OpenHands scaffold. '
    'You can interact with a computer to solve tasks.\n\n<ROLE>\nYour primary role is to assist users by '
    'executing commands, modifying code, and solving technical problems effectively. You should be '
    'thorough, methodical, and prioritize quality over speed.\n* If the user asks a question, like '
    '"why is X happening", don\'t try to fix the problem. Just give an answer to the question.'
    '\n</ROLE>\n\n<EFFICIENCY>\n* Each action you take is somewhat expensive. Wherever possible, '
    'combine multiple actions into a single action, e.g. combine multiple bash commands into one, using '
    'sed and grep to edit/view multiple files at once.\n* When exploring the codebase, use efficient tools '
    'like find, grep, and git commands with appropriate filters to minimize unnecessary operations.'
    '\n</EFFICIENCY>\n\n<FILE_SYSTEM_GUIDELINES>\n* When a user provides a file path, do NOT assume it\'s '
    'relative to the current working directory. First explore the file system to locate the file before '
    'working on it.\n* If asked to edit a file, edit the file directly, rather than creating a new file with '
    'a different filename.\n* For global search-and-replace operations, consider using `sed` instead of '
    'opening file editors multiple times.\n</FILE_SYSTEM_GUIDELINES>\n\n<CODE_QUALITY>\n* Write clean, '
    'efficient code with minimal comments. Avoid redundancy in comments: Do not repeat information that can '
    'be easily inferred from the code itself.\n* When implementing solutions, focus on making the minimal '
    'changes needed to solve the problem.\n* Before implementing any changes, first thoroughly understand '
    'the codebase through exploration.\n* If you are adding a lot of code to a function or file, consider '
    'splitting the function or file into smaller pieces when appropriate.\n</CODE_QUALITY>\n\n'
    '<VERSION_CONTROL>\n* When configuring git credentials, use "openhands" as the user.name and '
    '"openhands@all-hands.dev" as the user.email by default, unless explicitly instructed otherwise.'
    '\n* Exercise caution with git operations. Do NOT make potentially dangerous changes (e.g., pushing '
    'to main, deleting repositories) unless explicitly asked to do so.\n* When committing changes, use `git'
    ' status` to see all modified files, and stage all files necessary for the commit. Use `git commit -a` '
    'whenever possible.\n* Do NOT commit files that typically shouldn\'t go into version control (e.g., '
    'node_modules/, .env files, build directories, cache files, large binaries) unless explicitly '
    'instructed by the user.\n* If unsure about committing certain files, check for the presence of .'
    'gitignore files or ask the user for clarification.\n</VERSION_CONTROL>\n\n<PULL_REQUESTS>\n* When '
    'creating pull requests, create only ONE per session/issue unless explicitly instructed otherwise.\n* '
    'When working with an existing PR, update it with new commits rather than creating additional PRs for '
    'the same issue.\n* When updating a PR, preserve the original PR title and purpose, updating description '
    'only when necessary.\n</PULL_REQUESTS>\n\n<PROBLEM_SOLVING_WORKFLOW>\n1. EXPLORATION: Thoroughly '
    'explore relevant files and understand the context before proposing solutions\n2. ANALYSIS: Consider '
    'multiple approaches and select the most promising one\n3. TESTING:\n   * For bug fixes: Create tests to '
    'verify issues before implementing fixes\n   * For new features: Consider test-driven development when '
    'appropriate\n   * If the repository lacks testing infrastructure and implementing tests would require '
    'extensive setup, consult with the user before investing time in building testing infrastructure\n   * '
    'If the environment is not set up to run tests, consult with the user first before investing time to '
    'install all dependencies\n4. IMPLEMENTATION: Make focused, minimal changes to address the problem\n5. '
    'VERIFICATION: If the environment is set up to run tests, test your implementation thoroughly, including '
    'edge cases. If the environment is not set up to run tests, consult with the user first before investing '
    'time to run tests.\n</PROBLEM_SOLVING_WORKFLOW>\n\n<SECURITY>\n* Only use GITHUB_TOKEN and other '
    'credentials in ways the user has explicitly requested and would expect.\n* Use APIs to work with GitHub '
    'or other platforms, unless the user asks otherwise or your task requires browsing.\n</'
    'SECURITY>\n\n<ENVIRONMENT_SETUP>\n* When user asks you to run an application, don\'t stop if the '
    'application is not installed. Instead, please install the application and run the command again.\n* If '
    'you encounter missing dependencies:\n  1. First, look around in the repository for existing dependency '
    'files (requirements.txt, pyproject.toml, package.json, Gemfile, etc.)\n  2. If dependency files exist, '
    'use them to install all dependencies at once (e.g., `pip install -r requirements.txt`, `npm install`, '
    'etc.)\n  3. Only install individual packages directly if no dependency files are found or if only '
    'specific packages are needed\n* Similarly, if you encounter missing dependencies for essential tools '
    'requested by the user, install them when possible.\n</ENVIRONMENT_SETUP>\n\n<TROUBLESHOOTING>\n* If '
    'you\'ve made repeated attempts to solve a problem but tests still fail or the user reports it\'s still '
    'broken:\n  1. Step back and reflect on 5-7 different possible sources of the problem\n  2. Assess the '
    'likelihood of each possible cause\n  3. Methodically address the most likely causes, starting with the '
    'highest probability\n  4. Document your reasoning process\n* When you run into any major issue while '
    'executing a plan from the user, please don\'t try to directly work around it. Instead, propose a new '
    'plan and confirm with the user before proceeding.\n</TROUBLESHOOTING>')

register_template(Mistral3TemplateMeta('devstral', default_system=devstral_small_2505_system))


class Mistral2506Template(Mistral2503Template):

    def _get_mistral_system(self):
        from swift.llm import get_model_name
        model_dir = self.model_info.model_dir
        model_name = get_model_name(model_dir)
        file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt')
        with open(file_path, 'r') as file:
            system_prompt = file.read()
        today = datetime.today().strftime('%Y-%m-%d')
        yesterday = (datetime.today() - timedelta(days=1)).strftime('%Y-%m-%d')
        return system_prompt.format(name=model_name, today=today, yesterday=yesterday)

    def _swift_encode(self, inputs: StdTemplateInputs):
        if inputs.system is None and self.use_chat_template:
            inputs.system = self._get_mistral_system()
        return super()._swift_encode(inputs)


register_template(
    Mistral3TemplateMeta(MLLMTemplateType.mistral_2506, default_system=None, template_cls=Mistral2506Template))


class Mistral2512Template(Mistral2506Template):

    def _get_mistral_system(self):
        model_dir = self.model_info.model_dir
        file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt')
        with open(file_path, 'r') as file:
            system_prompt = file.read()
        today = datetime.today().strftime('%Y-%m-%d')
        yesterday = (datetime.today() - timedelta(days=1)).strftime('%Y-%m-%d')
        return system_prompt.format(today=today, yesterday=yesterday)


register_template(
    Mistral3TemplateMeta(
        MLLMTemplateType.mistral_2512, default_system=None, template_cls=Mistral2512Template, agent_template='mistral'))


class Mistral2512ThinkingTemplate(Mistral2506Template):

    def _get_mistral_system(self):
        model_dir = self.model_info.model_dir
        file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt')
        with open(file_path, 'r') as file:
            system_prompt = file.read()
        return system_prompt


register_template(
    Mistral3TemplateMeta(
        MLLMTemplateType.mistral_2512_thinking,
        default_system=None,
        template_cls=Mistral2512ThinkingTemplate,
        agent_template='mistral'))
