from dataclasses import dataclass
from datetime import datetime, timezone
import os
import logging
from pathlib import Path
import tempfile
import uuid
import json

from pydantic import BaseModel
from takk.models import PubSub
from takk.secrets import TakkLlmModels

from maur import git as git_ops
from maur.models import CodingTask, session
from maur.settings import GitlabSettings, GithubSettings, MaurSettings as CodeAgentSettings, MaurLLMSettings as CoderAgentLLMSettings
from maur.proc import await_command
from maur.opencode import complete_task

logger = logging.getLogger(__name__)

@dataclass
class ModelPrice:
    input_token: float
    output_token: float
    cache_read: float
    cache_write: float

# Price per million tokens in USD
_MODEL_PRICING: dict[TakkLlmModels, ModelPrice] = {
    "qwen3-coder-30b-a3b-instruct": ModelPrice(input_token=0.2,   output_token=0.8,  cache_read=0.0,  cache_write=0.0),
    "devstral-2-123b-instruct-2512": ModelPrice(input_token=0.4,   output_token=2.0,  cache_read=0.0,  cache_write=0.0),
}

def _compute_cost(
    model: str,
    input_tokens: int,
    output_tokens: int,
    cache_read_tokens: int,
    cache_write_tokens: int,
) -> float | None:
    pricing = next(
        (v for k, v in _MODEL_PRICING.items() if k in model or model in k),
        None,
    )
    if pricing is None:
        return None
    return (
        input_tokens * pricing.input_token
        + output_tokens * pricing.output_token
        + cache_read_tokens * pricing.cache_read
        + cache_write_tokens * pricing.cache_write
    ) / 1_000_000


def _format_cost_section(model: str, task: "CodingTask") -> str:
    lines = [
        "**Token usage:**",
        f"- Model: `{model}`",
        f"- Input: {task.tokens_input:,}",
        f"- Output: {task.tokens_output:,}",
    ]
    if task.tokens_cache_read:
        lines.append(f"- Cache read: {task.tokens_cache_read:,}")
    if task.tokens_cache_write:
        lines.append(f"- Cache write: {task.tokens_cache_write:,}")

    cost = _compute_cost(model, task.tokens_input, task.tokens_output, task.tokens_cache_read, task.tokens_cache_write)

    if cost is not None:
        lines.append(f"- Estimated cost: €{cost:.4f} EUR")

    return "\n".join(lines) + "\n\n"


class CodingTaskMessage(BaseModel):
    task_id: uuid.UUID


def resolve_git_settings() -> GitlabSettings | GithubSettings:
    try:
        return GitlabSettings.from_env()
    except Exception:
        return GithubSettings.from_env()


async def process_coding_task(msg: CodingTaskMessage) -> None:
    llm_settings = CoderAgentLLMSettings.from_env()

    git_settings = resolve_git_settings()

    async for db in session():
        task = await CodingTask.get(msg.task_id, db)
        if task is None:
            logger.error(f"Task {msg.task_id} not found")
            return

        task.started_at = datetime.now(tz=timezone.utc)
        await task.update(db)

        try:
            with tempfile.TemporaryDirectory() as workspace:
                repo_path = Path(workspace)

                if isinstance(git_settings, GitlabSettings):
                    repo_url = git_settings.gitlab_repo_url
                    git_token = git_settings.gitlab_token.get_secret_value()
                else:
                    repo_url = git_settings.github_repo_url
                    git_token = git_settings.github_token.get_secret_value()

                await git_ops.setup_workspace(
                    repo_url=repo_url,
                    branch=task.repo_branch,
                    task_id=str(task.id),
                    workspace_dir=repo_path,
                    git_token=git_token,
                )
                result_branch = f"maur/{task.id}"

                agent_output, step_token_lines = await complete_task(task, repo_path, llm_settings)

                if not await git_ops.has_changes(repo_path):
                    logger.info(f"Task {task.id}: no changes made, marking as completed without MR")
                    await task.update(db)
                    return

                _, _, _ = await await_command(
                    ["ruff", "check", ".", "--fix"],
                    cwd=repo_path,
                )

                summary = agent_output[:500] if agent_output else "No output"
                commit_message = f"chore: autonomous fix {task.id}\n\n{summary}"
                await git_ops.commit_and_push(repo_path, result_branch, commit_message)

                step_breakdown = "\n".join(step_token_lines) if step_token_lines else "  (none)"
                cost_section = _format_cost_section(llm_settings.maur_llm_model, task)

                mr_description = (
                    f"Automated fix for task `{task.id}`\n\n"
                    f"**Prompt:**\n{task.prompt}\n\n"
                    f"**OpenCode output:**\n```\n{agent_output[:2000]}\n```\n\n"
                    f"**Per-step token usage:**\n```\n{step_breakdown}\n```\n\n"
                    f"{cost_section}"
                )

                if isinstance(git_settings, GitlabSettings):
                    mr_url = await git_ops.create_gitlab_mr(
                        gitlab_url=git_settings.gitlab_url,
                        gitlab_token=git_token,
                        repo_url=repo_url,
                        title=task.source_id,
                        target_branch=task.repo_branch,
                        description=mr_description,
                        source_branch=result_branch,
                    )
                else:
                    mr_url = await git_ops.create_github_pr(
                        github_api_url=git_settings.github_api_url,
                        github_token=git_token,
                        repo_url=repo_url,
                        title=task.source_id,
                        target_branch=task.repo_branch,
                        description=mr_description,
                        source_branch=result_branch,
                    )

                task.result_branch = result_branch
                task.result_url = mr_url
                await task.update(db)

        except Exception as exc:
            logger.exception(f"Task {task.id} failed: {exc}")
            task.finnished_at = datetime.now(timezone.utc)
            await task.update(db)
