"""
GenerationSession — tracks all state for a single app generation run.
This is the central data object passed between all agents.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Any, Callable
from enum import Enum, auto


class AgentStatus(Enum):
    PENDING = auto()
    RUNNING = auto()
    DONE = auto()
    FAILED = auto()
    SKIPPED = auto()


@dataclass
class AgentRun:
    name: str
    engine: str
    status: AgentStatus = AgentStatus.PENDING
    started_at: Optional[float] = None
    ended_at: Optional[float] = None
    error: Optional[str] = None
    stats: dict = field(default_factory=dict)

    def start(self) -> None:
        self.status = AgentStatus.RUNNING
        self.started_at = time.time()

    def done(self, **stats: Any) -> None:
        self.status = AgentStatus.DONE
        self.ended_at = time.time()
        self.stats.update(stats)

    def fail(self, error: str) -> None:
        self.status = AgentStatus.FAILED
        self.ended_at = time.time()
        self.error = error

    @property
    def elapsed(self) -> float:
        if self.started_at is None:
            return 0.0
        end = self.ended_at or time.time()
        return end - self.started_at


@dataclass
class LogEntry:
    agent: str
    message: str
    level: str = "info"   # info | success | warning | error | highlight
    elapsed: float = 0.0


@dataclass
class AppPlan:
    """Structured plan produced by Agent 1 (Planner)."""
    app_name: str = ""
    slug: str = ""
    package_name: str = ""
    description: str = ""
    tagline: str = ""
    primary_color: str = "#2563eb"
    accent_color: str = "#f59e0b"
    background_color: str = "#0f172a"
    surface_color: str = "#1e293b"
    text_color: str = "#f8fafc"

    screens: list[dict] = field(default_factory=list)
    data_models: list[dict] = field(default_factory=list)
    features: list[str] = field(default_factory=list)
    user_stories: list[str] = field(default_factory=list)

    navigation: str = "tab"
    tab_screens: list[str] = field(default_factory=list)
    stack_screens: list[str] = field(default_factory=list)
    drawer_screens: list[str] = field(default_factory=list)

    storage: str = "asyncstorage"
    storage_key: str = "app_data"

    auth_required: bool = False
    auth_method: str = "none"

    external_apis: list[str] = field(default_factory=list)
    dependencies: list[str] = field(default_factory=list)
    dev_dependencies: list[str] = field(default_factory=list)

    target_platforms: list[str] = field(default_factory=lambda: ["ios", "android"])
    min_sdk_version: str = "13.0"

    clarifying_questions: list[str] = field(default_factory=list)
    assumptions: list[str] = field(default_factory=list)

    raw: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, d: dict) -> "AppPlan":
        import re
        slug = re.sub(r"[^a-z0-9-]", "-", d.get("appName", "my-app").lower().replace(" ", "-"))
        slug = re.sub(r"-+", "-", slug).strip("-")
        return cls(
            app_name=d.get("appName", "MyApp"),
            slug=slug,
            package_name=d.get("packageName", f"com.forge.{slug.replace('-','')}"),
            description=d.get("description", ""),
            tagline=d.get("tagline", ""),
            primary_color=d.get("primaryColor", "#2563eb"),
            accent_color=d.get("accentColor", "#f59e0b"),
            background_color=d.get("backgroundColor", "#0f172a"),
            surface_color=d.get("surfaceColor", "#1e293b"),
            text_color=d.get("textColor", "#f8fafc"),
            screens=d.get("screens", []),
            data_models=d.get("dataModels", []),
            features=d.get("features", []),
            user_stories=d.get("userStories", []),
            navigation=d.get("navigation", "tab"),
            tab_screens=d.get("tabScreens", []),
            stack_screens=d.get("stackScreens", []),
            drawer_screens=d.get("drawerScreens", []),
            storage=d.get("storage", "asyncstorage"),
            storage_key=d.get("storageKey", "app_data"),
            auth_required=d.get("authRequired", False),
            auth_method=d.get("authMethod", "none"),
            external_apis=d.get("externalAPIs", []),
            dependencies=d.get("dependencies", []),
            dev_dependencies=d.get("devDependencies", []),
            target_platforms=d.get("targetPlatforms", ["ios", "android"]),
            clarifying_questions=d.get("clarifyingQuestions", []),
            assumptions=d.get("assumptions", []),
            raw=d,
        )


@dataclass
class FileEntry:
    path: str
    content: str
    agent: str = "codegen"
    source: str = "nvidia"   # 'nvidia' | 'groq' | 'static' | 'template'
    size: int = 0
    lines: int = 0

    def __post_init__(self) -> None:
        self.size = len(self.content.encode())
        self.lines = self.content.count("\n") + 1


@dataclass
class GenerationSession:
    """Central state object for one full generation run."""

    prompt: str
    settings: Any  # Settings object
    session_id: str = ""
    started_at: float = field(default_factory=time.time)

    # Pipeline state
    plan: Optional[AppPlan] = None
    architecture: Optional[dict] = None
    files: dict[str, FileEntry] = field(default_factory=dict)
    test_files: dict[str, FileEntry] = field(default_factory=dict)
    ci_files: dict[str, FileEntry] = field(default_factory=dict)

    # Agent tracking
    agents: list[AgentRun] = field(default_factory=list)
    logs: list[LogEntry] = field(default_factory=list)

    # Clarification state
    clarification_answers: dict[str, str] = field(default_factory=dict)
    needs_clarification: bool = False

    # Output
    project_dir: Optional[Path] = None
    zip_path: Optional[Path] = None

    # Telegram
    telegram_message_id: Optional[int] = None

    # Callbacks (set by CLI/Telegram handler)
    on_log: Optional[Callable[[LogEntry], None]] = None
    on_agent_update: Optional[Callable[[AgentRun], None]] = None
    on_progress: Optional[Callable[[int, int, str], None]] = None

    def __post_init__(self) -> None:
        import uuid
        if not self.session_id:
            self.session_id = str(uuid.uuid4())[:8]

    # ── Logging ───────────────────────────────────────────────────────────────

    def log(self, agent: str, message: str, level: str = "info") -> None:
        entry = LogEntry(agent=agent, message=message, level=level, elapsed=self.elapsed)
        self.logs.append(entry)
        if self.on_log:
            self.on_log(entry)

    def log_success(self, agent: str, msg: str) -> None:
        self.log(agent, msg, "success")

    def log_warn(self, agent: str, msg: str) -> None:
        self.log(agent, msg, "warning")

    def log_error(self, agent: str, msg: str) -> None:
        self.log(agent, msg, "error")

    def log_highlight(self, agent: str, msg: str) -> None:
        self.log(agent, msg, "highlight")

    # ── Agent tracking ────────────────────────────────────────────────────────

    def start_agent(self, name: str, engine: str) -> AgentRun:
        run = AgentRun(name=name, engine=engine)
        run.start()
        self.agents.append(run)
        if self.on_agent_update:
            self.on_agent_update(run)
        return run

    def finish_agent(self, run: AgentRun, **stats: Any) -> None:
        run.done(**stats)
        if self.on_agent_update:
            self.on_agent_update(run)

    def fail_agent(self, run: AgentRun, error: str) -> None:
        run.fail(error)
        if self.on_agent_update:
            self.on_agent_update(run)

    # ── File management ───────────────────────────────────────────────────────

    def add_file(self, path: str, content: str, agent: str = "codegen", source: str = "nvidia") -> None:
        self.files[path] = FileEntry(path=path, content=content, agent=agent, source=source)

    def add_test(self, path: str, content: str) -> None:
        self.test_files[path] = FileEntry(path=path, content=content, agent="tester", source="groq")

    def add_ci(self, path: str, content: str) -> None:
        self.ci_files[path] = FileEntry(path=path, content=content, agent="packager", source="static")

    def all_files(self) -> dict[str, str]:
        result = {k: v.content for k, v in self.files.items()}
        result.update({k: v.content for k, v in self.test_files.items()})
        result.update({k: v.content for k, v in self.ci_files.items()})
        return result

    # ── Stats ─────────────────────────────────────────────────────────────────

    @property
    def elapsed(self) -> float:
        return time.time() - self.started_at

    @property
    def total_files(self) -> int:
        return len(self.files) + len(self.test_files) + len(self.ci_files)

    @property
    def total_lines(self) -> int:
        return sum(f.lines for f in self.files.values())

    @property
    def total_size_kb(self) -> float:
        return sum(f.size for f in self.files.values()) / 1024

    @property
    def nvidia_files(self) -> int:
        return sum(1 for f in self.files.values() if f.source == "nvidia")

    @property
    def groq_files(self) -> int:
        return sum(1 for f in self.files.values() if f.source == "groq")
