import os
from typing import Dict, Any, Optional
import time
import tempfile
from pathlib import Path
import subprocess
import json
import shlex
import celpy
import requests
import yaml
import logging
from environs import Env
from llmops.base import ExecutionResult, RequestType, Step, CliRequest

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

cel_env = celpy.Environment()


CLI_MAPPING = {
    "aws_cli": "cmd_aws",
    "azure_cli": "cmd_azure",
    "gcp_cli": "cmd_gcp",
    "mysqlsh": "cmd_mysqlsh",
}


class Executor:
    def __init__(self):
        # TODO: use shell script
        pass


class HttpExecutor(Executor):
    """only for bearer token now"""

    def __init__(
        self,
        base_url: str,
        api_token: Optional[str] = None,
        auth_type: Optional[str] = None,
    ):
        self.base_url = base_url.rstrip("/")
        self.requests = []
        self.auth_type = auth_type

        # 加载 token
        self.token = api_token

    def _get_headers(self, headers: Optional[str] = None) -> Dict[str, str]:
        """only for bearer token now"""
        result = {"Accept": "application/json"}

        # if headers:
        #     headers_dict = headers if isinstance(headers, dict) else json.loads(headers)
        #     result.update(headers_dict)

        if self.token:
            result["Authorization"] = f"bearer {self.token}"

        return result

    def _get_params(self, params: Optional[str] = None) -> Dict[str, Any]:
        if not params:
            return {}
        return json.loads(params)

    def _format_path(self, path: str, path_params: Optional[str] = None) -> str:
        if not path_params:
            return path

        params = json.loads(path_params)
        for key, value in params.items():
            path = path.replace(f"{{{key}}}", str(value))
        return path

    def execute_request(
        self,
        method: str,
        path: str,
        headers: Optional[str] = None,
        query_params: Optional[str] = None,
        path_params: Optional[str] = None,
        body: Optional[str] = None,
    ) -> ExecutionResult:
        try:
            # 准备请求参数
            url = f"{self.base_url}{self._format_path(path, path_params)}"
            headers = self._get_headers(headers)
            params = self._get_params(query_params)
            data = json.loads(body) if body else None
            self.requests.append(
                {
                    "url": url,
                    "method": method,
                    "headers": headers,
                    "params": params,
                    "body": data,
                }
            )

            # 发送请求
            response = requests.request(
                method=method, url=url, headers=headers, params=params, json=data
            )
            # 解析响应
            try:
                json_data = response.json() if response.text else {}
            except json.JSONDecodeError:
                json_data = {}

            return ExecutionResult(
                status_code=response.status_code,
                json=json_data,
                headers=headers,  # dict(response.headers)
            )

        except Exception as e:
            return ExecutionResult(status_code=500, json={}, headers={}, error=str(e))

    def execute_with_retry(
        self,
        method: str,
        path: str,
        headers: Optional[str] = None,
        query_params: Optional[str] = None,
        path_params: Optional[str] = None,
        body: Optional[str] = None,
        max_retries: int = 0,
        delay_after: int = 0,
        expect: Optional[str] = None,
    ) -> ExecutionResult:
        if max_retries == 0 and delay_after == 0:
            return self.execute_request(
                method=method,
                path=path,
                headers=headers,
                query_params=query_params,
                path_params=path_params,
                body=body,
            )

        if method.upper() != "GET":
            raise ValueError(f"Retry is only supported for GET requests, got {method}")

        for i in range(max_retries + 1):
            response = self.execute_request(
                method=method,
                path=path,
                headers=headers,
                query_params=query_params,
                path_params=path_params,
                body=body,
            )
            cel_src = expect
            cel_ast = cel_env.compile(cel_src)
            cel_prgm = cel_env.program(cel_ast)

            if (
                response.status_code == 200 and cel_prgm.evaluate(response.json)
            ) or i == max_retries:
                return response

            time.sleep(delay_after)

        return response


