import json
import time
from html import escape
from pathlib import Path
from typing import Literal

from loguru import logger
from pydantic import BaseModel
from xpdc.adb.device import ADBDevice
from xpdc.hdc.device import HDCDevice

from phone_copilot import image
from phone_copilot.action import Action, ActionResult
from phone_copilot.agent.context import Context
from phone_copilot.agent.prompt import EN_SYSTEM_PROMPT, ZH_SYSTEM_PROMPT
from phone_copilot.model_client import ModelClient, ModelResponse


class Agent:
    def __init__(
        self,
        *,
        device: ADBDevice | HDCDevice,
        system_prompt: str = None,
        language: Literal['zh', 'en'] = 'zh',
        max_steps: int = 100,
        model_client: ModelClient,
        ccf: int = 1000,
    ) -> None:
        self._device = device
        self._system_prompt = system_prompt or (ZH_SYSTEM_PROMPT if language == 'zh' else EN_SYSTEM_PROMPT)
        self._max_steps = max_steps
        self._model_client = model_client
        self._ccf = ccf

        self._context = Context()
        self._step_count = 0
        self._steps: list[StepResult] = []
        self._task = None

    def run(self, task: str) -> None:
        self.reset()
        self._task = task
        self._context.append_user_content(task)

        start_time = time.perf_counter()
        logger.info(f'Starting task: {task}')

        try:
            while self._step_count < self._max_steps:
                result = self._step()
                self._steps.append(result)

                if not result.model_response or not result.action_result:
                    logger.error('Model response or action result is None, agent stopped.')
                    break
                elif result.action_result and result.action_result.should_finish:
                    logger.info('Task finished by the agent.')
                    break
            else:
                logger.warning('Maximum steps reached, agent stopped.')

        except KeyboardInterrupt:
            logger.warning('Agent task interrupted by user.')

        logger.info(f'Task completed in {time.perf_counter() - start_time:.2f} seconds with {self._step_count} steps.')

    def _step(self) -> 'StepResult':
        self._step_count += 1

        try:
            screenshot = self._device.screenshot('phone-copilot-tmp-screenshot.png')
            screenshot._base64 = image.compress_base64_image(base64_image=screenshot.base64, quality=10)
            screenshot.unlink(missing_ok=True)

            screenshot_info = json.dumps({'current_app': self._device.get_current_app()}, ensure_ascii=False)
        except Exception as e:  # noqa: BLE001
            msg = f'Fail to capture screenshot: {e}'
            logger.error(msg)

            return StepResult(index=self._step_count, model_response=None, action_result=None)

        self._context.append_image_content(image_base64=screenshot.base64, image_info=screenshot_info)

        try:
            model_response = self._model_client.chat(self._context)
        except Exception as e:  # noqa: BLE001
            msg = f'Fail to get model response: {e}'
            logger.error(msg)

            return StepResult(
                index=self._step_count, model_response=None, action_result=None, screenshot_base64=screenshot.base64
            )

        self._context.remove_image_base64()

        action_result = Action(
            model_response=model_response,
            device=self._device,
            width=screenshot.width,
            height=screenshot.height,
            ccf=self._ccf,
        )()

        self._context.append_assistant_content(
            thinking=model_response.thinking,
            action=model_response.action,
            result=action_result.model_dump_json(exclude_none=True, ensure_ascii=False, exclude={'action'}),
        )

        return StepResult(
            index=self._step_count,
            model_response=model_response,
            action_result=action_result,
            screenshot_base64=screenshot.base64,
        )

    def reset(self) -> None:
        """Reset the agent's context and step count."""
        self._context.reset()
        self._step_count = 0
        self._steps = []
        self._context.append_system_content(self._system_prompt)
        self._task = None

    @property
    def max_steps(self) -> int:
        """Get the maximum number of steps allowed for the agent."""
        return self._max_steps

    @property
    def steps(self) -> list['StepResult']:
        """Get the list of step results."""
        return self._steps

    @property
    def context(self) -> Context:
        """Get the agent's context."""
        return self._context

    @property
    def step_count(self) -> int:
        """Get the current step count."""
        return self._step_count

    def export_json(self, path: Path) -> Path:
        """Export the agent's execution details to a JSON file."""
        path.write_text(self.json, encoding='utf-8')
        logger.info(f'Exporting JSON report to [{path}]')
        return path

    def export_html(self, path: Path) -> Path:
        """Export the agent's execution details to an HTML file."""
        path.write_text(self.html, encoding='utf-8')
        logger.info(f'Exporting HTML report to [{path}]')
        return path

    @property
    def json(self) -> str:
        """Get a JSON representation of the agent's execution."""
        logger.info('Generating JSON report...')

        return json.dumps(
            {
                'task': self._task,
                'device_id': self._device.device_id,
                'steps': [step.model_dump(mode='json') for step in self._steps],
            },
            ensure_ascii=False,
            indent=4,
        )

    @property
    def html(self) -> str:
        """Generate an HTML report of the agent's execution."""
        logger.info('Generating HTML report...')

        steps = '\n'.join(step.html for step in self._steps)

        return f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8" />
