from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from prompts import REASONING_SYSTEM_PROMPT


PROJECT_ROOT = Path(__file__).resolve().parent


@dataclass(frozen=True)
class TrainingModel:
    model_name: str = "Qwen/Qwen2.5-3B-Instruct"
    experiment: int | None = None
    checkpoint: int | str | None = None
    results_dir: str = "results"
    experiment_prefix: str = "DPO_"
    output_dir: str | None = None

    @property
    def experiment_name(self) -> str | None:
        if self.experiment is None:
            return None

        return f"{self.experiment_prefix}{self.experiment}"

    @property
    def experiment_dir(self) -> str | None:
        if self.output_dir:
            return self.output_dir

        if self.experiment_name is None:
            return None

        return str(Path(self.results_dir) / self.experiment_name)

    @property
    def checkpoint_path(self) -> str | None:
        if self.checkpoint is None:
            return None

        checkpoint = str(self.checkpoint)
        checkpoint_path = Path(checkpoint)

        if checkpoint_path.is_absolute() or "/" in checkpoint or "\\" in checkpoint:
            return checkpoint

        if checkpoint.isdigit():
            checkpoint = f"checkpoint-{checkpoint}"
        elif not checkpoint.startswith("checkpoint-"):
            checkpoint = f"checkpoint-{checkpoint}"

        if self.experiment_dir is None:
            return checkpoint

        return str(Path(self.experiment_dir) / checkpoint)

    def sft_defaults(self) -> dict[str, Any]:
        return {
            "model_name": self.model_name,
            "results_dir": self.results_dir,
            "experiment_prefix": self.experiment_prefix,
            "experiment_name": self.experiment_name,
            "output_dir": self.output_dir,
            "initial_checkpoint_path": self.checkpoint_path,
            "resume_from_checkpoint": None,
        }

    def webapp_defaults(self) -> dict[str, Any]:
        return {
            "modelName": self.model_name,
            "experiment": self.experiment,
            "experimentName": self.experiment_name,
            "experimentDir": self.experiment_dir,
            "checkpoint": self.checkpoint,
            "checkpointPath": self.checkpoint_path,
        }


@dataclass(frozen=True)
class WebAppSettings:
    endpoint: str = "http://127.0.0.1:8000/api/chat"
    stream_mode: str = "demo"
    demo_prompt: str = (
        "Write a short answer about why careful reasoning matters."
    )
    max_new_tokens: int = 512
    temperature: float = 0.2
    top_p: float = 0.9
    thinking_effort: int = 100
    max_tool_rounds: int = 2
    use_query_loop: bool = True
    show_injected_result: bool = True
    instruction_prompt: str = REASONING_SYSTEM_PROMPT

    def webapp_defaults(self) -> dict[str, Any]:
        data = asdict(self)
        return {
            "endpoint": data["endpoint"],
            "streamMode": data["stream_mode"],
            "demoPrompt": data["demo_prompt"],
            "maxNewTokens": data["max_new_tokens"],
            "temperature": data["temperature"],
            "topP": data["top_p"],
            "thinkingEffort": data["thinking_effort"],
            "maxToolRounds": data["max_tool_rounds"],
            "useQueryLoop": data["use_query_loop"],
            "showInjectedResult": data["show_injected_result"],
            "instructionPrompt": data["instruction_prompt"],
        }


# Edit these two objects for project-wide defaults.
TRAINING_MODEL = TrainingModel(
    model_name="Qwen/Qwen2.5-3B-Instruct",
    experiment=4,
    checkpoint=181,
)

WEBAPP = WebAppSettings(
    endpoint="http://127.0.0.1:8000/api/chat",
    stream_mode="http",
)


def sft_config_defaults() -> dict[str, Any]:
    return TRAINING_MODEL.sft_defaults()


def webapp_config() -> dict[str, Any]:
    return {
        **TRAINING_MODEL.webapp_defaults(),
        **WEBAPP.webapp_defaults(),
    }


def write_webapp_config(
    path: str | Path = PROJECT_ROOT / "production" / "webapp" / "project-config.js",
) -> Path:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    payload = json.dumps(webapp_config(), ensure_ascii=False, indent=2)
    out_path.write_text(f"window.PROJECT_SETTINGS = {payload};\n", encoding="utf-8")
    return out_path


if __name__ == "__main__":
    print(write_webapp_config())