import os
import requests
from datetime import datetime
from typing import Optional
import pytz


class Tools:
    def __init__(self):
        """
        Initialize the tools.
        """
        self.tool_endpoint = "http://100.84.94.73:8020"

    # Add your custom tools using pure Python code here, make sure to add type hints
    # Use Sphinx-style docstrings to document your tools, they will be used for generating tools specifications
    # Please refer to function_calling_filter_pipeline.py file from pipelines project for an example

    def calculator(self, equation: str) -> str:
        """
        Perform mathematical calculations on a given equation.

        Args:
            equation (str): A valid mathematical expression to evaluate. 
                           Example: "2 + 3 * (4 - 1)"

        Returns:
            str: The calculation result in format "equation = result" or error message if invalid
        """

        # Avoid using eval in production code
        # https://nedbatchelder.com/blog/201206/eval_really_is_dangerous.html
        try:
            result = eval(equation)
            return f"{equation} = {result}"
        except Exception as e:
            print(e)
            return "Invalid equation"

    def search_property_from_material_project(
        self,
        formula: str | list[str],
        chemsys: Optional[str | list[str] | None] = None,
        crystal_system: Optional[str | list[str] | None] = None,
        is_gap_direct: Optional[bool | None] = None,
        is_stable: Optional[bool | None] = None
    ) -> str:
        """
        Search material properties from Material Project database.

        Args:
            formula (str | list[str]): [Required] Chemical formula(s) to search. Example: "Fe2O3" or ["ABO3", "Si*"]
            chemsys (str | list[str] | None): [Optional] Chemical system(s) to search. Example: "Li-Fe-O" or ["Si-O", "Li-Fe-P"]
            crystal_system (str | list[str] | None): [Optional] Crystal system(s) to filter by. Example: "Cubic" or ["Hexagonal", "Tetragonal"]
            is_gap_direct (bool | None): [Optional] Filter by materials with direct band gap
            is_stable (bool | None): [Optional] Filter by thermodynamic stability

        Returns:
            str: JSON string containing material properties or error message
        """
        # Build request parameters
        params = {
            'chemsys': chemsys,
            'crystal_system': crystal_system,
            'formula': formula,
            'is_gap_direct': is_gap_direct,
            'is_stable': is_stable,
            'chunk_size': 5
        }
        
        # Filter out None values
        params = {k: v for k, v in params.items() if v is not None}
        
        try:
            response = requests.get(
                f"{self.tool_endpoint}/mp/search",
                params=params,
                timeout=30
            )
            response.raise_for_status()
            return str(response.json()['data'])
            
        except requests.exceptions.Timeout:
            return "Request timed out"
        except requests.exceptions.RequestException as e:
            return f"Request failed: {str(e)}"


    def search_from_oqmd_by_composition(self, composition: str) -> str:
        """
        Search materials from OQMD database by chemical composition.

        Args:
            composition (str): Chemical composition string. Example: "CsPbBr3"

        Returns:
            str: JSON string containing material data or error message
        """
        # 构建请求参数
        param = {
            'composition': composition
        }
        try:
            # 发送请求到/oqmd/search路由
            response = requests.get(
                self.tool_endpoint + "/oqmd/search",
                params=param
            )
            response.raise_for_status()
            return str(response.json()['data'])
        except requests.exceptions.RequestException as e:
            return f"Error: {str(e)}"
        
    
    def optimize_crystal_structure(self, markdown_block: str, input_format: str = "cif", output_format: str = "cif") -> str:
        """
        Optimize crystal structure using various file formats.

        Args:
            markdown_block (str): Markdown block containing structure data. 
                                Format: ```format\ncontent\n``` where format is one of ["cif", "poscar", "json", "xyz"]
            input_format (str): Input file format. Must be one of ["cif", "poscar", "json", "xyz"]
            output_format (str): Output file format. Must be one of ["cif", "poscar", "json", "xyz"]

        Returns:
            str: JSON string containing optimized structure data or error message
        """
        # Validate input/output formats
        valid_formats = ["cif", "poscar", "json", "xyz"]
        if input_format not in valid_formats:
            return f"Error: Invalid input format. Must be one of {valid_formats}"
        if output_format not in valid_formats:
            return f"Error: Invalid output format. Must be one of {valid_formats}"

        # Define format specific validation
        format_validators = {
            "cif": ("```cif\n", "\n```"),
            "poscar": ("```poscar\n", "\n```"), 
            "json": ("```json\n", "\n```"),
            "xyz": ("```xyz\n", "\n```")
        }

        try:
            # Validate markdown block format
            start_tag, end_tag = format_validators[input_format]
            if not markdown_block.startswith(start_tag) or not markdown_block.endswith(end_tag):
                return f"Error: Invalid {input_format} markdown format. Expected format: {start_tag}...{end_tag}"
            
            # Extract content
            content = markdown_block[len(start_tag):-len(end_tag)].strip()
            if not content:
                return f"Error: Empty {input_format} content"

            # Prepare request parameters
            params = {
                'input_format': input_format,
                'output_format': output_format
            }

            # Send request to optimization service
            try:
                response = requests.post(
                    self.tool_endpoint + "/fairchem/optimize_structure",
                    data=content,
                    headers={"Content-Type": "text/plain"},
                    params=params,
                    timeout=30  # Add timeout
                )
                response.raise_for_status()
                
                # 直接返回完整响应
                return str(response.json()['data'])
                
            except requests.exceptions.Timeout:
                return "Error: Optimization request timed out"
            except requests.exceptions.RequestException as e:
                return f"Error: Optimization request failed - {str(e)}"
            except ValueError as e:
                return f"Error: Invalid JSON response - {str(e)}"
                
        except Exception as e:
            return f"Error: {str(e)}"