<title>Phone Copilot Task Report</title>
<style>
    body {{ font-family: -apple-system, system-ui, sans-serif; margin: 24px; color: #111; }}
    header {{ border-bottom: 1px solid #ddd; margin-bottom: 16px; padding-bottom: 8px; }}
    section {{ margin-bottom: 20px; }}
        .step {{ border: 1px solid #e1e1e1; border-radius: 8px; padding: 12px; margin: 12px 0; }}
        .meta {{ color: #555; font-size: 14px; }}
        .badge {{ display: inline-block; padding: 2px 8px; border-radius: 999px; font-size: 12px; }}
        .badge.success {{ background: #e6ffed; color: #036b26; }}
        .badge.failure {{ background: #ffe6e6; color: #8a1111; }}
        .step-content {{ display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 12px; align-items: start; }}
        .step-col pre {{ margin: 0; }}
    pre {{ white-space: pre-wrap; background: #f7f7f7; padding: 8px; border-radius: 6px; }}
    img {{ max-width:100%; max-height:480px; border:1px solid #ddd; border-radius:6px; }}
</style>
</head>
<body>
<header>
    <h1>Phone Copilot Task Report</h1>
    <div class="meta">Task: {escape(self._task) if self._task else 'None'}</div>
    <div class="meta">Device ID: {escape(self._device.device_id)}</div>
    <div class="meta">Total steps: {len(self.steps)}</div>
</header>
{steps}
</body>
</html>
"""  # noqa: E501


class StepResult(BaseModel):
    index: int
    model_response: ModelResponse | None
    action_result: ActionResult | None
    screenshot_base64: str | None = None

    def _mark_screenshot(self) -> None:
        """
        Mark the screenshot with the action performed in this step,
        such as drawing a red point for tap actions or a red arrow for swipe actions.

        This helps to visually identify the action on the screenshot when viewing the HTML report.
        """
        if not self.action_result or not self.screenshot_base64:
            return

        try:
            if str(self.action_result.action) in ('Tap', 'DoubleTap', 'LongPress'):
                element = self.action_result.action.param.element
                self.screenshot_base64 = image.draw_red_point_base64(self.screenshot_base64, *element)
            elif str(self.action_result.action) == 'Swipe':
                start = self.action_result.action.param.start
                end = self.action_result.action.param.end
                self.screenshot_base64 = image.draw_red_arrow_base64(self.screenshot_base64, *start, *end)
        except Exception as e:  # noqa: BLE001
            logger.error(f'Failed to mark screenshot for step {self.index}: {e}')

    @property
    def html(self) -> str:
        """Generate HTML representation of the step result."""
        self._mark_screenshot()

        badge = 'success' if self.action_result and self.action_result.success else 'failure'
        time_to_first_token = self.model_response.time_to_first_token if self.model_response else None
        time_to_last_token = self.model_response.time_to_last_token if self.model_response else None
        screenshot_html = (
            (
                f'<div><img src="data:image/png;base64,{self.screenshot_base64}" '
                f'style="max-width:600px; border:1px solid #ccc;"/></div>'
            )
            if self.screenshot_base64
            else ''
        )

        return f"""
<section class="step">
    <div class="meta">Step {self.index} - <span class="badge {badge}">{badge}</span></div>
    <div class="meta">Time to first token: {time_to_first_token} seconds</div>
    <div class="meta">Time to last token: {time_to_last_token} seconds</div>
    <div class="step-content">
        <div class="step-col">
            <div><strong>Content</strong></div>
            <pre>{escape(self.model_response.content) if self.model_response else None}</pre>
        </div>
        <div class="step-col">
            <div><strong>Action</strong></div>
            <pre>{escape(self.model_response.action) if self.model_response else None}</pre>
        </div>
        <div class="step-col">
            <div><strong>Result</strong></div>
            <pre>{escape(self.action_result.model_dump_json(exclude_none=True, ensure_ascii=False, indent=4)) if self.action_result else None}</pre>
        </div>
        <div class="step-col">
            <div><strong>Screenshot</strong></div>
            {screenshot_html}
        </div>
    </div>
</section>
"""  # noqa: E501