class CompoundExecutor(Executor):
    """Executor that uses shell script to execute requests"""

    def __init__(
        self,
        base_url: str,
        api_token: Optional[str | Dict[str, str]] = None,
        auth_type: Optional[str] = None,
    ):
        self.base_url = base_url.rstrip("/")
        self.api_token = api_token
        self.auth_type = auth_type if auth_type else "bearer"
        self.requests = []  # 存储所有请求，用于重放
        self.script_path = Path(__file__).parent / "do-action.sh"
        if not self.script_path.exists():
            raise RuntimeError(f"Script not found: {self.script_path}")

    def _format_path(self, path: str, path_params: Optional[str] = None) -> str:
        """Format path with path parameters"""
        if not path_params:
            return path
        try:
            params = json.loads(path_params)
            for key, value in params.items():
                path = path.replace(f"{{{key}}}", str(value))
            return path
        except json.JSONDecodeError:
            return path

    def execute_http_request(
        self,
        method: str,
        path: str,
        headers: Optional[str] = None,
        query_params: Optional[str] = None,
        path_params: Optional[str] = None,
        body: Optional[str] = None,
        auth_type: Optional[str] = None,
    ) -> ExecutionResult:
        """Execute request by calling execute_with_retry with no retries"""
        return self.execute_http_with_retry(
            method=method,
            path=path,
            headers=headers,
            query_params=query_params,
            path_params=path_params,
            body=body,
            max_retries=0,
            delay_after=0,
            auth_type=self.auth_type or auth_type,
        )

    def execute_with_retry(self, step: Step) -> ExecutionResult:
        if step.request_type == RequestType.HTTP:
            return self.execute_http_with_retry(
                method=step.request.method,
                path=step.request.path,
                headers=step.request.headers,
                query_params=step.request.query_params,
                path_params=step.request.path_params,
                body=step.request.body,
                max_retries=step.max_retries,
                delay_after=step.delay_after,
                expect=step.expect,
                auth_type=step.request.auth.auth_type,
                role=step.request.auth.content,
            )
        elif step.request_type in [RequestType.AWS_CLI, RequestType.GCP_CLI, RequestType.AZURE_CLI, RequestType.MYSQLSH]:
            return self.execute_cli_with_retry(
                request_type=step.request_type.value,
                request=step.request,
                max_retries=step.max_retries,
                delay_after=step.delay_after,
                expect=step.expect,
            )
        else:
            raise ValueError(f"Unsupported request type: {step.request_type}")

    def execute_http_with_retry(
        self,
        method: str,
        path: str,
        headers: Optional[str] = None,
        query_params: Optional[str] = None,
        path_params: Optional[str] = None,
        body: Optional[str] = None,
        max_retries: int = 0,
        delay_after: int = 0,
        expect: Optional[str] = None,
        auth_type: Optional[str] = None,
        role: Optional[str] = None,
    ) -> ExecutionResult:
        """Execute request using shell script with retries"""
        if isinstance(self.api_token, dict):
            if role in self.api_token:
                request_api_token = self.api_token[role]
            else:
                if len(self.api_token) == 1:
                    request_api_token = list(self.api_token.values())[0]
                else:
                    logger.warning("Multiple tokens found, but no role specified. Using the first token.")
                    request_api_token = list(self.api_token.values())[0]
        else:
            request_api_token = self.api_token
        # store request for replay
        self.requests.append(
            {
                "method": method,
                "path": path,
                "headers": headers,
                "query_params": query_params,
                "path_params": path_params,
                "body": body,
                "auth_type": self.auth_type or auth_type,
            }
        )

        # prepare request yaml
        request_obj = {
            "method": method,
            "path": self._format_path(path, path_params).rstrip("/"),
            "headers": json.loads(headers) if headers else {},
            "query_params": json.loads(query_params) if query_params else {},
            "body": json.loads(body) if body else {},
            "auth": {
                "auth_type": auth_type,
                "content": role if role else "placeholder",
            },
        }

        # Append query parameters to path if any
        if request_obj["query_params"]:
            from urllib.parse import urlencode
            query_string = urlencode(request_obj["query_params"], doseq=True)
            request_obj["path"] = f"{request_obj['path']}?{query_string}"
            request_obj["query_params"] = {}  # Clear query_params since they're now in the path

        if expect is not None and method == "GET":  # polling get
            request_data = {
                "request_type": "http",  # TODO: don't hardcode http
                "request": json.dumps(request_obj),
                "expect": expect,
                "max_retries": 60,
                "delay_after": 30,
            }
        else:
            request_data = {
                "request_type": "http",
                "request": json.dumps(request_obj),
                "expect": expect,
                "max_retries": max_retries,
                "delay_after": delay_after,
            }
        self.reqeust_data = request_data

        # create temp yaml file
        try:
            tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
            yaml.dump(request_data, tmp)
            tmp.close()

            # build command
            cmd = f"bash {self.script_path} -i {tmp.name} -a {self.auth_type or auth_type} -t {request_api_token} -H {self.base_url}"

            # execute command
            try:
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

                try:
                    response_data = json.loads(result.stdout)
                    if result.returncode == 0:
                        # Success case
                        return ExecutionResult(
                            status_code=response_data.get("status_code", 500),
                            headers=response_data.get("headers", {}),
                            json=response_data.get("body", {}),
                        )
                    else:
                        # Retry failed case
                        return ExecutionResult(
                            status_code=response_data.get("status_code", 500),
                            headers=response_data.get("headers", {}),
                            json=response_data.get("body", {}),
                            error=f"Command failed with return code {result.returncode}: {result.stderr}",
                        )
                except json.JSONDecodeError as e:
                    logger.error(
                        f"Failed to parse response JSON: {e}, request: {request_data}, response: {result.stdout}"
                    )
                    return ExecutionResult(
                        status_code=500,
                        headers={},
                        json={},
                        error=str(e),
                    )
            except Exception as e:
                logger.error(
                    f"Failed to execute command: {e}, request: {request_data}, response: {result}"
                )
                return ExecutionResult(
                    status_code=500,
                    headers={},
                    json={},
                    error=str(e),
                )
        finally:
            try:
                os.unlink(tmp.name)
            except FileNotFoundError:
                pass

    def execute_cli_with_retry(
        self,
        request_type: str,
        request: CliRequest,
        max_retries: int = 0,
        delay_after: int = 0,
        expect: Optional[str] = None,
    ) -> ExecutionResult:
        """Execute request using shell script with retries"""
        self.requests.append(
            {
                "service": request.service,
                "command": request.command,
                "args": request.args,
                "args_params": request.args_params,
            }
        )

        safe_args = [shlex.quote(arg) for arg in request.args]

        request_obj = {
            "command": request.command,
            "args": safe_args,
            "args_params": request.args_params,
        }

        if expect is not None:  # and method == "GET":  # polling get
            request_data = {
                "request_type": CLI_MAPPING[request_type],
                "request": json.dumps(request_obj),
                "expect": expect,
                "max_retries": 60,
                "delay_after": 30,
            }
        else:
            request_data = {
                "request_type": CLI_MAPPING[request_type],
                "request": json.dumps(request_obj),
                "expect": expect,
                "max_retries": max_retries,
                "delay_after": delay_after,
            }
        self.reqeust_data = request_data

        # create temp yaml file
        try:
            tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
            yaml.dump(request_data, tmp)
            tmp.close()

            # aws_profile = "assume-prod"

            # build command
            # cmd = f"bash {self.script_path} -i {tmp.name} -p {aws_profile}"
            cmd = f"bash {self.script_path} -i {tmp.name}"

            # execute command
            try:
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

                try:
                    if result.returncode == 0:
                        # response_data = json.loads({"output": result.stdout})
                        # Success case
                        return ExecutionResult(
                            status_code=result.returncode,
                            headers={},
                            json={"output": result.stdout},  #response_data,
                            error=result.stderr
                        )
                    else:
                        # Retry failed case
                        return ExecutionResult(
                            status_code=result.returncode,
                            headers={},
                            json={"output": result.stdout},
                            error=f"Command failed with return code {result.returncode}: {result.stderr}",
                        )
                except json.JSONDecodeError as e:
                    logger.error(
                        f"Failed to parse response JSON: {e}, request: {request_data}, response: {result.stdout}"
                    )
                    return ExecutionResult(
                        status_code=result.returncode,
                        headers={},
                        json={"output": result.stdout},
                        error=str(e),
                    )
            except Exception as e:
                logger.error(
                    f"Failed to execute command: {e}, request: {request_data}, response: {result}"
                )
                return ExecutionResult(
                    status_code=result.returncode,
                    headers={},
                    json={"output": result.stdout},
                    error=str(e),
                )
        finally:
            try:
                os.unlink(tmp.name)
            except FileNotFoundError:
                pass