if __name__ == '__main__':
    tools = Tools()
    # print(tools.search_property_from_material_project(formula='CsPbBr3'))
    # print(tools.search_from_material_project())
    print(tools.search_from_oqmd_by_composition("CsPbBr3"))
#     print(tools.optimize_crystal_structure("""```cif\n# generated using pymatgen
# data_CsPbBr3
# _symmetry_space_group_name_H-M   Pnma
# _cell_length_a   8.45384704
# _cell_length_b   11.87891123
# _cell_length_c   8.10107841
# _cell_angle_alpha   90.00000000
# _cell_angle_beta   90.00000000
# _cell_angle_gamma   90.00000000
# _symmetry_Int_Tables_number   62
# _chemical_formula_structural   CsPbBr3
# _chemical_formula_sum   'Cs4 Pb4 Br12'
# _cell_volume   813.53053480
# _cell_formula_units_Z   4
# loop_
#  _symmetry_equiv_pos_site_id
#  _symmetry_equiv_pos_as_xyz
#   1  'x, y, z'
#   2  '-x, -y, -z'
#   3  '-x+1/2, -y, z+1/2'
#   4  'x+1/2, y, -z+1/2'
#   5  'x+1/2, -y+1/2, -z+1/2'
#   6  '-x+1/2, y+1/2, z+1/2'
#   7  '-x, y+1/2, -z'
#   8  'x, -y+1/2, z'
# loop_
#  _atom_type_symbol
#  _atom_type_oxidation_number
#   Cs+  1.0
#   Pb2+  2.0
#   Br-  -1.0
# loop_
#  _atom_site_type_symbol
#  _atom_site_label
#  _atom_site_symmetry_multiplicity
#  _atom_site_fract_x
#  _atom_site_fract_y
#  _atom_site_fract_z
#  _atom_site_occupancy
#   Cs+  Cs0  4  0.06016347  0.75000000  0.01945212  1
#   Pb2+  Pb1  4  0.00000000  0.00000000  0.50000000  1
#   Br-  Br2  8  0.20102217  0.03170128  0.80279189  1
#   Br-  Br3  4  0.00625945  0.25000000  0.44033452  1
# \n```""", input_format="cif", output_format="poscar"))