import pygame

from engine import collection
from engine.growth import MAX_LEVEL
from engine.plants import POOL, RARITY_COLORS, plant_rarity
from engine.skills import skill_from_dict
from engine.unit import ROLE_NAMES
from ui.menu import _bracket_text
from ui.scaling import make_window, request_quit, to_logical

WIN_W, WIN_H = 1040, 720
BG = (24, 30, 24)
TEXT = (230, 235, 230)
MUTED = (150, 155, 150)
HINT = (255, 220, 120)
PANEL = (32, 38, 32)
BTN = (60, 90, 60)
BTN_HOVER = (85, 125, 85)
ACCENT = (120, 220, 190)
DANGER = (150, 70, 70)

ORDINALS = {1: "2nd", 2: "3rd"}


class ExtractionScreen:
    def __init__(self):
        pygame.init()
        self.window, self.scale = make_window(WIN_W, WIN_H)
        self.screen = pygame.Surface((WIN_W, WIN_H))
        pygame.display.set_caption("Plants — Ability Extraction")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.SysFont("segoeui", 30)
        self.bold = pygame.font.SysFont("segoeui", 20, bold=True)
        self.small = pygame.font.SysFont("segoeui", 16)
        self.tiny = pygame.font.SysFont("segoeui", 13)
        self.back_rect = pygame.Rect(16, 16, 90, 34)
        self.step_rect = pygame.Rect(116, 16, 120, 34)
        self.tab_rects = {
            "extract": pygame.Rect(WIN_W // 2 - 210, 78, 200, 36),
            "inject": pygame.Rect(WIN_W // 2 + 10, 78, 200, 36),
        }
        self.tab = "extract"
        self.phase = "list"   # extract: list -> skills ; inject: ability -> plant -> slot
        self.source = None    # selected source plant (extract)
        self.ability = None   # selected extracted-ability index (inject)
        self.target = None    # selected target plant (inject)
        self.item_rects = []  # (key, rect) for the current phase list
        self.confirm = None   # dict, or None
        self.confirm_ok = pygame.Rect(0, 0, 0, 0)
        self.confirm_no = pygame.Rect(0, 0, 0, 0)
        self.mouse_pos = (0, 0)

    # ---- helpers -------------------------------------------------------
    def _deep(self):
        if self.tab == "extract":
            return self.phase == "skills"
        return self.phase in ("plant", "slot")

    def _reset(self, tab):
        self.tab = tab
        self.phase = "list" if tab == "extract" else "ability"
        self.source = None
        self.ability = None
        self.target = None
        self.confirm = None

    # ---- click ---------------------------------------------------------
    def _handle(self, key):
        if self.tab == "extract":
            if self.phase == "list":
                self.source = key
                self.phase = "skills"
            elif self.phase == "skills":
                self.confirm = {"kind": "extract", "name": self.source, "index": key}
        else:
            if self.phase == "ability":
                self.ability = key
                self.phase = "plant"
            elif self.phase == "plant":
                self.target = key
                self.phase = "slot"
            elif self.phase == "slot":
                self.confirm = {"kind": "inject", "ability": self.ability,
                                "target": self.target, "slot": key}

    def _confirm_ok(self):
        c = self.confirm
        if c["kind"] == "extract":
            if collection.extract_skill(c["name"], c["index"]):
                self.phase = "list"
                self.source = None
        else:
            if collection.inject_skill(c["ability"], c["target"], c["slot"]):
                self.phase = "ability"
                self.ability = None
                self.target = None
        self.confirm = None

    def run(self):
        running = True
        while running:
            self.clock.tick(60)
            self.mouse_pos = to_logical(pygame.mouse.get_pos(), self.scale)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    request_quit()
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    pos = to_logical(event.pos, self.scale)
                    if self.confirm is not None:
                        if self.confirm_ok.collidepoint(pos):
                            self._confirm_ok()
                        elif self.confirm_no.collidepoint(pos):
                            self.confirm = None
                        continue
                    if self.back_rect.collidepoint(pos):
                        running = False
                        continue
                    if self.step_rect.collidepoint(pos) and self._deep():
                        self._step_back()
                        continue
                    if self.tab_rects["extract"].collidepoint(pos):
                        if self.tab != "extract":
                            self._reset("extract")
                        continue
                    if self.tab_rects["inject"].collidepoint(pos):
                        if self.tab != "inject":
                            self._reset("inject")
                        continue
                    for key, rect in self.item_rects:
                        if rect.collidepoint(pos):
                            self._handle(key)
                            break
            self._draw()
            pygame.transform.smoothscale(self.screen, self.window.get_size(), self.window)
            pygame.display.flip()

    def _step_back(self):
        if self.tab == "extract":
            self.phase = "list"
            self.source = None
        else:
            if self.phase == "slot":
                self.phase = "plant"
                self.target = None
            elif self.phase == "plant":
                self.phase = "ability"
                self.ability = None

    # ---- drawing -------------------------------------------------------
    def _draw(self):
        self.item_rects = []
        self.screen.fill(BG)
        pygame.draw.rect(self.screen, BTN, self.back_rect, border_radius=6)
        t = self.small.render("Back", True, TEXT)
        self.screen.blit(t, t.get_rect(center=self.back_rect.center))
        if self._deep():
            pygame.draw.rect(self.screen, BTN, self.step_rect, border_radius=6)
            st = self.small.render("← Back", True, TEXT)
            self.screen.blit(st, st.get_rect(center=self.step_rect.center))
        title = self.font.render("Ability Extraction", True, HINT)
        self.screen.blit(title, title.get_rect(center=(WIN_W // 2, 38)))
        sub = self.small.render("Sacrifice a maxed plant to copy an ability; inject it into a same-role plant.",
                                True, MUTED)
        self.screen.blit(sub, sub.get_rect(center=(WIN_W // 2, 68)))
        for name, rect in self.tab_rects.items():
            active = self.tab == name
            pygame.draw.rect(self.screen, BTN_HOVER if active else BTN, rect, border_radius=8)
            pygame.draw.rect(self.screen, (150, 190, 150) if active else (70, 80, 70),
                             rect, 2 if active else 1, border_radius=8)
            tt = self.small.render(name.capitalize(), True, TEXT)
            self.screen.blit(tt, tt.get_rect(center=rect.center))
        if self.confirm is not None:
            self._draw_confirm()
        elif self.tab == "extract":
            self._draw_extract()
        else:
            self._draw_inject()

    # ---- extract tab ---------------------------------------------------
    def _draw_extract(self):
        if self.phase == "list":
            sources = sorted(n for n in collection.owned() if collection.can_extract(n) == "")
            if not sources:
                t = self.small.render(
                    f"No plant is ready yet. A plant needs Lv {MAX_LEVEL} and 2 evolutions.", True, MUTED)
                self.screen.blit(t, t.get_rect(center=(WIN_W // 2, 220)))
                return
            rows = []
            for n in sources:
                spec = POOL[n]
                rows.append((n, n, RARITY_COLORS.get(plant_rarity(n), TEXT),
                             f"{ROLE_NAMES.get(spec['role'], spec['role'])}  ·  Lv {MAX_LEVEL} · 2 evolutions"))
            self._draw_items(rows, "Plants ready to extract from:")
        else:
            self._draw_extract_skills()

    def _draw_extract_skills(self):
        rows = []
        for i, sk in enumerate(POOL[self.source]["skills"]):
            reason = collection.extraction_blocked_reason(self.source, i, sk)
            if reason:
                rows.append((i, sk.name, MUTED, f"Not extractable — {reason}", False))
            else:
                rows.append((i, sk.name, ACCENT, _skill_sub(sk)))
        self._draw_items(rows, f"Extract from {self.source}:")

    # ---- inject tab ----------------------------------------------------
    def _draw_inject(self):
        if self.phase == "ability":
            ext = collection.extracted_skills()
            if not ext:
                t = self.small.render("No extracted abilities yet — extract one first.", True, MUTED)
                self.screen.blit(t, t.get_rect(center=(WIN_W // 2, 220)))
                return
            rows = [(i, r["skill_name"], ACCENT,
                     f"from {r['source']} · {ROLE_NAMES.get(r['role'], r['role'])}")
                    for i, r in enumerate(ext)]
            self._draw_items(rows, "Pick an extracted ability:")
        elif self.phase == "plant":
            rec = collection.extracted_skills()[self.ability]
            targets = sorted(n for n in collection.owned() if collection.can_inject(n, rec) == "")
            if not targets:
                t = self.small.render("No eligible plants (same role, no injection).", True, MUTED)
                self.screen.blit(t, (60, 150))
                return
            rows = [(n, n, RARITY_COLORS.get(plant_rarity(n), TEXT),
                     ROLE_NAMES.get(POOL[n]["role"], POOL[n]["role"])) for n in targets]
            self._draw_items(rows, f"Inject {rec['skill_name']} into:")
        else:
            self._draw_inject_slots()

    def _draw_inject_slots(self):
        rec = collection.extracted_skills()[self.ability]
        spec = POOL[self.target]
        rows = []
        for slot in (1, 2):
            if slot >= len(spec["skills"]):
                continue
            sk = spec["skills"][slot]
            reason = collection.injection_slot_reason(self.target, rec, slot)
            if reason:
                rows.append((slot, f"Replace {ORDINALS[slot]} skill ({sk.name})", MUTED,
                             f"Not allowed — {reason}", False))
            else:
                rows.append((slot, f"Replace {ORDINALS[slot]} skill ({sk.name})", HINT,
                             f"CD {sk.cooldown} · {sk.description}"))
        self._draw_items(rows, f"Inject {rec['skill_name']} into {self.target} — replace:")

    # ---- shared list ---------------------------------------------------
    def _draw_items(self, rows, header):
        h = self.bold.render(header, True, MUTED)
        self.screen.blit(h, (60, 118))
        y = 146
        for row in rows:
            key, label, color, sub = row[0], row[1], row[2], row[3]
            enabled = row[4] if len(row) > 4 else True
            rect = pygame.Rect(60, y, 560, 46)
            if enabled:
                self.item_rects.append((key, rect))
            hover = enabled and rect.collidepoint(self.mouse_pos)
            pygame.draw.rect(self.screen, BTN_HOVER if hover else PANEL, rect, border_radius=6)
            pygame.draw.rect(self.screen, (70, 90, 70) if hover else (60, 70, 60),
                             rect, 2 if hover else 1, border_radius=6)
            lt = self.bold.render(label, True, color)
            self.screen.blit(lt, (rect.x + 12, rect.y + 5))
            if sub:
                st = self.tiny.render(sub, True, MUTED)
                self.screen.blit(st, (rect.x + 12, rect.y + 26))
            y += 54

    # ---- confirm overlay -----------------------------------------------
    def _draw_confirm(self):
        shade = pygame.Surface((WIN_W, WIN_H), pygame.SRCALPHA)
        shade.fill((0, 0, 0, 170))
        self.screen.blit(shade, (0, 0))
        panel = pygame.Rect(WIN_W // 2 - 250, 200, 500, 200)
        pygame.draw.rect(self.screen, PANEL, panel, border_radius=12)
        pygame.draw.rect(self.screen, (90, 110, 90), panel, 2, border_radius=12)
        c = self.confirm
        if c["kind"] == "extract":
            sk = POOL[c["name"]]["skills"][c["index"]]
            title = f"Extract {sk.name}?"
            body = [f"from {c['name']} — this resets it to Lv 1 / Stage 1.",
                    "The ability is stored for one injection."]
        else:
            rec = collection.extracted_skills()[c["ability"]]
            spec = POOL[c["target"]]
            title = f"Inject {rec['skill_name']} into {c['target']}?"
            body = [f"Replaces the {ORDINALS[c['slot']]} skill ({spec['skills'][c['slot']].name}).",
                    "Consumes the extracted ability.",
                    "Removing it later destroys the ability."]
        tt = self.bold.render(title, True, HINT)
        self.screen.blit(tt, tt.get_rect(center=(panel.centerx, panel.y + 40)))
        yy = panel.y + 76
        for line in body:
            b = self.small.render(line, True, TEXT)
            self.screen.blit(b, b.get_rect(center=(panel.centerx, yy)))
            yy += 28
        self.confirm_ok = pygame.Rect(panel.centerx - 145, panel.bottom - 52, 135, 38)
        self.confirm_no = pygame.Rect(panel.centerx + 10, panel.bottom - 52, 135, 38)
        pygame.draw.rect(self.screen, BTN_HOVER, self.confirm_ok, border_radius=8)
        ot = self.bold.render("Confirm", True, TEXT)
        self.screen.blit(ot, ot.get_rect(center=self.confirm_ok.center))
        pygame.draw.rect(self.screen, DANGER, self.confirm_no, border_radius=8)
        nt = self.bold.render("Cancel", True, TEXT)
        self.screen.blit(nt, nt.get_rect(center=self.confirm_no.center))


def _skill_sub(sk) -> str:
    bracket = _bracket_text(sk)
    head = f"CD {sk.cooldown}"
    if bracket:
        head += f"  [{bracket}]"
    return f"{head} · {sk.description}"