if __name__ == "__main__":
    import json
    import logging
    import time
    from datetime import datetime
    from pathlib import Path
    from environs import Env

    src_parent_path = Path(__file__).parent.parent.parent.parent

    env = Env()
    env.read_env(
        src_parent_path
        / "examples"
        / "tidbcloud_dedicated_sut"
        / "configs"
        / "sut"
        / ".env"
    )

    api_token = env("SUT_TOKEN")

    # Setup logging
    logging.basicConfig(level=logging.DEBUG)
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)

    # Test request
    request_data = {
        "request_type": "http",
        "request": json.dumps(
            {
                "method": "GET",
                "path": "/clusters",
                "headers": {},
                "query_params": {},
                "body": {},
                "auth": {"auth_type": "digest", "content": api_token},
            }
        ),
        "expect": json.dumps(
            {
                "status_code": 200,
                # 'cel': "body.state == 'ACTIVE'"
            }
        ),
        "max_retries": 5,
        "delay_after": 5,
    }

    # Create executor
    executor = CompoundExecutor(
        base_url="https://dedicated.staging.tidbapi.com/v1beta1",
        api_token=api_token,
        auth_type="digest",
    )

    # Create temp yaml file for testing
    import tempfile
    import yaml

    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as tmp:
        # Write request data to temp file
        yaml.dump(request_data, tmp)
        tmp.flush()

        # Build command
        script_path = Path(__file__).parent / "do-action.sh"
        cmd = f"bash {script_path} -i {tmp.name} -a digest -t {executor.api_token} -H {executor.base_url}"

        # Execute command with debug output
        logger.debug(f"Start time: {datetime.now()}")
        start_time = time.time()

        logger.debug(f"Executing command: {cmd}")
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

        end_time = time.time()
        duration = end_time - start_time
        logger.debug(f"End time: {datetime.now()}")
        logger.debug(f"Duration: {duration:.2f} seconds")

        logger.debug(f"Command return code: {result.returncode}")
        logger.debug(
            f"Command stdout: {result.stdout!r}"
        )  # Use !r to show raw string with escapes
        logger.debug(f"Command stderr: {result.stderr!r}")

        # Try to parse response
        try:
            response_data = json.loads(result.stdout)
            logger.debug(f"Parsed response: {json.dumps(response_data, indent=2)}")
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse response JSON: {e}")
            logger.debug(
                f"Raw response: {result.stdout!r}"
            )  # Use !r to show raw string